"""What a connector produces: entity references, temporal claims, relations, events, prices, benchmark results.
Connectors never touch the database — they return `Facts`; the writer resolves, versions and records them."""
from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime
from typing import Any

CONFIDENCE = ("verified", "high", "medium", "low", "conflicted")


@dataclass
class EntityRef:
    entity_type: str
    name: str
    identifiers: dict[str, str] = field(default_factory=dict)   # scheme -> value  (hf_repo, github_repo, arxiv, doi, domain, pypi, provider_model_id…)
    aliases: list[str] = field(default_factory=list)
    slug_hint: str | None = None
    organization: EntityRef | None = None                       # developer / publisher / owner
    description: str | None = None
    status: str | None = None
    attributes: dict[str, Any] = field(default_factory=dict)    # convenience: each item becomes a Claim
    first_seen_hint: datetime | None = None                     # e.g. release date, to backdate `first_seen_at` for historical backfill
    # canonical hierarchy hints (model_family → model → artifact): the writer materialises `entities.family_id` / `canonical_id`
    family: EntityRef | None = None                             # model → its model_family ("Qwen3.6", "Llama 4", "Claude")
    canonical: EntityRef | None = None                          # artifact → the model it packages/quantises; alias entity → canonical entity
    artifact_kind: str | None = None                            # artifacts only: checkpoint|quantization|conversion|packaging
    identity_confidence: str | None = None                      # high|medium|low — how sure the connector is that this is one real thing
    id: str | None = field(default=None, compare=False)

    def key(self) -> str:
        if self.identifiers:
            scheme, value = sorted(self.identifiers.items())[0]
            return f"{self.entity_type}:{scheme}={value}"
        return f"{self.entity_type}:name={self.name.strip().lower()}"


@dataclass
class Claim:
    entity: EntityRef
    property: str
    value: Any
    unit: str | None = None
    confidence: str | None = None            # default from source tier
    observed_at: datetime | None = None
    effective_at: datetime | None = None
    source_url: str | None = None


@dataclass
class Relation:
    subject: EntityRef
    predicate: str
    object: EntityRef
    attributes: dict[str, Any] = field(default_factory=dict)
    confidence: str | None = None
    source_url: str | None = None


@dataclass
class Event:
    event_type: str
    category: str
    summary: str
    entity: EntityRef | None = None
    old_value: Any = None
    new_value: Any = None
    importance: int = 2
    effective_at: datetime | None = None
    dedupe_key: str | None = None
    source_url: str | None = None
    meta: dict[str, Any] = field(default_factory=dict)


@dataclass
class PriceObs:
    model: EntityRef
    provider: EntityRef
    provider_model_id: str | None = None
    input_per_mtok: float | None = None
    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
    per_request: float | None = None
    currency: str = "USD"
    context_length: int | None = None
    max_output_tokens: int | None = None
    features: dict[str, Any] = field(default_factory=dict)
    source_url: str | None = None
    meta: dict[str, Any] = field(default_factory=dict)

    def price_tuple(self) -> tuple[Any, ...]:
        return (self.input_per_mtok, self.output_per_mtok, self.cached_input_per_mtok, self.cache_write_per_mtok,
                self.batch_input_per_mtok, self.batch_output_per_mtok, self.per_image, self.per_request, self.currency)


@dataclass
class ResultObs:
    model: EntityRef
    benchmark: EntityRef
    score: float
    metric: str | None = None
    unit: str | None = "%"
    higher_is_better: bool = True
    config: dict[str, Any] = field(default_factory=dict)
    evaluated_at: datetime | None = None
    source_url: str | None = None
    confidence: str | None = None
    trust_level: str | None = None           # ontology.benchmarks.TRUST_LEVELS; derived from the source when None
    variant: str | None = None               # benchmark variant label (GPQA Diamond, Verified…); derived from config when None
    run_group: str | None = None             # evaluation run/release the row belongs to; derived from config when None


@dataclass
class Target:
    """Something to fetch. Connectors return targets from `discover()` and may add more from `extract()`."""
    url: str
    doc_type: str = "page"
    entity: EntityRef | None = None
    meta: dict[str, Any] = field(default_factory=dict)
    accept: str | None = None
    min_bytes: int = 64
    escalate: bool = False
    needs_llm: bool = False
    priority: int = 2
    key: str | None = None                   # short handle for --file overrides / fixtures
    rate_per_min: int | None = None


