from __future__ import annotations

import types
import typing
from typing import Any

from pydantic import BaseModel, Field, model_validator

OPENNESS = ("open-weights", "open-source", "proprietary", "restricted", "unknown")
MODALITIES = ("text", "image", "audio", "video", "code", "embedding", "3d", "multimodal")


def _is_list_type(tp: Any) -> bool:
    origin = typing.get_origin(tp)
    if origin is list:
        return True
    if origin in (typing.Union, types.UnionType):
        return any(_is_list_type(a) for a in typing.get_args(tp))
    return False


def _is_type(tp: Any, target: type) -> bool:
    if tp is target:
        return True
    if typing.get_origin(tp) in (typing.Union, types.UnionType):
        return any(a is target for a in typing.get_args(tp))
    return False


class Tolerant(BaseModel):
    """LLM outputs are messy: null for lists, "yes"/"no" for booleans, "70B" for integers. Coerce instead of failing (and escalating)."""

    @model_validator(mode="before")
    @classmethod
    def _coerce(cls, data: Any) -> Any:
        if not isinstance(data, dict):
            return data
        out = dict(data)
        for name, field in cls.model_fields.items():
            if name not in out:
                continue
            v = out[name]
            tp = field.annotation
            if _is_list_type(tp):
                if v is None or v == "" or v == "null":
                    out[name] = []
                elif isinstance(v, (str, int, float)):
                    out[name] = [x.strip() for x in str(v).split(",") if x.strip()] if isinstance(v, str) else [v]
                elif isinstance(v, dict):
                    out[name] = [str(x) for x in v.values() if x not in (None, "")]
            elif _is_type(tp, bool) and isinstance(v, str):
                low = v.strip().lower()
                out[name] = True if low in ("yes", "true", "supported", "y") else False if low in ("no", "false", "unsupported", "n", "none") else None
            elif _is_type(tp, int) and isinstance(v, str):
                from aiatlas.sdk.extract.numbers import parse_context_length, parse_param_count

                parsed = parse_param_count(v) if "param" in name else parse_context_length(v)
                out[name] = parsed if parsed is not None else None
            elif _is_type(tp, float) and isinstance(v, str):
                import re

                m = re.search(r"-?\d+(?:\.\d+)?", v.replace(",", ""))
                out[name] = float(m.group(0)) if m else None
            elif _is_type(tp, str) and isinstance(v, (int, float)) and not isinstance(v, bool):
                out[name] = str(v)
            elif v == "null":
                out[name] = None
        return out


class DocumentClassification(Tolerant):
    doc_class: str = Field(description="one of: model_release, model_update, pricing, research_paper, company_news, product_launch, "
                                       "framework_release, dataset_release, benchmark, hardware, regulation, incident, funding, acquisition, other")
    is_primary_source: bool | None = None
    mentioned_models: list[str] = Field(default_factory=list)
    mentioned_companies: list[str] = Field(default_factory=list)
    language: str | None = None
    relevance: float = Field(default=0.5, ge=0, le=1, description="relevance to the AI ecosystem")


class ModelPassport(Tolerant):
    name: str | None = None
    developer: str | None = Field(default=None, description="organization that developed the model")
    family: str | None = None
    version: str | None = None
    release_date: str | None = Field(default=None, description="ISO date if stated (YYYY-MM-DD, YYYY-MM or YYYY)")
    status: str | None = Field(default=None, description="available | preview | deprecated | retired | announced")
    openness: str | None = Field(default=None, description=" | ".join(OPENNESS))
    license: str | None = None
    architecture: str | None = Field(default=None, description="e.g. transformer decoder, MoE, diffusion, state-space")
    parameter_count: int | None = Field(default=None, description="total parameters, as an integer (e.g. 70000000000)")
    active_parameter_count: int | None = None
    is_moe: bool | None = None
    modalities_input: list[str] = Field(default_factory=list)
    modalities_output: list[str] = Field(default_factory=list)
    context_length: int | None = Field(default=None, description="tokens")
    max_output_tokens: int | None = None
    knowledge_cutoff: str | None = None
    languages: list[str] = Field(default_factory=list)
    tool_calling: bool | None = None
    structured_output: bool | None = None
    reasoning: bool | None = None
    vision: bool | None = None
    audio: bool | None = None
    fine_tuning_available: bool | None = None
    tokenizer: str | None = None
    training_data_notes: str | None = None
    predecessor: str | None = None
    base_model: str | None = None
    quantizations: list[str] = Field(default_factory=list)
    paper_url: str | None = None
    model_card_url: str | None = None
    repository_url: str | None = None
    official_page_url: str | None = None
    hardware_requirements: str | None = None
    safety_notes: str | None = None
    evidence: list[str] = Field(default_factory=list, description="short quotes supporting the most important fields")


