"""
RAG Quality Audit for ModuleDesk AI Draft feature.

Measures RETRIEVAL quality separately from draft quality.
For each eval ticket, determines whether the retrieval returned a genuinely
relevant precedent, and classifies each ticket into:
  (a) bad-retrieval   — a relevant precedent EXISTS but wasn't retrieved (or only irrelevant ones were)
  (b) good-retrieval-weak-draft — relevant precedent WAS retrieved but draft still diverged
  (c) no-precedent    — no genuinely similar resolved ticket exists in the corpus

Brute-force scan: for each ticket, we embed the inbound question and cosine-compare
against ALL resolved thread embeddings to detect hidden relevant precedents.

Usage:
    venv/bin/python scripts/eval_rag_quality.py [--n 8] [--schema tenant_internal]
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

# ---------------------------------------------------------------------------
# Pricing (for embedding/judge cost tracking)
# ---------------------------------------------------------------------------
USD_EUR = 0.92
PRICE_USD_PER_1M = {
    "gpt-4o-mini":              {"in": 0.15, "out": 0.60},
    "text-embedding-3-small":   {"in": 0.02, "out": 0.00},
}

def tokens_to_eur(model: str, prompt_tokens: int, completion_tokens: int) -> float:
    p = PRICE_USD_PER_1M.get(model, {"in": 0.0, "out": 0.0})
    usd = (prompt_tokens * p["in"] + completion_tokens * p["out"]) / 1_000_000
    return usd * USD_EUR


@dataclass
class AuditCosts:
    calls: List[Dict] = field(default_factory=list)

    def record(self, call_type: str, model: str, prompt_tokens: int, completion_tokens: int) -> float:
        cost = tokens_to_eur(model, prompt_tokens, completion_tokens)
        self.calls.append({"type": call_type, "model": model, "pt": prompt_tokens, "ct": completion_tokens, "cost": cost})
        return cost

    def total(self) -> float:
        return sum(c["cost"] for c in self.calls)


# ---------------------------------------------------------------------------
# LLM relevance judge
# ---------------------------------------------------------------------------

RELEVANCE_JUDGE_SYSTEM = """\
You are an expert support ticket analyst. You will be given two customer support questions:
- Question A: a new (incoming) support question
- Question B: a previously-resolved support question from the knowledge base

Your task: determine whether Question B is GENUINELY RELEVANT to Question A.
Genuinely relevant means: if the answer to B were shown to the agent, it would materially
help them answer A (same problem, same module feature, same error, same workflow).

Answer with JSON ONLY:
{
  "relevant": true or false,
  "confidence": 0-100 (how confident you are),
  "reason": "one sentence"
}