@dataclass
class Facts:
    entities: list[EntityRef] = field(default_factory=list)
    claims: list[Claim] = field(default_factory=list)
    relations: list[Relation] = field(default_factory=list)
    events: list[Event] = field(default_factory=list)
    prices: list[PriceObs] = field(default_factory=list)
    results: list[ResultObs] = field(default_factory=list)
    targets: list[Target] = field(default_factory=list)       # follow-up fetches discovered while extracting
    document_title: str | None = None
    document_entity: EntityRef | None = None                 # main entity described by the document
    llm_hint: str | None = None                              # ask the LLM factory for a specific extraction task

    # ---------------------------------------------------------------------------------------------- builders
    def entity(self, entity_type: str, name: str, **kw: Any) -> EntityRef:
        ref = EntityRef(entity_type=entity_type, name=name.strip(), **kw)
        self.entities.append(ref)
        return ref

    def claim(self, entity: EntityRef, property: str, value: Any, **kw: Any) -> Claim | None:
        if value is None or value == "" or value == [] or value == {}:
            return None
        c = Claim(entity=entity, property=property, value=value, **kw)
        self.claims.append(c)
        return c

    def relate(self, subject: EntityRef, predicate: str, obj: EntityRef, **kw: Any) -> Relation:
        r = Relation(subject=subject, predicate=predicate, object=obj, **kw)
        self.relations.append(r)
        return r

    def event(self, event_type: str, category: str, summary: str, **kw: Any) -> Event:
        e = Event(event_type=event_type, category=category, summary=summary[:500], **kw)
        self.events.append(e)
        return e

    def price(self, **kw: Any) -> PriceObs:
        p = PriceObs(**kw)
        self.prices.append(p)
        return p

    def result(self, **kw: Any) -> ResultObs:
        r = ResultObs(**kw)
        self.results.append(r)
        return r

    def follow(self, url: str, **kw: Any) -> Target:
        t = Target(url=url, **kw)
        self.targets.append(t)
        return t

    def extend(self, other: Facts) -> None:
        self.entities += other.entities
        self.claims += other.claims
        self.relations += other.relations
        self.events += other.events
        self.prices += other.prices
        self.results += other.results
        self.targets += other.targets

    def is_empty(self) -> bool:
        return not (self.entities or self.claims or self.relations or self.events or self.prices or self.results)


# Property → change-event mapping (material properties emit events; noisy metrics never do).
MATERIAL_PROPERTIES: dict[str, tuple[str, int]] = {
    "context_length": ("CONTEXT_CHANGED", 2),
    "max_output_tokens": ("MAX_OUTPUT_CHANGED", 1),
    "status": ("STATUS_CHANGED", 2),
    "license": ("LICENSE_CHANGED", 2),
    "parameter_count": ("PARAMETERS_CHANGED", 1),
    "active_parameter_count": ("PARAMETERS_CHANGED", 1),
    "weights_availability": ("OPENNESS_CHANGED", 3),
    "openness": ("OPENNESS_CHANGED", 3),
    "license_key": ("LICENSE_CHANGED", 2),
    "release_date": ("RELEASE_DATE_CHANGED", 1),
    "knowledge_cutoff": ("KNOWLEDGE_CUTOFF_CHANGED", 1),
    "latest_version": ("VERSION_RELEASED", 2),
    "deprecation_date": ("DEPRECATION_ANNOUNCED", 3),
    "retirement_date": ("RETIREMENT_ANNOUNCED", 3),
    "modalities": ("CAPABILITIES_CHANGED", 2),
    "capabilities": ("CAPABILITIES_CHANGED", 1),
    "rate_limit": ("RATE_LIMIT_CHANGED", 1),
    "memory_gb": ("SPEC_CHANGED", 1),
    "price_usd": ("PRICE_CHANGED", 2),
}

NOISY_PREFIXES = ("metric.", "stats.", "counts.")

EVENT_CATEGORY_BY_TYPE: dict[str, str] = {
    "model": "model", "model_family": "model", "artifact": "model", "company": "company", "organization": "company", "lab": "company", "university": "company",
    "paper": "paper", "dataset": "dataset", "benchmark": "benchmark", "provider": "provider", "framework": "framework",
    "library": "framework", "repository": "repository", "tool": "tool", "agent": "tool", "hardware": "hardware",
    "runtime": "framework", "quantization": "model", "regulation": "regulation", "incident": "incident", "release": "release",
    "mcp_server": "tool", "researcher": "company", "license": "model", "product": "tool", "robot": "hardware",
}


# ---------------------------------------------------------------------------------------------- JSON round-trip (quarantine)
# Quarantined runs hold their Facts in `quarantined_runs.facts` (jsonb) until an operator releases or discards them. Datetimes are
# tagged so they come back as real datetime objects (asyncpg needs them, not ISO strings). Entity references are inlined; the
# resolver caches by `EntityRef.key()`, so equal refs resolve to the same entity after a round-trip.
_DT_TAG = "__datetime__"