class PriceLine(Tolerant):
    model: str
    provider_model_id: str | None = None
    input_per_mtok: float | None = Field(default=None, description="USD per 1M input tokens")
    output_per_mtok: float | None = None
    cached_input_per_mtok: float | None = None
    cache_write_per_mtok: float | None = None
    batch_input_per_mtok: float | None = None
    batch_output_per_mtok: float | None = None
    per_image: float | None = None
    context_length: int | None = None
    max_output_tokens: int | None = None
    notes: str | None = None


class PricingExtraction(Tolerant):
    provider: str | None = None
    currency: str = "USD"
    effective_date: str | None = None
    prices: list[PriceLine] = Field(default_factory=list)


class CompanyPassport(Tolerant):
    name: str | None = None
    legal_name: str | None = None
    country: str | None = Field(default=None, description="ISO 3166-1 alpha-2 if determinable")
    headquarters: str | None = None
    founded: str | None = None
    founders: list[str] = Field(default_factory=list)
    leadership: list[str] = Field(default_factory=list)
    website: str | None = None
    description: str | None = None
    products: list[str] = Field(default_factory=list)
    models: list[str] = Field(default_factory=list)
    investors: list[str] = Field(default_factory=list)
    parent_company: str | None = None
    subsidiaries: list[str] = Field(default_factory=list)
    employee_count: int | None = None
    funding_total_usd: float | None = None


class PaperPassport(Tolerant):
    title: str | None = None
    authors: list[str] = Field(default_factory=list)
    affiliations: list[str] = Field(default_factory=list)
    date: str | None = None
    field: str | None = None
    summary: str | None = Field(default=None, description="2-3 sentence factual summary")
    methods: list[str] = Field(default_factory=list)
    models: list[str] = Field(default_factory=list, description="models introduced or evaluated")
    datasets: list[str] = Field(default_factory=list)
    benchmarks: list[str] = Field(default_factory=list)
    key_claims: list[str] = Field(default_factory=list)
    results: list[str] = Field(default_factory=list)
    limitations: list[str] = Field(default_factory=list)
    code_url: str | None = None


class BenchmarkRow(Tolerant):
    model: str
    score: float
    metric: str | None = None
    config: str | None = None


class BenchmarkResultExtraction(Tolerant):
    benchmark: str | None = None
    metric: str | None = None
    higher_is_better: bool | None = True
    evaluated_at: str | None = None
    rows: list[BenchmarkRow] = Field(default_factory=list)


class HardwareSpec(Tolerant):
    name: str | None = None
    manufacturer: str | None = None
    kind: str | None = Field(default=None, description="gpu | cpu | npu | tpu | asic | soc | accelerator | system")
    architecture: str | None = None
    release_date: str | None = None
    memory_gb: float | None = None
    memory_type: str | None = None
    memory_bandwidth_gbs: float | None = None
    compute_fp16_tflops: float | None = None
    compute_fp8_tflops: float | None = None
    compute_int8_tops: float | None = None
    tdp_watts: float | None = None
    form_factor: str | None = None
    price_usd: float | None = None
    interconnect: str | None = None


class ReleaseAnnouncement(Tolerant):
    kind: str | None = Field(default=None, description="new_model | model_update | pricing_change | deprecation | new_product | framework_release | research | other")
    title: str | None = None
    date: str | None = None
    organization: str | None = None
    models: list[str] = Field(default_factory=list)
    summary: str | None = None
    facts: list[str] = Field(default_factory=list, description="atomic factual statements (one per item)")
    model_passport: ModelPassport | None = None
    pricing: PricingExtraction | None = None


TASKS: dict[str, tuple[type[BaseModel], str]] = {
    "classify": (DocumentClassification, "small"),
    "model_passport": (ModelPassport, "medium"),
    "pricing": (PricingExtraction, "medium"),
    "company_passport": (CompanyPassport, "medium"),
    "paper_passport": (PaperPassport, "medium"),
    "benchmark_results": (BenchmarkResultExtraction, "medium"),
    "hardware_spec": (HardwareSpec, "medium"),
    "release_announcement": (ReleaseAnnouncement, "medium"),
}


def schema_for(task: str) -> tuple[type[BaseModel], str]:
    return TASKS[task]