Be strict. A vague thematic overlap is NOT relevant. Relevance requires the same concrete problem.
"""


def judge_relevance(
    question_a: str,
    question_b: str,
    client,
    costs: AuditCosts,
) -> Dict[str, Any]:
    """Ask gpt-4o-mini whether question_b is relevant to question_a."""
    user_content = (
        f"Question A (new ticket):\n---\n{question_a[:1000]}\n---\n\n"
        f"Question B (precedent):\n---\n{question_b[:1000]}\n---\n\n"
        "Respond with JSON only."
    )
    try:
        import openai as _openai
        resp = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": RELEVANCE_JUDGE_SYSTEM},
                {"role": "user", "content": user_content},
            ],
            temperature=0.0,
            max_tokens=150,
            response_format={"type": "json_object"},
        )
        usage = resp.usage
        if usage:
            costs.record("judge_relevance", "gpt-4o-mini", usage.prompt_tokens, usage.completion_tokens)
        result = json.loads(resp.choices[0].message.content)
        return {
            "relevant": bool(result.get("relevant", False)),
            "confidence": int(result.get("confidence", 0)),
            "reason": str(result.get("reason", "")),
        }
    except Exception as exc:
        return {"relevant": False, "confidence": 0, "reason": f"Judge failed: {exc}"}


# ---------------------------------------------------------------------------
# Main audit
# ---------------------------------------------------------------------------

def run_rag_audit(n: int = 8, schema: str = "tenant_internal") -> None:
    print(f"\n{'='*65}")
    print(f"ModuleDesk RAG Quality Audit")
    print(f"Schema: {schema} | N tickets: {n}")
    print(f"{'='*65}\n")

    from supporthub.app.config import config
    from supporthub.app.db import session_scope, set_worker_tenant_schema

    if not config.AI_API_KEY:
        print("ERROR: AI_API_KEY not set.")
        sys.exit(1)

    set_worker_tenant_schema(schema)

    import openai
    client = openai.OpenAI(api_key=config.AI_API_KEY, base_url=config.AI_BASE_URL)
    costs = AuditCosts()

    # Load embedding index
    from supporthub.app.services.embedding_service import embedding_index, EmbeddingService
    print("Loading embedding index...")
    embedding_index.load_from_db(schema=schema)
    print(f"  Loaded {embedding_index.count} thread embeddings\n")

    embed_svc = EmbeddingService()
    # Use tracked client
    embed_svc.client = client

    # Pull eval tickets (same logic as eval_draft_quality.py)
    from scripts.eval_draft_quality import select_eval_tickets
    print(f"Selecting {n} eval tickets from {schema}...")
    with session_scope(schema=schema) as session:
        tickets = select_eval_tickets(session, n=n)
        session.expunge_all()
    print(f"  Selected {len(tickets)} tickets\n")

    if not tickets:
        print("ERROR: No eligible tickets found.")
        sys.exit(1)

    # Build brute-force corpus: all resolved threads with embeddings + outbound replies
    # "resolved" = has at least one outbound message (same as what get_similar_threads requires)
    print("Building brute-force corpus from embedding index...")
    with session_scope(schema=schema) as session:
        from sqlalchemy import text
        # Get all threads that have an outbound reply AND an embedding
        corpus_rows = session.execute(text("""
            SELECT DISTINCT te.thread_id, te.summary_text,
                   (SELECT body_text FROM addons_messages
                    WHERE thread_id = te.thread_id AND direction = 'inbound'
                    ORDER BY created_at ASC LIMIT 1) AS first_inbound,
                   (SELECT body_text FROM addons_messages
                    WHERE thread_id = te.thread_id AND direction = 'outbound'
                    ORDER BY created_at ASC LIMIT 1) AS first_outbound
            FROM thread_embeddings te
            JOIN addons_messages m ON m.thread_id = te.thread_id AND m.direction = 'outbound'
            WHERE te.embedding_blob IS NOT NULL
        """)).fetchall()

    corpus = {}
    for row in corpus_rows:
        if row.first_inbound and row.first_outbound:
            corpus[row.thread_id] = {
                "thread_id": row.thread_id,
                "summary_text": row.summary_text or "",
                "first_inbound": row.first_inbound[:800],
                "first_outbound": row.first_outbound[:500],
            }

    corpus_thread_ids = list(corpus.keys())
    print(f"  Corpus size: {len(corpus_thread_ids)} threads with embeddings + outbound replies\n")

    # Build a matrix of all corpus vectors for brute-force scan
    with embedding_index._lock:
        idx_thread_ids = list(embedding_index._thread_ids)
        idx_vectors = embedding_index._vectors.copy() if embedding_index._vectors is not None else None

    if idx_vectors is None or len(idx_thread_ids) == 0:
        print("ERROR: Embedding index is empty.")
        sys.exit(1)

    # Map thread_id -> vector index position
    tid_to_pos = {tid: i for i, tid in enumerate(idx_thread_ids)}

    # Build corpus-only matrix (only threads that have outbound replies)
    corpus_positions = [tid_to_pos[tid] for tid in corpus_thread_ids if tid in tid_to_pos]
    corpus_tids_ordered = [corpus_thread_ids[i] for i, tid in enumerate(corpus_thread_ids) if tid in tid_to_pos]

    if not corpus_positions:
        print("ERROR: No corpus threads found in embedding index.")
        sys.exit(1)

    corpus_matrix = idx_vectors[corpus_positions]  # (C, D)
    # Normalize for cosine similarity
    norms = np.linalg.norm(corpus_matrix, axis=1, keepdims=True) + 1e-10
    corpus_matrix_normed = corpus_matrix / norms
    print(f"  Corpus matrix: {corpus_matrix_normed.shape}\n")

    # Process each eval ticket
    results = []
    BRUTE_FORCE_THRESHOLD = 0.70    # sim >= this → "relevant precedent exists"
    RETRIEVAL_THRESHOLD = config.RAG_SIMILARITY_THRESHOLD  # what get_similar_threads uses
    TOP_K_BRUTE = 5                 # top-K from brute-force to judge

    for i, item in enumerate(tickets, 1):
        print(f"[{i}/{len(tickets)}] ticket={item['ticket_id']} thread={item['thread_id']} | {item['product_name']}")

        # Load inbound text
        with session_scope(schema=schema) as session:
            from sqlalchemy import text
            irow = session.execute(text(
                "SELECT body_text FROM addons_messages WHERE id = :id"
            ), {"id": item["inbound_msg_id"]}).one_or_none()

        if not irow or not irow.body_text:
            print("  SKIP: no inbound body")
            continue

        inbound_text = irow.body_text

        # 1. Embed the inbound question
        try:
            resp = client.embeddings.create(
                input=inbound_text[:8000],
                model=config.EMBEDDING_MODEL,
                dimensions=config.EMBEDDING_DIMENSIONS,
            )
            if resp.usage:
                costs.record("embed_query", config.EMBEDDING_MODEL, resp.usage.prompt_tokens, 0)
            query_vec = np.array(resp.data[0].embedding, dtype=np.float32)
        except Exception as exc:
            print(f"  ERROR embedding query: {exc}")
            continue

        query_norm = query_vec / (np.linalg.norm(query_vec) + 1e-10)

        # Load product_id for this ticket (needed for the product-aware floor in get_similar_threads)
        with session_scope(schema=schema) as session:
            from sqlalchemy import text as _text2
            from supporthub.app.models import AddonsThreads
            th = session.query(AddonsThreads).filter_by(id=item["thread_id"]).one_or_none()
            ticket_product_id = th.id_product if th else None

        # 2. Run standard retrieval (what the app actually uses — WITH product_id, matching live behavior)
        similar = embed_svc.get_similar_threads(
            index=embedding_index,
            query_text=None,
            query_vec=query_vec,
            product_id=ticket_product_id,  # use real product_id to match live app behavior
            exclude_thread_id=item["thread_id"],
            top_k=5,
        )

        retrieved_thread_ids = [s["thread_id"] for s in similar]
        retrieved_scores = {s["thread_id"]: s["similarity"] for s in similar}
        print(f"  Retrieved {len(similar)} threads from standard retrieval")
        for s in similar:
            print(f"    thread={s['thread_id']} sim={s['similarity']}% | {s.get('customer_asked','')[:60]}")

        # 3. Brute-force scan: top-K from entire corpus (excluding the ticket's own thread)
        exclude_set = {item["thread_id"]}
        sims_all = corpus_matrix_normed @ query_norm  # (C,)

        # Rank all corpus threads
        ranked_indices = np.argsort(sims_all)[::-1]
        brute_top: List[Tuple[int, float]] = []
        for idx in ranked_indices:
            tid = corpus_tids_ordered[idx]
            if tid in exclude_set:
                continue
            score = float(sims_all[idx])
            brute_top.append((tid, score))
            if len(brute_top) >= TOP_K_BRUTE:
                break

        print(f"  Brute-force top-{TOP_K_BRUTE}:")
        for tid, sc in brute_top:
            print(f"    thread={tid} sim={sc:.3f} | {corpus.get(tid, {}).get('first_inbound','')[:60]}")

        # 4. Judge relevance:
        # 4a. Are any RETRIEVED threads relevant?
        retrieved_relevant = []
        for s in similar[:3]:  # check top-3 retrieved
            tid = s["thread_id"]
            if tid not in corpus:
                continue
            precedent_q = corpus[tid]["first_inbound"]
            verdict = judge_relevance(inbound_text, precedent_q, client, costs)
            print(f"  JUDGE retrieved thread={tid}: relevant={verdict['relevant']} conf={verdict['confidence']} | {verdict['reason'][:70]}")
            if verdict["relevant"]:
                retrieved_relevant.append({"thread_id": tid, "similarity": s["similarity"], **verdict})

        # 4b. Are any BRUTE-FORCE top threads relevant (to detect missed precedents)?
        brute_relevant = []
        brute_threads_to_judge = [(tid, sc) for tid, sc in brute_top if tid not in retrieved_scores]
        for tid, sc in brute_threads_to_judge[:3]:  # check up to 3 un-retrieved brute-force hits
            if tid not in corpus:
                continue
            precedent_q = corpus[tid]["first_inbound"]
            verdict = judge_relevance(inbound_text, precedent_q, client, costs)
            print(f"  JUDGE brute thread={tid} (sim={sc:.3f}): relevant={verdict['relevant']} conf={verdict['confidence']} | {verdict['reason'][:70]}")
            if verdict["relevant"]:
                brute_relevant.append({"thread_id": tid, "brute_sim": sc, **verdict})

        # 5. Classify into bucket
        has_retrieved_relevant = len(retrieved_relevant) > 0
        has_brute_missed = len(brute_relevant) > 0

        if has_brute_missed and not has_retrieved_relevant:
            bucket = "bad-retrieval"
            bucket_note = f"Missed relevant thread(s): {[b['thread_id'] for b in brute_relevant]}"
        elif has_retrieved_relevant:
            bucket = "good-retrieval-weak-draft"
            bucket_note = f"Relevant thread(s) retrieved: {[r['thread_id'] for r in retrieved_relevant]}"
        else:
            bucket = "no-precedent"
            bucket_note = "No genuinely relevant precedent found in corpus"

        print(f"  => BUCKET: {bucket} | {bucket_note}\n")

        results.append({
            "ticket_id": item["ticket_id"],
            "thread_id": item["thread_id"],
            "product_name": item["product_name"],
            "inbound_preview": inbound_text[:200],
            "n_retrieved": len(similar),
            "retrieved_thread_ids": retrieved_thread_ids,
            "retrieved_top_sim": similar[0]["similarity"] if similar else None,
            "n_retrieved_relevant": len(retrieved_relevant),
            "n_brute_missed": len(brute_relevant),
            "retrieved_relevant": retrieved_relevant,
            "brute_missed": brute_relevant,
            "bucket": bucket,
            "bucket_note": bucket_note,
        })

    # ---------------------------------------------------------------------------
    # Aggregate
    # ---------------------------------------------------------------------------
    bucket_counts = {"bad-retrieval": 0, "good-retrieval-weak-draft": 0, "no-precedent": 0}
    for r in results:
        bucket_counts[r["bucket"]] = bucket_counts.get(r["bucket"], 0) + 1

    total_cost = costs.total()
    print(f"\n{'='*65}")
    print("RAG QUALITY AUDIT RESULTS")
    print(f"{'='*65}")
    print(f"Tickets audited: {len(results)}")
    print(f"\nBucket split:")
    for b, c in bucket_counts.items():
        pct = c / len(results) * 100 if results else 0
        print(f"  {b}: {c}/{len(results)} ({pct:.0f}%)")

    print(f"\nTotal audit cost: €{total_cost:.5f}")
    print(f"Cost breakdown:")
    call_agg: Dict[str, float] = {}
    for c in costs.calls:
        key = f"{c['type']}:{c['model']}"
        call_agg[key] = call_agg.get(key, 0.0) + c["cost"]
    for k, v in sorted(call_agg.items()):
        print(f"  {k}: €{v:.5f}")

    print(f"\nPer-ticket detail:")
    for r in results:
        print(f"  ticket={r['ticket_id']} [{r['bucket']}] retrieved={r['n_retrieved']} (relevant={r['n_retrieved_relevant']}) missed={r['n_brute_missed']}")
        print(f"    {r['bucket_note']}")

    # ---------------------------------------------------------------------------
    # Recommendation
    # ---------------------------------------------------------------------------
    bad_retrieval_count = bucket_counts["bad-retrieval"]
    good_retrieval_count = bucket_counts["good-retrieval-weak-draft"]
    no_precedent_count = bucket_counts["no-precedent"]

    print(f"\n{'='*65}")
    print("DECISION")
    print(f"{'='*65}")
    if bad_retrieval_count > no_precedent_count and bad_retrieval_count >= (len(results) / 3):
        print(f"RETRIEVAL IS THE BOTTLENECK ({bad_retrieval_count}/{len(results)} bad-retrieval).")
        print("Recommended lever: Re-embed corpus on richer text (seller reply + full thread)")
        print("or try text-embedding-3-large / higher dimensions.")
    elif no_precedent_count >= (len(results) / 2):
        print(f"NO-PRECEDENT DOMINATES ({no_precedent_count}/{len(results)}).")
        print("Retrieval is near its ceiling. Lever: grow corpus/guides + confidence flagging.")
        print("Do NOT churn retrieval — it cannot retrieve what isn't there.")
    else:
        print(f"MIXED: bad-retrieval={bad_retrieval_count}, good-retrieval-weak-draft={good_retrieval_count}, no-precedent={no_precedent_count}")
        print("Focus on whichever bucket is largest above.")

    # ---------------------------------------------------------------------------
    # Write results
    # ---------------------------------------------------------------------------
    out_dir = Path(__file__).resolve().parents[1] / "tmp"
    out_dir.mkdir(exist_ok=True)
    json_path = out_dir / "eval-rag-audit.json"
    with open(json_path, "w", encoding="utf-8") as f:
        json.dump({
            "run_date": time.strftime('%Y-%m-%d %H:%M UTC', time.gmtime()),
            "schema": schema,
            "n_tickets": len(results),
            "bucket_counts": bucket_counts,
            "total_cost_eur": round(total_cost, 6),
            "corpus_size": len(corpus_thread_ids),
            "retrieval_threshold": RETRIEVAL_THRESHOLD,
            "per_ticket": results,
        }, f, indent=2, ensure_ascii=False)

    print(f"\nResults written to: {json_path}")
    return bucket_counts, total_cost, results


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="RAG Quality Audit for ModuleDesk")
    parser.add_argument("--n", type=int, default=8, help="Number of eval tickets")
    parser.add_argument("--schema", default="tenant_internal", help="Tenant schema")
    args = parser.parse_args()

    run_rag_audit(n=args.n, schema=args.schema)