def _encode(value: Any) -> Any:
    if isinstance(value, datetime):
        return {_DT_TAG: value.isoformat()}
    if isinstance(value, EntityRef):
        return {"entity_type": value.entity_type, "name": value.name, "identifiers": dict(value.identifiers), "aliases": list(value.aliases),
                "slug_hint": value.slug_hint, "organization": _encode(value.organization), "description": value.description, "status": value.status,
                "attributes": _encode(value.attributes), "first_seen_hint": _encode(value.first_seen_hint), "family": _encode(value.family),
                "canonical": _encode(value.canonical), "artifact_kind": value.artifact_kind, "identity_confidence": value.identity_confidence, "id": value.id}
    if isinstance(value, dict):
        return {str(k): _encode(v) for k, v in value.items()}
    if isinstance(value, (list, tuple, set)):
        return [_encode(v) for v in value]
    if isinstance(value, (str, int, float, bool)) or value is None:
        return value
    return str(value)


def _decode(value: Any) -> Any:
    if isinstance(value, dict):
        if set(value) == {_DT_TAG}:
            return datetime.fromisoformat(value[_DT_TAG])
        return {k: _decode(v) for k, v in value.items()}
    if isinstance(value, list):
        return [_decode(v) for v in value]
    return value


def _ref_from(d: dict[str, Any] | None) -> EntityRef | None:
    if not d:
        return None
    return EntityRef(entity_type=d["entity_type"], name=d["name"], identifiers=dict(d.get("identifiers") or {}), aliases=list(d.get("aliases") or []),
                     slug_hint=d.get("slug_hint"), organization=_ref_from(d.get("organization")), description=d.get("description"), status=d.get("status"),
                     attributes=_decode(d.get("attributes") or {}), first_seen_hint=_decode(d.get("first_seen_hint")), family=_ref_from(d.get("family")),
                     canonical=_ref_from(d.get("canonical")), artifact_kind=d.get("artifact_kind"), identity_confidence=d.get("identity_confidence"), id=d.get("id"))


def facts_to_json(facts: Facts) -> dict[str, Any]:
    """Serialise a Facts object to a JSON-compatible dict (see `facts_from_json`)."""
    def obj(o: Any) -> dict[str, Any]:
        return {k: _encode(v) for k, v in o.__dict__.items()}

    return {
        "version": 1,
        "entities": [_encode(e) for e in facts.entities],
        "claims": [obj(c) for c in facts.claims],
        "relations": [obj(r) for r in facts.relations],
        "events": [obj(e) for e in facts.events],
        "prices": [obj(p) for p in facts.prices],
        "results": [obj(r) for r in facts.results],
        "targets": [obj(t) for t in facts.targets],
        "document_title": facts.document_title,
        "document_entity": _encode(facts.document_entity),
        "llm_hint": facts.llm_hint,
    }


def facts_from_json(data: dict[str, Any]) -> Facts:
    """Inverse of `facts_to_json`."""
    def ref_fields(d: dict[str, Any], *keys: str) -> dict[str, Any]:
        out = {k: _decode(v) for k, v in d.items() if k not in keys}
        for k in keys:
            out[k] = _ref_from(d.get(k))
        return out

    facts = Facts()
    facts.entities = [r for r in (_ref_from(e) for e in data.get("entities") or []) if r]
    facts.claims = [Claim(**ref_fields(c, "entity")) for c in data.get("claims") or []]
    facts.relations = [Relation(**ref_fields(r, "subject", "object")) for r in data.get("relations") or []]
    facts.events = [Event(**ref_fields(e, "entity")) for e in data.get("events") or []]
    facts.prices = [PriceObs(**ref_fields(p, "model", "provider")) for p in data.get("prices") or []]
    facts.results = [ResultObs(**ref_fields(r, "model", "benchmark")) for r in data.get("results") or []]
    facts.targets = [Target(**ref_fields(t, "entity")) for t in data.get("targets") or []]
    facts.document_title = data.get("document_title")
    facts.document_entity = _ref_from(data.get("document_entity"))
    facts.llm_hint = data.get("llm_hint")
    return facts


__all__ = [
    "CONFIDENCE",
    "EVENT_CATEGORY_BY_TYPE",
    "MATERIAL_PROPERTIES",
    "NOISY_PREFIXES",
    "Claim",
    "EntityRef",
    "Event",
    "Facts",
    "PriceObs",
    "Relation",
    "ResultObs",
    "Target",
    "facts_from_json",
    "facts_to_json",
]