"""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):
    from tools.sync.sync_indexes import source_files
    version = 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)]
    # 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)]
    annotatable=[u for u in bounded if u['kind'] not in {'reference','lifecycle'}]
    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 {'reference','lifecycle'})
    order={u['id']:i for i,u in enumerate(bounded)}
    extracted.sort(key=lambda u:order[u['id']])
    index = {'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))
    with publication_lock(data_dir, exclusive=True):
        if read_version(data_dir) != version or source_files(wiki_root) != sources:
            raise RuntimeError('Source changed during extraction; checkpoint not activated.')
        if graph:
            from tools.knowledge.store import write_graph
            write_graph(index)
        durable_json(data_dir / 'knowledge/index.json', index)
    return {'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()