"""File the questions that sit on a bare system, using retrieval then a chooser.

1,455 questions have no condition inside their own system, so nothing could be
promoted — they are genuinely unclassified below the system. General Pediatrics
alone holds 396 and has no subcategories at all, so its questions are almost
always about something that already exists elsewhere in the tree.

Two stages, deliberately:

  1. *Retrieval* ranks existing categories against the question's stored
     embedding and keeps the nearest few.
  2. *A chooser* picks one of those, or says NONE.

The model never sees the whole tree and never writes a category name — it returns
an id from the shortlist, and anything not on that shortlist is discarded. So a
category it invents cannot survive, exactly as with AI-mode citations. Nothing
new is created either: this files questions into categories an educator already
has.

A question whose best candidate is weak is left alone rather than forced, and the
old system is kept as a cross-link so no existing filter narrows.

`--force` files the leftovers on the nearest candidate even when the chooser
declines, so nothing is left sitting on a bare system. It is a weaker signal than
a chosen match, so those questions are marked for review — but the system is
still kept as a cross-link, so a wrong pick narrows nothing.

    docker compose exec backend python -m scripts.classify_unfiled_questions
    docker compose exec backend python -m scripts.classify_unfiled_questions --limit 50
    docker compose exec backend python -m scripts.classify_unfiled_questions --apply
    docker compose exec backend python -m scripts.classify_unfiled_questions --apply --force
"""
import json
import sys

from sqlalchemy import text as sa_text

from app.config import settings
from app.database import SessionLocal
from app.services.embedding_service import generate_embedding

SHORTLIST = 8
# Below this cosine the nearest category is not a real match, so do not ask.
MIN_SIMILARITY = 0.45

PROMPT = """You are filing a pediatric exam question into an existing category.

Question:
{question}

Candidate categories (choose the single best fit):
{candidates}

Reply with only the numeric id of the best category, or NONE if none of them fit
the question's main topic. Do not invent a category or explain."""


def cosine(a, b):
    dot = sum(x * y for x, y in zip(a, b))
    na = sum(x * x for x in a) ** 0.5
    nb = sum(y * y for y in b) ** 0.5
    return dot / (na * nb) if na and nb else 0.0


def category_paths(db):
    """Every category below the top level, labelled with its full trail."""
    rows = db.execute(sa_text("SELECT id, name, parent_id FROM question_categories")).fetchall()
    name = {cid: n for cid, n, _p in rows}
    parent = {cid: p for cid, _n, p in rows}

    def trail(cid):
        parts, cursor, guard = [], cid, 0
        while cursor is not None and guard < 6:
            parts.append(name[cursor])
            cursor, guard = parent.get(cursor), guard + 1
        return " > ".join(reversed(parts))

    return {cid: trail(cid) for cid, _n, p in rows if p is not None}


def choose(question_text, candidates):
    """Ask the model to pick one id from the shortlist. Returns an id or None."""
    import httpx

    listing = "\n".join(f"{cid}: {path}" for cid, path, _score in candidates)
    body = {
        "model": settings.LITELLM_MODEL,
        "messages": [{"role": "user", "content": PROMPT.format(
            question=question_text[:1500], candidates=listing)}],
        "temperature": 0,
        "max_tokens": 12,
    }
    base = (settings.LITELLM_API_BASE or "").rstrip("/").removesuffix("/v1")
    response = httpx.post(
        f"{base}/v1/chat/completions",
        headers={"Authorization": f"Bearer {settings.LITELLM_API_KEY}"},
        json=body, timeout=45,
    )
    response.raise_for_status()
    reply = response.json()["choices"][0]["message"]["content"].strip()
    allowed = {cid for cid, _p, _s in candidates}
    for token in reply.replace(":", " ").split():
        if token.isdigit() and int(token) in allowed:
            return int(token)
    return None  # NONE, or anything not on the shortlist


def main():
    apply_changes = "--apply" in sys.argv
    force = "--force" in sys.argv
    limit = None
    if "--limit" in sys.argv:
        limit = int(sys.argv[sys.argv.index("--limit") + 1])

    db = SessionLocal()
    try:
        paths = category_paths(db)
        systems = {row[0] for row in db.execute(sa_text(
            "SELECT id FROM question_categories WHERE parent_id IS NULL")).fetchall()}

        pending = db.execute(sa_text("""
            SELECT q.id, q.question_text, q.question_category_id
            FROM questions q
            WHERE q.question_category_id IN (SELECT id FROM question_categories WHERE parent_id IS NULL)
            ORDER BY q.id
        """)).fetchall()
        if limit:
            pending = pending[:limit]
        print(f"  questions on a bare system: {len(pending)}")

        print("  embedding category labels…", flush=True)
        vectors = {}
        for cid, path in paths.items():
            embedding = generate_embedding(path)
            if embedding:
                vectors[cid] = embedding
        print(f"  categories embedded: {len(vectors)}", flush=True)

        filed, skipped, forced, no_candidates = [], [], [], 0
        for index, (qid, text, system_id) in enumerate(pending, start=1):
            row = db.execute(sa_text(
                "SELECT embedding FROM questions WHERE id = :q"), {"q": qid}).scalar()
            if row is None:
                no_candidates += 1
                continue
            question_vector = [float(x) for x in str(row).strip("[]").split(",")]

            ranked = sorted(
                ((cid, paths[cid], cosine(question_vector, vec)) for cid, vec in vectors.items()),
                key=lambda item: -item[2],
            )[:SHORTLIST]
            if not ranked:
                no_candidates += 1
                continue
            if ranked[0][2] < MIN_SIMILARITY and not force:
                no_candidates += 1
                continue

            try:
                picked = choose(text, ranked)
            except Exception as error:
                print(f"    chooser failed on question {qid}: {error}")
                continue
            if picked is None or picked in systems:
                if not force:
                    skipped.append((qid, ranked[0][1]))
                    continue
                # Nearest neighbour, on the same shortlist the chooser saw.
                picked = ranked[0][0]
                forced.append(qid)
            filed.append((qid, picked, paths[picked], system_id))

            if apply_changes:
                db.execute(sa_text(
                    "UPDATE questions SET question_category_id = :c WHERE id = :q"),
                    {"c": picked, "q": qid})
                db.execute(sa_text("""
                    INSERT INTO question_category_links (question_id, category_id)
                    SELECT :q, :sys WHERE NOT EXISTS (
                        SELECT 1 FROM question_category_links
                        WHERE question_id = :q AND category_id = :sys)
                """), {"q": qid, "sys": system_id})
                if index % 25 == 0:
                    db.commit()
                    print(f"    {index}/{len(pending)} processed, {len(filed)} filed", flush=True)

        if apply_changes:
            db.commit()

        print(f"\n{'APPLIED' if apply_changes else 'DRY RUN'}")
        print(f"  filed          : {len(filed)}")
        print(f"  of those forced: {len(forced)}")
        print(f"  chooser said no: {len(skipped)}")
        print(f"  no candidate   : {no_candidates}\n")
        print("  Sample of what was chosen:")
        for qid, _cid, path, _sys in filed[:15]:
            print(f"    q{qid} -> {path}")
        if not apply_changes:
            print("\n  Re-run with --apply to write these.")
    finally:
        db.close()


if __name__ == "__main__":
    sys.exit(main())