"""Corpus preparation and graph building (file-based; no database)."""

from __future__ import annotations

import json
from collections import Counter
from pathlib import Path

from .canonical import Canonicalizer
from .chunking import chunk_document
from .config import ProjectConfig, load_yaml, manifest, write_json
from .corpus import metric_hits, select_corpus
from .data.qasper import load_split
from .graph import KG, build_graph, save_graph
from .schemas import Chunk, Document, paper_prefix


def read_jsonl(path: str | Path) -> list[dict]:
    p = Path(path)
    if not p.exists():
        return []
    return [json.loads(x) for x in p.read_text(encoding="utf-8").splitlines() if x.strip()]


def write_jsonl(path: str | Path, rows) -> None:
    p = Path(path)
    p.parent.mkdir(parents=True, exist_ok=True)
    with p.open("w", encoding="utf-8") as f:
        for r in rows:
            f.write((r.model_dump_json() if hasattr(r, "model_dump_json") else json.dumps(r, ensure_ascii=False)) + "\n")


def prepare_corpus(cfg: ProjectConfig) -> dict:
    docs: dict[str, Document] = {}
    for split in cfg.corpus.splits:
        d, _ = load_split(cfg.paths.raw, split)
        docs.update(d)
    sel = select_corpus(docs, cfg.corpus.n_papers, cfg.corpus.min_mentions)
    papers = [docs[p] for p in sorted(sel.papers)]
    chunks: list[Chunk] = []
    for doc in papers:
        chunks.extend(chunk_document(doc, cfg.chunking, id_prefix=paper_prefix(doc.doc_id)))
    out = Path(cfg.paths.processed)
    write_jsonl(out / "papers.jsonl", papers)
    write_jsonl(out / "chunks.jsonl", chunks)
    metric_df = Counter(m for d in papers for m, n in metric_hits(d).items() if n >= cfg.corpus.min_mentions)
    toks = sorted(c.n_tokens for c in chunks)
    stats = {
        **sel.stats,
        "selection_order": sel.papers,
        "chunks": len(chunks),
        "chunk_tokens_median": toks[len(toks) // 2] if toks else 0,
        "chunk_tokens_p90": toks[int(0.9 * len(toks))] if toks else 0,
        "corpus_tokens": sum(toks),
        "metric_lexicon_df": dict(metric_df.most_common()),
        "paper_datasets": sel.mentions,
    }
    write_json(out / "corpus_manifest.json", manifest({"corpus_config": cfg.corpus.model_dump(), "chunking": cfg.chunking.model_dump(), "stats": stats}))
    return stats


def load_corpus(cfg: ProjectConfig) -> tuple[dict[str, Document], list[Chunk]]:
    out = Path(cfg.paths.processed)
    papers = {d["doc_id"]: Document(**d) for d in read_jsonl(out / "papers.jsonl")}
    chunks = [Chunk(**c) for c in read_jsonl(out / "chunks.jsonl")]
    if not papers:
        raise FileNotFoundError(f"no processed corpus in {out}; run `graphtrace corpus` first")
    return papers, chunks


def build_canonicalizer(cfg: ProjectConfig, papers: dict[str, Document]) -> Canonicalizer:
    curated = load_yaml(cfg.extraction.curated_aliases).get("aliases", {}) if Path(cfg.extraction.curated_aliases).exists() else {}
    texts = [p for d in papers.values() for s in d.sections for p in s.paragraphs]
    return Canonicalizer.build(curated, texts)


def build(cfg: ProjectConfig, score_confidence: bool | None = None) -> dict:
    papers, _ = load_corpus(cfg)
    records = read_jsonl(Path(cfg.paths.graph) / "extractions.jsonl")
    if not records:
        raise FileNotFoundError("no extractions found; run `graphtrace extract` first")
    canon = build_canonicalizer(cfg, papers)
    build_stats: dict = {}
    nodes, edges, results = build_graph(records, papers, canon, stats=build_stats)
    confidence: dict[str, float] = {}
    if score_confidence if score_confidence is not None else cfg.extraction.nli_confidence:
        from .support import score_edges
        from .verification import NLIVerifier

        confidence = score_edges(KG(nodes, edges, results), NLIVerifier(cfg.extraction.nli_model))
        edges = [e.model_copy(update={"confidence": confidence.get(e.edge_id)}) for e in edges]
    save_graph(cfg.paths.graph, nodes, edges, results, canon.merges)
    kg = KG(nodes, edges, results)
    stats = kg.stats()
    write_json(
        Path(cfg.paths.graph) / "graph_manifest.json",
        manifest(
            {
                "stats": stats,
                **build_stats,
                "alias_merges": len(canon.merges),
                "nli_model": cfg.extraction.nli_model if confidence else None,
                "min_edge_confidence_for_questions_and_methods": cfg.bench.min_edge_confidence,
                "edges_below_min_confidence": len(kg.edges) - len(kg.filtered(cfg.bench.min_edge_confidence).edges),
            }
        ),
    )
    return stats