"""normalize.py —— 上游响应 → 契约形态的归一化器（API.md §7/§8/§9）。

真实上游键以 DISCOVERY.md §4.1 为准；别名规则按 API.md 逐条实现。
本模块是纯函数集合（不发起 IO），供 qa 路由（A2）与测试共用。
"""
from __future__ import annotations

from typing import Any


def pick(obj: dict, *names: str, default=None):
    """取第一个存在且非 None 的键（归一化核心工具）。保留 PM 脚手架语义。"""
    for n in names:
        if isinstance(obj, dict) and n in obj and obj[n] is not None:
            return obj[n]
    return default


# ---------------------------------------------------------------------------
# §8 Source 归一化（缺失字段键必须存在，值为 null）
# ---------------------------------------------------------------------------

def normalize_source(item: dict[str, Any] | None) -> dict[str, Any]:
    """单个 source item → API.md §8 形态（所有键存在，缺失为 null）。

    别名映射（按 API.md §8 精确实现）：
    - content|evidence|text|snippet → snippet
    - doc_id|paper_id|pubmed_id|id|chunk_id → paper_id（chunk_id 单独保留）
    - final_score|score|confidence|rerank_score → score（真实键 final_score）
    - title|paper_title|source → title（真实 item 的 source 是断言主体实体名；
      title 存在时优先取 title，两者并存时 source 落入 metadata.source_entity）
    - url|doi_url → url；只有 doi 时派生 url = https://doi.org/{doi}
    - evidence_id|id → id；relation_id|assertion_id → metadata.assertion_id
    - output_layer / claim_scope / cluster_id / support_count / section /
      species / vector_score / rrf_score / rerank_score / retrieval_methods /
      retrieval_queries / cluster_member_* 原样透传（或进 metadata）
    """
    if not isinstance(item, dict):
        item = {}

    source = pick(item, "source")            # 断言主体实体名（上游语义）
    title = pick(item, "title", "paper_title")

    doi = pick(item, "doi")
    url = pick(item, "url", "doi_url")
    if url is None and doi:
        url = f"https://doi.org/{doi}"

    metadata: dict[str, Any] = {}
    if isinstance(pick(item, "metadata"), dict):
        metadata.update(item["metadata"])
    if source is not None and title is not None and source != title:
        metadata.setdefault("source_entity", source)
    for key in (
        "vector_score", "rrf_score", "rerank_score", "kb",
        "retrieval_methods", "retrieval_queries",
        "cluster_member_count", "cluster_member_assertion_ids",
        "cluster_member_paper_ids", "cluster_member_pmids",
        "output_layer", "claim_scope", "species",
    ):
        if key in item and item[key] is not None:
            metadata.setdefault(key, item[key])

    assertion_id = pick(item, "relation_id", "assertion_id")
    if assertion_id is not None:
        metadata.setdefault("assertion_id", assertion_id)

    return {
        "id": pick(item, "evidence_id", "id"),
        "title": title,
        "paper_id": pick(item, "doc_id", "paper_id", "pubmed_id", "id", "chunk_id"),
        "chunk_id": pick(item, "chunk_id"),
        "pmid": pick(item, "pmid"),
        "pmcid": pick(item, "pmcid"),
        "doi": doi,
        "url": url,
        "snippet": pick(item, "content", "evidence", "text", "snippet"),
        "score": pick(item, "final_score", "score", "confidence", "rerank_score"),
        "locator": pick(item, "locator"),
        "section": pick(item, "section"),
        "support_count": pick(item, "support_count"),
        "type": pick(item, "type"),
        "metadata": metadata,
    }


# ---------------------------------------------------------------------------
# §9 图谱：hop / chain / {nodes,edges,chains}
# ---------------------------------------------------------------------------

def normalize_hop(hop: dict[str, Any]) -> dict[str, Any]:
    """hop 归一化：
    from_name|source → source、to_name|target → target、
    from_id|source_canonical_id → source 节点 id、to_id|target_canonical_id → target 节点 id、
    relation 原样、source_type|target_type → source/target 节点 type、
    paper_id / pmid / pmcid / doi 透传。
    """
    if not isinstance(hop, dict):
        hop = {}
    return {
        "source": pick(hop, "from_name", "source"),
        "target": pick(hop, "to_name", "target"),
        "source_id": pick(hop, "from_id", "source_canonical_id"),
        "target_id": pick(hop, "to_id", "target_canonical_id"),
        "source_type": pick(hop, "source_type"),
        "target_type": pick(hop, "target_type"),
        "relation": pick(hop, "relation"),
        "assertion_id": pick(hop, "assertion_id"),
        "paper_id": pick(hop, "paper_id"),
        "pmid": pick(hop, "pmid"),
        "pmcid": pick(hop, "pmcid"),
        "doi": pick(hop, "doi"),
    }


def normalize_chain_item(item: dict[str, Any]) -> dict[str, Any]:
    """单个 chain item（真实形态：chain_id / hops[] 等）→ {chain_id, hops}。"""
    if not isinstance(item, dict):
        item = {}
    hops = [normalize_hop(h) for h in item.get("hops", []) if isinstance(h, dict)]
    return {"chain_id": pick(item, "chain_id"), "hops": hops}


def extract_chains(raw: dict[str, Any] | None) -> list[dict[str, Any]]:
    """从可能的两类源提取 chain 列表：
    - `native_graph_chains`：列表（真实形态，305 条）
    - `graph_knowledge_chains`：dict 形态 `{mode, chains[], chain_count, …}`，取 `.chains`
    两种形态都要能提取出 hop 列表。
    """
    if not isinstance(raw, dict):
        return []
    chains = pick(raw, "native_graph_chains", "graph_knowledge_chains")
    if isinstance(chains, dict):
        chains = chains.get("chains")
    if not isinstance(chains, list):
        return []
    return [normalize_chain_item(c) for c in chains]


