"""Data lineage contracts for risk graph observations."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any


@dataclass(slots=True)
class RiskDataLineage:
    observation_key: str
    data_tables: list[str] = field(default_factory=list)
    document_ids: list[str] = field(default_factory=list)
    ocr_job_ids: list[str] = field(default_factory=list)
    agent_run_ids: list[str] = field(default_factory=list)
    tool_call_ids: list[str] = field(default_factory=list)
    rule_versions: list[str] = field(default_factory=list)
    ontology_version: str = ""
    algorithm_version: str = ""
    source_event_ids: list[str] = field(default_factory=list)
    quality_gates: list[str] = field(default_factory=list)

    def as_dict(self) -> dict[str, Any]:
        return {
            "observation_key": self.observation_key,
            "data_tables": list(self.data_tables),
            "document_ids": list(self.document_ids),
            "ocr_job_ids": list(self.ocr_job_ids),
            "agent_run_ids": list(self.agent_run_ids),
            "tool_call_ids": list(self.tool_call_ids),
            "rule_versions": list(self.rule_versions),
            "ontology_version": self.ontology_version,
            "algorithm_version": self.algorithm_version,
            "source_event_ids": list(self.source_event_ids),
            "quality_gates": list(self.quality_gates),
        }


class RiskDataLineageBuilder:
    def build_from_observation(
        self,
        observation: dict[str, Any],
        *,
        source_event_ids: list[str] | None = None,
    ) -> RiskDataLineage:
        evidence = [item for item in observation.get("evidence", []) if isinstance(item, dict)]
        ontology_json = observation.get("ontology_json") or {}
        decision_trace = observation.get("decision_trace") or {}
        data_tables = ["risk_observations"]
        if observation.get("claim_id"):
            data_tables.extend(["expense_claims", "expense_claim_items"])
        if evidence:
            data_tables.append("risk_observation_evidence")

        return RiskDataLineage(
            observation_key=str(observation.get("observation_key") or ""),
            data_tables=_unique(data_tables),
            document_ids=_evidence_values(evidence, ["document_id", "doc_id", "file_id"]),
            ocr_job_ids=_evidence_values(evidence, ["ocr_job_id", "ocr_run_id"]),
            agent_run_ids=_unique(
                [
                    str(observation.get("run_id") or "").strip(),
                    str(decision_trace.get("agent_run_id") or "").strip(),
                ]
            ),
            tool_call_ids=_evidence_values(evidence, ["tool_call_id"]),
            rule_versions=_unique(
                [
                    *_evidence_values(evidence, ["rule_version"]),
                    str(decision_trace.get("rule_version") or "").strip(),
                ]
            ),
            ontology_version=str(ontology_json.get("ontology_version") or "").strip(),
            algorithm_version=str(observation.get("algorithm_version") or "").strip(),
            source_event_ids=_unique(source_event_ids or []),
            quality_gates=_quality_gates(decision_trace),
        )


def _evidence_values(evidence: list[dict[str, Any]], keys: list[str]) -> list[str]:
    values: list[str] = []
    for item in evidence:
        metadata = item.get("metadata") if isinstance(item.get("metadata"), dict) else {}
        for key in keys:
            value = str(item.get(key) or metadata.get(key) or "").strip()
            if value:
                values.append(value)
    return _unique(values)


def _quality_gates(decision_trace: dict[str, Any]) -> list[str]:
    gates = [
        str(decision_trace.get("evidence_source_gate") or "").strip(),
        str(decision_trace.get("data_quality_gate") or "").strip(),
    ]
    sampling = decision_trace.get("sampling_strategy")
    if isinstance(sampling, dict):
        gates.append(str(sampling.get("strategy") or "").strip())
    return _unique([item for item in gates if item and item != "passed"])


def _unique(values: list[str]) -> list[str]:
    return list(dict.fromkeys(str(item).strip() for item in values if str(item).strip()))
