"""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
from itertools import islice

from tools.knowledge.extraction import extract_batch, extraction_key, PROMPT
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) >= 64 or size + len(unit["text"]) > 16000):
            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.
    # Only the model's annotation view may be shortened; evidence never is.
    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 = {"reference", "lifecycle", "work_record", "meeting", "consultant"}
    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)
    # Reuse annotations independently of dates/line offsets. These are parsed
    # locally, not inferred by the LLM. Legacy checkpoints are accepted only
    # when their complete prompt/model/input fingerprint verifies.
    proposal_fields = {
        "concepts",
        "relations",
        "discarded_concepts",
        "discarded_relations",
        "annotation_input_limited",
    }
    proposals = {}
    for file in cache.glob("*.json"):
        previous = json.loads(file.read_text())
        inputs = [{k: v for k, v in row.items() if k not in proposal_fields} for row in previous]
        legacy = identifier(
            PROMPT,
            json.dumps(inputs, sort_keys=True, ensure_ascii=False),
            json.dumps(model_identity, sort_keys=True),
        )
        if file.stem not in {legacy, extraction_key(inputs, model_identity)}:
            continue
        for row in previous:
            proposals[extraction_key([row], model_identity)] = {
                k: v for k, v in row.items() if k in proposal_fields
            }

    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]:
                raise RuntimeError("Invalid extraction cache identity.")
            rows = [
                {**unit, **{k: v for k, v in row.items() if k in proposal_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_batch(missing, complete)} 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:
        pending = iter(batches(annotatable))
        while window := list(islice(pending, workers)):
            # Bounded submission: an API failure cannot leave hundreds of
            # queued requests running after a failed extraction.
            for rows, cached in pool.map(process, window):
                extracted.extend(rows)
                hits += int(cached)
                calls += int(not cached)
                if progress:
                    progress(calls + hits)
    # A batch can straddle two profiles. Give profiles with no accepted
    # annotation their own bounded pass rather than silently losing them.
    absent = {u["person"] for u in extracted} - {u["person"] for u in extracted if u["concepts"]}
    mixed = {
        u["person"]
        for batch in batches(annotatable)
        if len({r["person"] for r in batch}) > 1
        for u in batch
    }
    absent &= mixed
    by_id = {u["id"]: u for u in extracted}
    for person in sorted(absent):
        own = [u for u in annotatable if u["person"] == person and u["kind"] != "supervisor"][:32]
        if own:
            rows, cached = process(own)
            by_id.update({u["id"]: u for u in rows})
            calls += int(not cached)
            hits += int(cached)
    extracted = [by_id[u["id"]] for u in extracted]
    extracted.extend(
        {**u, "concepts": [], "relations": []} 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,
        "index_version": version,
        "model": model_identity,
        "source_scope": "published_student_profiles",
        "profiles": len(documents),
        "units": extracted,
        "graph_ready": graph,
    }
    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,
        "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)
    args = parser.parse_args()
    complete, identity = configured_extractor()
    print(
        json.dumps(
            build(
                args.wiki,
                args.data,
                complete,
                identity,
                progress=lambda count: (
                    print(f"Validated extraction batches: {count}", flush=True)
                    if count % 10 == 0
                    else None
                ),
            ),
            ensure_ascii=False,
        )
    )


if __name__ == "__main__":
    main()