def normalize_graph(raw: dict[str, Any] | None) -> dict[str, Any]:
    """从 `knowledge_graph` / `graph_edges` / `native_graph_chains`（列表形态）/
    `graph_knowledge_chains`（dict 形态）提取并归一化为 {nodes, edges, chains}。

    - 节点去重（按 id；id 缺失的 hop 跳过）；label 取 from_name/to_name、type 取
      source_type/target_type。
    - 边携带 chain_id 与 hop 序号（hop_index 缺失时按列表序补）。
    - 链保持 chain_id + hops 列表。
    """
    raw = raw if isinstance(raw, dict) else {}
    chains = extract_chains(raw)

    nodes: dict[str, dict[str, Any]] = {}
    edges: list[dict[str, Any]] = []

    def add_node(node_id: Any, name: Any, typ: Any) -> None:
        if not node_id:
            return
        key = str(node_id)
        if key not in nodes:
            nodes[key] = {"id": node_id, "label": name, "type": typ}

    for chain in chains:
        for idx, hop in enumerate(chain["hops"]):
            hop_index = hop.get("hop_index")
            if hop_index is None:
                hop_index = idx
            # 节点 id 优先 canonical id；label 优先名字，退化到 id
            src_id = pick(hop, "source_id", "source") or hop.get("source")
            tgt_id = pick(hop, "target_id", "target") or hop.get("target")
            src_label = hop.get("source") or src_id
            tgt_label = hop.get("target") or tgt_id
            add_node(src_id, src_label, hop.get("source_type"))
            add_node(tgt_id, tgt_label, hop.get("target_type"))
            if src_id and tgt_id:
                edges.append({
                    "source": src_id,
                    "target": tgt_id,
                    "relation": hop.get("relation"),
                    "label": hop.get("relation"),
                    "chain_id": chain.get("chain_id"),
                    "hop": hop_index,
                    "paper_id": hop.get("paper_id"),
                    "pmid": hop.get("pmid"),
                    "pmcid": hop.get("pmcid"),
                    "doi": hop.get("doi"),
                })

    # 兼容显式 {nodes, edges} 形态（knowledge_graph / graph_edges 直接给边列表时）
    explicit_nodes = raw.get("nodes")
    explicit_edges = raw.get("edges")
    if isinstance(explicit_nodes, list):
        for n in explicit_nodes:
            if isinstance(n, dict) and n.get("id"):
                key = str(n["id"])
                if key not in nodes:
                    nodes[key] = {
                        "id": n["id"],
                        "label": pick(n, "label"),
                        "type": pick(n, "type"),
                    }
    if isinstance(explicit_edges, list):
        edges.extend(
            {
                "source": pick(e, "source"),
                "target": pick(e, "target"),
                "relation": pick(e, "relation"),
                "label": pick(e, "label", "relation"),
                "chain_id": pick(e, "chain_id"),
                "hop": pick(e, "hop"),
                "paper_id": pick(e, "paper_id"),
                "pmid": pick(e, "pmid"),
                "pmcid": pick(e, "pmcid"),
                "doi": pick(e, "doi"),
            }
            for e in explicit_edges
            if isinstance(e, dict)
        )

    return {
        "nodes": list(nodes.values()),
        "edges": edges,
        "chains": chains,
    }


# ---------------------------------------------------------------------------
# §7 归一化 QA 响应（/chat、/chat/stream complete、/search 共用）
# ---------------------------------------------------------------------------

def normalize_qa_response(raw: dict[str, Any] | None) -> dict[str, Any]:
    """上游 31 键响应 → API.md §7 形态。

    别名（真实键优先）：
    - answer|response|message → answer
    - sources|citations → sources；used_sources|references → used_sources
    - query_plan|agentic_plan|retrieval_plan → retrieval_plan（真实键 query_plan）
    - retrieval_notes|notes → retrieval_notes
    - timings|timing → timing（真实键 timings）
    - citation_validation|citation_map → citation_map（真实键 citation_validation）
    - 图谱四源 → graph + chains（见 normalize_graph）
    - success 透传；raw 保存原始完整响应
    """
    raw = raw if isinstance(raw, dict) else {}

    sources = [normalize_source(s) for s in (pick(raw, "sources", "citations") or [])]
    used_sources = [
        normalize_source(s) for s in (pick(raw, "used_sources", "references") or [])
    ]
    graph = normalize_graph(raw)
    retrieval_plan = pick(raw, "query_plan", "agentic_plan", "retrieval_plan")

    return {
        "query_id": pick(raw, "query_id"),
        "version": pick(raw, "version"),
        "mode": pick(raw, "mode", default="chat"),
        "answer": pick(raw, "answer", "response", "message", default=""),
        "source_count": pick(raw, "source_count", default=len(sources)),
        "sources": sources,
        "used_sources": used_sources,
        "chains": graph["chains"],
        "graph": {"nodes": graph["nodes"], "edges": graph["edges"], "chains": graph["chains"]},
        "retrieval_plan": retrieval_plan,
        "retrieval_notes": pick(raw, "retrieval_notes", "notes"),
        "timing": pick(raw, "timings", "timing"),
        "citation_map": pick(raw, "citation_validation", "citation_map"),
        "success": pick(raw, "success", default=True),
        "raw": raw,
    }