"""CLI: python -m tools.knowledge.build (same private .env as the backend)."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, wait, FIRST_COMPLETED

from tools.knowledge.extraction import extract_batch, extraction_key, ANNOTATION_FIELDS, MAX_BATCH_UNITS, MAX_BATCH_CHARS, ExtractionRejected
from tools.knowledge.evidence import EXTRACTION_SCHEMA, EXTRACTION_POLICY, STRUCTURAL_KINDS
from tools.knowledge.sources import PROFILE, profile_units, identifier
from tools.rag.common import ZPWIKI_ROOT, DATA_DIR, load_zpwiki_page
from tools.sync.index_state import read_version, durable_json, publication_lock


def batches(units):
    batch, size = [], 0
    for unit in units:
        if batch and (len(batch) >= MAX_BATCH_UNITS or size + len(unit["text"]) > MAX_BATCH_CHARS
                      or (unit["source_path"], unit.get("work_label")) != (batch[0]["source_path"], batch[0].get("work_label"))):
            yield batch
            batch, size = [], 0
        batch.append(unit)
        size += len(unit["text"])
    if batch:
        yield batch


def build(
    wiki_root: Path,
    data_dir: Path,
    complete,
    model_identity: dict,
    *,
    graph=True,
    vectors=True,
    workers=2,
    progress=None,
) -> dict:
    from tools.rag.rebuild_index import acquire_reindex_lock

    with acquire_reindex_lock(data_dir / "knowledge/.build.lock"):
        return _build(
            wiki_root,
            data_dir,
            complete,
            model_identity,
            graph=graph,
            vectors=vectors,
            workers=workers,
            progress=progress,
        )


def _build(wiki_root, data_dir, complete, model_identity, *, graph, vectors, workers, progress):
    import tempfile
    import shutil
    from tools.sync.index_state import index_snapshot
    from tools.knowledge.publication import checkpoint_state, publish_checkpoint, PENDING_FILE
    from tools.knowledge.store import create_constraints
    from tools.graphrag.graph_db import create_neo4j_driver, get_neo4j_settings
    from tools.sync.sync_indexes import source_files
    with index_snapshot(data_dir) as version:
        if version is None:
            raise RuntimeError("Build the shared source indexes before content extraction.")
    index, summary = prepare_index(wiki_root, data_dir, complete, model_identity,
        graph=graph, vectors=vectors, workers=workers, progress=progress, source_version=version)
    stage = Path(tempfile.mkdtemp(prefix=".knowledge-stage-", dir=data_dir))
    try:
        checkpoint_state(stage, index)
        with publication_lock(data_dir, exclusive=True):
            source_digest=hashlib.sha256(json.dumps(source_files(wiki_root),sort_keys=True).encode()).hexdigest()
            if read_version(data_dir) != version or source_digest != version['source_sha256']:
                raise RuntimeError("Source changed during extraction; checkpoint not activated.")
            if graph:
                settings = get_neo4j_settings()
                with create_neo4j_driver(settings) as driver, driver.session(database=settings.database) as session:
                    driver.verify_connectivity()
                    create_constraints(session)
                    publish_checkpoint(data_dir, stage, index, session)
            else:
                publish_checkpoint(data_dir, stage, index)
        return summary
    finally:
        if not (data_dir / PENDING_FILE).exists():
            shutil.rmtree(stage)


def prepare_index(wiki_root, data_dir, complete, model_identity, *, graph=True, vectors=True, workers=2, progress=None, source_version=None):
    from tools.sync.sync_indexes import source_files

    version = source_version or read_version(data_dir)
    if not version:
        raise RuntimeError("Build the shared source indexes before content extraction.")
    sources = source_files(wiki_root)
    digest = hashlib.sha256(json.dumps(sources, sort_keys=True).encode()).hexdigest()
    if digest != version["source_sha256"]:
        raise RuntimeError("Wiki differs from the active index; synchronize first.")
    documents = []
    for path in sources:
        if PROFILE.fullmatch(path):
            document = load_zpwiki_page(wiki_root / path, wiki_root=wiki_root)
            raw = (wiki_root / path).read_text()
            document["body_line_offset"] = (
                raw[: raw.index(document["content"])].count("\n") if document["content"] else 0
            )
            documents.append(document)
    documents = [d for d in documents if d.get("published") is not False]
    if not documents:
        raise RuntimeError("No published student profiles to extract.")
    units = [unit for d in documents for unit in profile_units(d, include_records=True)]
    # Keep whole statements, including trailing negations/qualifications.
    # Over-budget source units are retained with an explicit uncertain decision.
    previous = data_dir / "knowledge/index.json"
    prior = json.loads(previous.read_text()).get("units", []) if previous.exists() else []

    def source_identity(unit):
        return (unit["source_path"], unit.get("work_label"), unit["line"], unit["text"])

    prior_ids = {source_identity(u): u["id"] for u in prior}
    # Inserting a status note must not invalidate every subsequent profile's
    # model cache just because its position in the corpus has shifted.
    bounded = [
        {
            **unit,
            "id": prior_ids.get(source_identity(unit), identifier(unit["id"], i, unit["text"])),
        }
        for i, unit in enumerate(units)
    ]
    local_kinds = STRUCTURAL_KINDS | {"lifecycle"}
    annotatable = [u for u in bounded if u["kind"] not in local_kinds]
    cache = data_dir / "knowledge/cache"
    cache.mkdir(parents=True, exist_ok=True)
    # Every cached decision is bound to this policy, verifier, complete source
    # scope, actor, dates and model identity. Legacy omissions are never reused.
    proposals = {}
    for file in cache.glob("*.json"):
        try:
            previous = json.loads(file.read_text())
            if (not isinstance(previous,list) or not previous
                    or any(row.get("annotation_policy") != EXTRACTION_POLICY
                           or row.get("annotation_status") not in {"knowledge","irrelevant","uncertain"}
                           for row in previous)
                    or file.stem != extraction_key(previous,model_identity)):
                continue
        except (ValueError, OSError, AttributeError, TypeError):
            continue
        for row in previous:
            proposals[extraction_key([row],model_identity)] = {k:v for k,v in row.items() if k in ANNOTATION_FIELDS}
    import threading
    requests = 0
    request_lock = threading.Lock()
    def counted_complete(messages):
        nonlocal requests
        with request_lock:
            requests += 1
        return complete(messages)

    def extract_complete_batch(batch):
        try:
            return extract_batch(batch,counted_complete)
        except (ExtractionRejected, RuntimeError) as error:
            if isinstance(error, RuntimeError) and str(error) != "Incomplete answer model output":
                raise
            if len(batch) == 1:
                if isinstance(error, ExtractionRejected):
                    raise ExtractionRejected([f"{batch[0]['source_path']}:{batch[0]['line']}",str(error)]) from error
                raise
            # Every child still needs a decision for every input. A missing ID
            # can never turn into an empty successful annotation.
            middle = len(batch)//2
            return [*extract_complete_batch(batch[:middle]),*extract_complete_batch(batch[middle:])]

    def process(batch):
        key = extraction_key(batch, model_identity)
        file = cache / (key + ".json")
        if file.exists():
            rows = json.loads(file.read_text())
            # Cache is private and keyed by complete source, prompt and model.
            if ([r["id"] for r in rows] != [r["id"] for r in batch]
                    or any(r.get("annotation_policy") != EXTRACTION_POLICY or r.get("annotation_status") not in {"knowledge","irrelevant","uncertain"}
                           or extraction_key([r],model_identity) != extraction_key([u],model_identity)
                           for r,u in zip(rows,batch))):
                raise RuntimeError("Invalid extraction cache identity.")
            rows = [
                {**unit, **{k: v for k, v in row.items() if k in ANNOTATION_FIELDS}}
                for unit, row in zip(batch, rows)
            ]
            cached = True
        else:
            missing = [u for u in batch if extraction_key([u], model_identity) not in proposals]
            fresh = {u["id"]: u for u in extract_complete_batch(missing)} if missing else {}
            rows = [
                fresh.get(u["id"]) or {**u, **proposals[extraction_key([u], model_identity)]}
                for u in batch
            ]
            durable_json(file, rows)
            cached = not missing
        return rows, cached

    extracted, calls, hits = [], 0, 0
    if not 1 <= workers <= 4:
        raise ValueError("Extraction workers must be between 1 and 4.")
    with ThreadPoolExecutor(max_workers=workers) as pool:
        source_batches = iter(batches(annotatable))
        pending = set()
        for _ in range(workers):
            if (batch := next(source_batches, None)) is not None:
                pending.add(pool.submit(process, batch))
        while pending:
            done, pending = wait(pending, return_when=FIRST_COMPLETED)
            # At most `workers` requests are in flight. A slow batch does not
            # block starting the next one; failures never queue the whole corpus.
            completed = [future.result() for future in done]
            for rows, cached in completed:
                extracted.extend(rows)
                hits += int(cached)
                calls += int(not cached)
                if progress:
                    progress(calls + hits)
                if (batch := next(source_batches, None)) is not None:
                    pending.add(pool.submit(process, batch))
    extracted.extend(
        {**u, "concepts": [], "relations": [], "annotation_status": "structural",
         "annotation_reason": "Literal structural, role, reference or lifecycle record",
         "annotation_policy": EXTRACTION_POLICY}
        for u in bounded if u["kind"] in local_kinds
    )
    order = {u["id"]: i for i, u in enumerate(bounded)}
    extracted.sort(key=lambda u: order[u["id"]])
    from tools.knowledge.academic import completion_event

    for unit in extracted:
        if event := completion_event(unit):
            unit.update(
                completion_status=event["status"],
                completion_scope=event["scope"],
                completion_event_year=event["event_year"],
                completion_event_date=event["event_date"],
            )
    index = {
        "schema": 1,
        "academic_schema": 1,
        "extraction_schema": EXTRACTION_SCHEMA,
        "extraction_policy": EXTRACTION_POLICY,
        "index_version": version,
        "model": model_identity,
        "source_scope": "published_student_profiles",
        "profiles": len(documents),
        "units": extracted,
        "graph_ready": graph,
    }
    from collections import Counter, defaultdict
    decision_counts = Counter(u["annotation_status"] for u in extracted)
    per_scope = defaultdict(Counter)
    for unit in extracted:
        per_scope[(unit["person"],unit.get("work_label"),unit["phase"])][unit["annotation_status"]] += 1
    index["coverage"] = {"decisions":dict(decision_counts), "units":len(extracted),
        "by_person_work_phase":[{"person":key[0],"work_label":key[1],"phase":key[2],"decisions":dict(value)}
                                for key,value in sorted(per_scope.items(),key=lambda pair:str(pair[0]))]}
    if len(extracted) != len(bounded) or {u["id"] for u in extracted} != {u["id"] for u in bounded}:
        raise RuntimeError("Missing or duplicate decisions in extraction checkpoint")
    if vectors:
        index.update(embedding_index(extracted, data_dir))
    index["version"] = identifier(json.dumps(index, ensure_ascii=False, sort_keys=True))
    if source_files(wiki_root) != sources:
        raise RuntimeError("Source changed during extraction; checkpoint not activated.")
    return index, {

        "profiles": len(documents),
        "evidence_units": len(extracted),
        "llm_batches": calls,
        "cached_batches": hits,
        "model_requests": requests,
        "decisions": dict(decision_counts),
        "uncertain_relations": sum(r.get("uncertain_relations",0) for r in extracted),
        "concepts": len({c.casefold() for r in extracted for c in r["concepts"]}),
        "discarded_concepts": sum(r.get("discarded_concepts", 0) for r in extracted),
        "discarded_relations": sum(r.get("discarded_relations", 0) for r in extracted),
        "version": index["version"],
        "graph_ready": graph,
    }


def embedding_index(units: list[dict], data_dir: Path) -> dict:
    import base64
    import zlib
    import numpy as np
    from tools.rag.embedding_cache import EmbeddingCache, cache_key
    from tools.rag.embedding_utils import (
        embedding_model_name,
        model_profile,
        split_passage,
        embed_passages,
    )

    profile = model_profile(embedding_model_name())
    cache = EmbeddingCache(data_dir / "knowledge/embedding_cache.sqlite")
    inputs = [split_passage(u["text"]) for u in units]
    keys = [[cache_key(profile, s) for s in pieces] for pieces in inputs]
    found = cache.load([key for group in keys for key in group])
    shared = data_dir / "embedding_cache.sqlite"
    if shared.exists():
        import sqlite3

        try:
            found.update(
                EmbeddingCache(shared, readonly=True).load(
                    [key for group in keys for key in group if key not in found]
                )
            )
        except sqlite3.Error:
            pass  # A cache is an optimization, never the source of truth.
    missing = {
        key: text
        for group, pieces in zip(keys, inputs)
        for key, text in zip(group, pieces)
        if key not in found
    }
    if missing:
        fresh = dict(zip(missing, embed_passages(list(missing.values()))))
        cache.save(fresh)
        found.update(fresh)
    vectors = []
    for group in keys:
        value = np.mean([found[key] for key in group], axis=0)
        value /= np.linalg.norm(value)
        vectors.append(value)
    matrix = np.asarray(vectors, dtype="<f4")
    return {
        "embedding_fingerprint": profile["fingerprint"],
        "vector_shape": list(matrix.shape),
        "vectors": base64.b64encode(zlib.compress(matrix.tobytes())).decode("ascii"),
    }


def configured_extractor():
    from functools import partial
    from tools.answers.model_client import complete_with_configured_model

    provider = os.getenv("ANSWER_PROVIDER", "openai_compatible")
    legacy = bool(os.getenv("OPENWEBUI_API_KEY"))
    identity = {
        "provider": provider,
        "model": os.getenv("ANSWER_MODEL")
        or ("model120-fast" if provider == "openai_compatible" and legacy else ""),
        "endpoint": os.getenv("ANSWER_LLM_URL")
        or (
            "https://api.anthropic.com/v1/messages"
            if provider == "anthropic"
            else "https://ui.tukekemt.xyz/api/chat/completions" if legacy else ""
        ),
        "max_tokens": 4096,
    }
    return partial(complete_with_configured_model, max_tokens=4096), identity


def refresh_existing(wiki_root: Path, data_dir: Path) -> dict | None:
    """Once enabled by an initial build, keep content aligned on later syncs.

    Cached unchanged inputs need no LLM calls. A failed refresh leaves content
    unavailable instead of serving the previous snapshot as current evidence.
    """
    if not (data_dir / "knowledge/index.json").exists():
        return None
    complete, identity = configured_extractor()
    return build(wiki_root, data_dir, complete, identity)


def main():
    # Dotenv is intentionally local to the CLI; services already load env.
    from dotenv import load_dotenv

    load_dotenv(Path(__file__).resolve().parents[2] / ".env", override=False)
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--wiki", type=Path, default=Path(os.getenv("ZPWIKI_ROOT", str(ZPWIKI_ROOT)))
    )
    parser.add_argument("--data", type=Path, default=DATA_DIR)
    parser.add_argument("--workers", type=int, choices=range(1,5), default=2)
    args = parser.parse_args()
    complete, identity = configured_extractor()
    print(
        json.dumps(
            build(
                args.wiki,
                args.data,
                complete,
                identity,
                workers=args.workers,
                progress=lambda count: (
                    print(f"Validated extraction batches: {count}", flush=True)
                    if count % 10 == 0
                    else None
                ),
            ),
            ensure_ascii=False,
        )
    )


if __name__ == "__main__":
    main()