import os
import re
import time
import json
import hashlib
import requests
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import gradio as gr
from collections import defaultdict
from bs4 import BeautifulSoup
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from rank_bm25 import BM25Okapi
from sentence_transformers import CrossEncoder


# ═══════════════════════════════════════════════════════
#  CONFIG
# ═══════════════════════════════════════════════════════

PERSIST_DIR = "k8s_chroma_db"

URLS = {
    "pods":               "https://kubernetes.io/docs/concepts/workloads/pods/",
    "deployments":        "https://kubernetes.io/docs/concepts/workloads/controllers/deployment/",
    "services":           "https://kubernetes.io/docs/concepts/services-networking/service/",
    "namespaces":         "https://kubernetes.io/docs/concepts/overview/working-with-objects/namespaces/",
    "nodes":              "https://kubernetes.io/docs/concepts/architecture/nodes/",
    "statefulsets":       "https://kubernetes.io/docs/concepts/workloads/controllers/statefulset/",
    "rbac":               "https://kubernetes.io/docs/reference/access-authn-authz/rbac/",
    "persistent-volumes": "https://kubernetes.io/docs/concepts/storage/persistent-volumes/",
    "ingress":            "https://kubernetes.io/docs/concepts/services-networking/ingress/",
    "autoscaling":        "https://kubernetes.io/docs/tasks/run-application/horizontal-pod-autoscale/",
    "docker-overview":    "https://docs.docker.com/get-started/overview/",
    "docker-images":      "https://docs.docker.com/get-started/docker-concepts/the-basics/what-is-an-image/",
    "docker-containers":  "https://docs.docker.com/get-started/docker-concepts/the-basics/what-is-a-container/",
    "docker-volumes":     "https://docs.docker.com/storage/volumes/",
    "docker-networking":  "https://docs.docker.com/network/",
    "docker-compose":     "https://docs.docker.com/compose/",
}

# Cost estimates (llama-3.1-8b-instruct on OpenRouter, per 1K tokens)
COST_PER_1K_INPUT_TOKENS  = 0.00006
COST_PER_1K_OUTPUT_TOKENS = 0.00006
AVG_CHARS_PER_TOKEN       = 4

MAX_RETRIES        = 3
RETRY_BACKOFF_BASE = 1.5
MAX_HISTORY_TURNS  = 4

CONFIDENCE_THRESHOLD   = 0.35
GROUNDEDNESS_THRESHOLD = 40


# ═══════════════════════════════════════════════════════
#  LLM GATEWAY — Caching, Retry, Fallback, Cost Tracking
# ═══════════════════════════════════════════════════════

_response_cache: dict = {}
CACHE_STATS = {"hits": 0, "misses": 0}
COST_LOG: list = []


def _cache_key(prompt: str) -> str:
    return hashlib.sha256(prompt.encode()).hexdigest()

def cache_get(prompt: str):
    return _response_cache.get(_cache_key(prompt))

def cache_set(prompt: str, answer: str) -> None:
    _response_cache[_cache_key(prompt)] = answer

def estimate_tokens(text: str) -> int:
    return max(1, len(text) // AVG_CHARS_PER_TOKEN)

def log_cost(query: str, prompt: str, answer: str, cached: bool = False) -> float:
    input_tok  = estimate_tokens(prompt)
    output_tok = estimate_tokens(answer)
    cost_usd   = (input_tok / 1000 * COST_PER_1K_INPUT_TOKENS +
                  output_tok / 1000 * COST_PER_1K_OUTPUT_TOKENS)
    COST_LOG.append({
        "query": query, "input_tokens": input_tok,
        "output_tokens": output_tok,
        "cost_usd": round(cost_usd, 6), "cached": cached,
    })
    return cost_usd


def _build_fallback(context_chunks) -> str:
    if not context_chunks:
        return "⚠️ The language model is unavailable and no docs were found. Please try again."
    snippets = "\n\n".join(f"• {c[:300]}..." for c in context_chunks[:3])
    return ("⚠️ The language model is temporarily unavailable. "
            "Here are the most relevant documentation excerpts:\n\n"
            + snippets + "\n\n_Please try again for a full AI-generated answer._")


def call_llm(prompt: str, context_chunks=None, max_tokens: int = 400) -> tuple:
    """
    LLM Gateway: Cache → Retry (exponential backoff) → Fallback.
    Returns (answer, was_cached).
    """
    api_key = os.getenv("OPENROUTER_API_KEY")
    if not api_key:
        return _build_fallback(context_chunks), False

    cached = cache_get(prompt)
    if cached:
        CACHE_STATS["hits"] += 1
        log_cost("(cached)", prompt, cached, cached=True)
        return cached, True
    CACHE_STATS["misses"] += 1

    last_error = None
    for attempt in range(1, MAX_RETRIES + 1):
        try:
            res = requests.post(
                "https://openrouter.ai/api/v1/chat/completions",
                headers={
                    "Authorization": f"Bearer {api_key}",
                    "HTTP-Referer": "https://huggingface.co/",
                    "X-Title": "Kubernetes RAG Assistant",
                },
                json={
                    "model": "meta-llama/llama-3.1-8b-instruct",
                    "messages": [{"role": "user", "content": prompt}],
                    "max_tokens": max_tokens,
                    "temperature": 0.0,
                },
                timeout=60,
            )
            res.raise_for_status()
            answer = res.json()["choices"][0]["message"]["content"]
            cache_set(prompt, answer)
            return answer, False
        except Exception as e:
            last_error = e
            if attempt < MAX_RETRIES:
                wait = RETRY_BACKOFF_BASE ** attempt
                print(f"[WARN] LLM attempt {attempt} failed: {e} — retrying in {wait:.1f}s")
                time.sleep(wait)

    print(f"[ERROR] All {MAX_RETRIES} LLM attempts failed: {last_error}. Using fallback.")
    return _build_fallback(context_chunks), False


# ═══════════════════════════════════════════════════════
#  GUARDRAILS
# ═══════════════════════════════════════════════════════

PII_PATTERNS = {
    "email":  r"[a-zA-Z0-9_.+-]+@[a-zA-Z0-9-]+\.[a-zA-Z0-9-.]+",
    "phone":  r"\+?\d[\d\s\-]{8,15}\d",
    "ssn":    r"\b\d{3}-\d{2}-\d{4}\b",
    "credit": r"\b(?:\d[ -]?){13,16}\b",
    "ip":     r"\b(?:\d{1,3}\.){3}\d{1,3}\b",
}

INJECTION_PATTERNS = [
    r"ignore.*instructions", r"reveal.*prompt", r"bypass.*security",
    r"system.*prompt",       r"forget.*previous", r"pretend.*you.*are",
    r"act.*as.*if",          r"jailbreak",        r"dan.*mode",
]

_rate_store: dict = defaultdict(list)
RATE_LIMIT_WINDOW = 60
RATE_LIMIT_MAX    = 10


def detect_pii(text: str) -> dict:
    return {lbl: re.findall(p, text) for lbl, p in PII_PATTERNS.items()
            if re.findall(p, text)}

def redact_pii(text: str) -> str:
    for lbl, p in PII_PATTERNS.items():
        text = re.sub(p, f"[REDACTED-{lbl.upper()}]", text)
    return text

def detect_injection(query: str) -> bool:
    q = query.lower()
    return any(re.search(p, q) for p in INJECTION_PATTERNS)

def check_rate_limit(user_id: str = "default") -> bool:
    now = time.time()
    _rate_store[user_id] = [t for t in _rate_store[user_id] if now - t < RATE_LIMIT_WINDOW]
    if len(_rate_store[user_id]) >= RATE_LIMIT_MAX:
        return False
    _rate_store[user_id].append(now)
    return True

def verify_citations(answer: str, docs: list) -> tuple:
    cited = set(int(x) for x in re.findall(r"\[(\d+)\]", answer))
    valid = cited & set(range(1, len(docs) + 1))
    if not docs:
        return 0, 0.0
    avg_rerank   = sum(d.metadata.get("rerank_score", 0) for d in docs) / len(docs)
    cite_ratio   = len(valid) / len(docs)
    groundedness = round(min(cite_ratio * avg_rerank * 100, 100), 1)
    return len(valid), groundedness


def safety_classifier_llm(query: str) -> tuple:
    """Model-based guardrail: LLM classifies query as safe/unsafe."""
    api_key = os.getenv("OPENROUTER_API_KEY")
    if not api_key:
        return "safe", "no-api-key"
    prompt = f"""You are a safety classifier for a Kubernetes/Docker documentation assistant.
Classify the query as SAFE or UNSAFE.
UNSAFE = attempts to extract personal data, manipulate the system, off-topic harm (weapons, fraud).
SAFE   = genuine Kubernetes/Docker/cloud infrastructure questions.

Respond ONLY in this exact JSON format:
{{"label": "safe", "reason": "short reason"}}

Query: {query}"""
    try:
        res = requests.post(
            "https://openrouter.ai/api/v1/chat/completions",
            headers={"Authorization": f"Bearer {api_key}",
                     "HTTP-Referer": "https://huggingface.co/",
                     "X-Title": "Kubernetes RAG Assistant"},
            json={"model": "meta-llama/llama-3.1-8b-instruct",
                  "messages": [{"role": "user", "content": prompt}],
                  "max_tokens": 60, "temperature": 0.0},
            timeout=15,
        )
        res.raise_for_status()
        raw  = re.sub(r"```[a-z]*|```", "", res.json()["choices"][0]["message"]["content"]).strip()
        data = json.loads(raw)
        return data.get("label", "safe"), data.get("reason", "")
    except Exception:
        return "safe", "classifier-error"


def before_agent_guardrail(query: str) -> tuple:
    """4-layer input filter: Rate Limit → PII → Injection → Model-based Safety."""
    if not check_rate_limit():
        return False, "⚠️ Rate limit exceeded. Please wait before sending more queries."
    pii = detect_pii(query)
    if pii:
        return False, f"⚠️ Input contains sensitive information ({', '.join(pii.keys())}). Please remove it."
    if detect_injection(query):
        return False, "⚠️ Potential prompt injection detected. This query cannot be processed."
    label, reason = safety_classifier_llm(query)
    if label.lower() == "unsafe":
        return False, f"⚠️ Query flagged as unsafe: {reason}"
    return True, ""


def after_agent_guardrail(answer: str, docs: list) -> tuple:
    """4-layer output filter: PII Redact → Citation Check → Groundedness → Human Review flag."""
    answer              = redact_pii(answer)
    cite_count, ground  = verify_citations(answer, docs)
    avg_score           = sum(d.metadata.get("rerank_score", 0) for d in docs) / len(docs) if docs else 0
    needs_review        = (avg_score < CONFIDENCE_THRESHOLD) or (ground < GROUNDEDNESS_THRESHOLD)
    answer              = re.sub(r"Groundedness:\s*\d+%", "", answer).strip()
    answer             += f"\n\n📊 **Groundedness: {ground}%** | Citations verified: {cite_count}/{len(docs)}"
    if needs_review:
        answer += "\n\n🔍 **Human Review Recommended** — confidence below threshold."
    return answer, needs_review, ground


# ═══════════════════════════════════════════════════════
#  KNOWLEDGE BASE
# ═══════════════════════════════════════════════════════

def scrape_page(name, url):
    try:
        r    = requests.get(url, timeout=20)
        r.raise_for_status()
        soup = BeautifulSoup(r.text, "html.parser")
        content = (soup.find("div", class_="td-content") or
                   soup.find("div", class_="docs-content"))
        if not content:
            return None
        return Document(page_content=content.get_text(separator="\n").strip(),
                        metadata={"doc_id": name, "url": url})
    except Exception as e:
        print(f"[ERROR] scraping {url}: {e}")
        return None


def build_or_load_kb():
    emb = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    if os.path.isdir(PERSIST_DIR):
        print("[INFO] Loading existing DB...")
        vectordb = Chroma(embedding_function=emb, persist_directory=PERSIST_DIR)
        raw      = vectordb._collection.get(include=["documents", "metadatas"])
        chunks   = [Document(page_content=d, metadata=m)
                    for d, m in zip(raw["documents"], raw["metadatas"])]
        return vectordb, chunks
    print("[INFO] No DB found — scraping docs...")
    docs     = [d for name, url in URLS.items() if (d := scrape_page(name, url))]
    splitter = RecursiveCharacterTextSplitter(chunk_size=900, chunk_overlap=200)
    chunks   = splitter.split_documents(docs)
    vectordb = Chroma.from_documents(chunks, emb, persist_directory=PERSIST_DIR)
    vectordb.persist()
    print(f"[INFO] DB created with {len(chunks)} chunks.")
    return vectordb, chunks


vectordb, chunks = build_or_load_kb()
bm25      = BM25Okapi([c.page_content.split() for c in chunks])
reranker  = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-12-v2")
retriever = vectordb.as_retriever(
    search_type="similarity_score_threshold",
    search_kwargs={"k": 8, "score_threshold": 0.35},
)


# ═══════════════════════════════════════════════════════
#  RETRIEVAL — Hybrid BM25 + Vector + CrossEncoder Rerank
# ═══════════════════════════════════════════════════════

def hybrid_search(query: str, top_k: int = 5) -> list:
    vector_results = retriever.invoke(query)
    bm_scores      = bm25.get_scores(query.lower().split())
    bm_results     = [d for _, d in sorted(zip(bm_scores, chunks), key=lambda x: x[0], reverse=True)[:top_k]]
    combined       = vector_results + bm_results
    seen, unique   = set(), []
    for d in combined:
        key = (d.metadata.get("doc_id"), d.page_content[:80])
        if key not in seen:
            seen.add(key); unique.append(d)
    if not unique:
        return []
    scores = reranker.predict([(query, d.page_content) for d in unique])
    ranked = sorted(zip(scores, unique), key=lambda x: x[0], reverse=True)[:top_k]
    for s, doc in ranked:
        doc.metadata["rerank_score"] = float(s)
    return [doc for _, doc in ranked]


# ═══════════════════════════════════════════════════════
#  METRICS
# ═══════════════════════════════════════════════════════

METRICS = {
    "q": [], "lat": [], "tok": [], "g": [], "r": [],
    "c": [], "t":   [], "review": [], "guardrail": [],
    "cached": [], "cost": [],
}

def classify_query(q: str) -> str:
    q = q.lower()
    if "how" in q:                    return "how-to"
    if "error" in q or "fail" in q:   return "debug"
    return "general"


# ═══════════════════════════════════════════════════════
#  CONVERSATION MEMORY
# ═══════════════════════════════════════════════════════

def build_prompt_with_history(query: str, history: list, context: str) -> str:
    recent = history[-(MAX_HISTORY_TURNS * 2):]
    history_text = ""
    for msg in recent:
        role, content = msg.get("role", ""), msg.get("content", "")
        if role == "user":
            history_text += f"User: {content}\n"
        elif role == "assistant":
            history_text += f"Assistant: {content[:400]}...\n"

    history_block = (f"Previous conversation (for context only):\n{history_text}\n---\n"
                     if history_text else "")

    return f"""{history_block}Answer using ONLY the context below.
Every sentence MUST include at least one citation like [1] or [2].
Do NOT add any groundedness score — that is computed separately.

Current question: {query}

Context:
{context}
"""


# ═══════════════════════════════════════════════════════
#  FEATURE 1 — SUGGESTED FOLLOW-UP QUESTIONS
# ═══════════════════════════════════════════════════════

def generate_followup_questions(query: str, answer: str) -> list:
    """
    Asks the LLM to suggest 3 natural follow-up questions based on
    what was just asked and answered. Returns list of 3 strings.
    Falls back to empty list on any error so the main flow is never broken.
    """
    api_key = os.getenv("OPENROUTER_API_KEY")
    if not api_key:
        return []

    prompt = f"""A user asked a Kubernetes/Docker assistant:
"{query}"

The assistant answered it. Suggest exactly 3 short follow-up questions the user might ask next.
Rules:
- Must be about Kubernetes or Docker only
- Keep each question under 12 words
- Return ONLY a valid JSON array of 3 strings, no other text
- Example: ["Question one?", "Question two?", "Question three?"]

JSON array:"""

    try:
        res = requests.post(
            "https://openrouter.ai/api/v1/chat/completions",
            headers={"Authorization": f"Bearer {api_key}",
                     "HTTP-Referer": "https://huggingface.co/",
                     "X-Title": "Kubernetes RAG Assistant"},
            json={"model": "meta-llama/llama-3.1-8b-instruct",
                  "messages": [{"role": "user", "content": prompt}],
                  "max_tokens": 120, "temperature": 0.7},
            timeout=15,
        )
        res.raise_for_status()
        raw       = re.sub(r"```[a-z]*|```", "",
                           res.json()["choices"][0]["message"]["content"]).strip()
        questions = json.loads(raw)
        if isinstance(questions, list):
            return [str(q) for q in questions[:3]]
    except Exception as e:
        print(f"[WARN] follow-up generation failed: {e}")
    return []


# ═══════════════════════════════════════════════════════
#  FEATURE 2 — RETRIEVAL DETAILS PANEL
#  Shows exactly which chunks were retrieved, their reranker
#  scores, and whether each was cited. Makes the RAG pipeline
#  visible — this is what separates the app from ChatGPT.
# ═══════════════════════════════════════════════════════

def build_retrieval_details(docs: list, answer: str) -> str:
    cited_indices = set(int(x) for x in re.findall(r"\[(\d+)\]", answer))
    lines = ["**🔍 Retrieval Details** — how this answer was found\n",
             "| # | Source | Rerank Score | Cited? | Snippet |",
             "|---|--------|-------------|--------|---------|"]
    for i, doc in enumerate(docs, 1):
        score     = doc.metadata.get("rerank_score", 0)
        url       = doc.metadata.get("url", "unknown")
        cited     = "✅ Yes" if i in cited_indices else "❌ No"
        snippet   = doc.page_content[:120].replace("\n", " ").strip() + "..."
        short_url = (url.replace("https://kubernetes.io/docs/", "k8s/")
                        .replace("https://docs.docker.com/", "docker/"))
        lines.append(f"| [{i}] | `{short_url}` | `{score:.3f}` | {cited} | {snippet} |")
    lines.append(f"\n**Chunks retrieved:** {len(docs)} &nbsp;|&nbsp; "
                 f"**Chunks cited:** {len(cited_indices)}")
    return "\n".join(lines)


# ═══════════════════════════════════════════════════════
#  MAIN CHAT HANDLER
# ═══════════════════════════════════════════════════════

def answer_question(query: str, history: list) -> tuple:
    """
    Returns:
        history        — updated chat history
        clear_input    — empty string to clear the textbox
        fq1, fq2, fq3 — three follow-up question strings (may be empty)
        retrieval_md   — markdown string for the retrieval details panel
    """
    start       = time.time()
    NO_EXTRAS   = ("", "", "", "")   # fq1, fq2, fq3, retrieval_md

    # ── Before-Agent Guardrail ────────────────────────
    is_safe, block_reason = before_agent_guardrail(query)
    if not is_safe:
        _log(query, time.time()-start, 0, 0, 0, 0, blocked=True)
        history.append({"role": "user",      "content": query})
        history.append({"role": "assistant", "content": block_reason})
        return history, "", *NO_EXTRAS

    # ── Retrieval ─────────────────────────────────────
    docs = hybrid_search(query)
    if not docs:
        history.append({"role": "user",      "content": query})
        history.append({"role": "assistant",
                        "content": "❌ Not found in the knowledge base. "
                                   "Please ask about Kubernetes or Docker."})
        return history, "", *NO_EXTRAS

    # ── Build context & prompt (with conversation memory) ─
    ctx, sources = "", []
    for i, d in enumerate(docs, 1):
        ctx += f"[{i}] {d.page_content[:900]}\nSource: {d.metadata['url']}\n\n"
        sources.append(f"[{i}] → {d.metadata['url']}")

    prompt = build_prompt_with_history(query, history, ctx)

    # ── LLM call (cache → retry → fallback) ──────────
    raw_answer, was_cached = call_llm(prompt, context_chunks=[d.page_content for d in docs])
    cost    = log_cost(query, prompt, raw_answer, cached=was_cached)
    latency = time.time() - start

    # ── After-Agent Guardrail ─────────────────────────
    final_answer, needs_review, groundedness = after_agent_guardrail(raw_answer, docs)

    cache_tag     = " ⚡ _Served from cache_" if was_cached else ""
    final_answer += "\n\n---\n**Sources:**\n" + "\n".join(sources) + cache_tag
    if needs_review:
        final_answer += "\n\n> ℹ️ A human expert should verify this answer."

    # ── Feature 1: Follow-up questions ───────────────
    followups = generate_followup_questions(query, raw_answer)
    while len(followups) < 3:
        followups.append("")
    fq1, fq2, fq3 = followups[0], followups[1], followups[2]

    # ── Feature 2: Retrieval details panel ───────────
    retrieval_md = build_retrieval_details(docs, raw_answer)

    # ── Log metrics ───────────────────────────────────
    avg_score  = sum(d.metadata.get("rerank_score", 0) for d in docs) / len(docs)
    cite_count = len(set(re.findall(r"\[(\d+)\]", raw_answer)))
    _log(query, latency, groundedness, avg_score, cite_count,
         cost, tokens=len(raw_answer.split()),
         review=needs_review, cached=was_cached)

    history.append({"role": "user",      "content": query})
    history.append({"role": "assistant", "content": final_answer})
    return history, "", fq1, fq2, fq3, retrieval_md


def _log(query, latency, ground, avg_score, cite_count, cost,
         tokens=0, review=False, cached=False, blocked=False):
    METRICS["q"].append(query)
    METRICS["lat"].append(round(latency, 2))
    METRICS["tok"].append(tokens)
    METRICS["g"].append(ground)
    METRICS["r"].append(round(avg_score, 3))
    METRICS["c"].append(cite_count)
    METRICS["t"].append(classify_query(query))
    METRICS["review"].append(review)
    METRICS["guardrail"].append(blocked)
    METRICS["cached"].append(cached)
    METRICS["cost"].append(round(cost, 6))


def use_followup(question: str, history: list) -> tuple:
    """Called when user clicks a follow-up suggestion button."""
    if not question.strip():
        return history, "", "", "", "", ""
    return answer_question(question, history)


# ═══════════════════════════════════════════════════════
#  ANALYTICS CHARTS
# ═══════════════════════════════════════════════════════

def make_groundedness_chart():
    fig, ax = plt.subplots(figsize=(7, 3))
    if not METRICS["g"]:
        ax.text(0.5, 0.5, "No data yet", ha="center", va="center", transform=ax.transAxes)
        return fig
    idx    = list(range(1, len(METRICS["g"]) + 1))
    colors = ["#2ecc71" if g >= 60 else "#f39c12" if g >= 40 else "#e74c3c"
              for g in METRICS["g"]]
    ax.bar(idx, METRICS["g"], color=colors, width=0.6, edgecolor="none")
    ax.axhline(40, color="#e74c3c", linestyle="--", lw=1, label="Review threshold (40%)")
    ax.axhline(60, color="#2ecc71", linestyle="--", lw=1, label="Good (60%)")
    ax.set_xlabel("Query #"); ax.set_ylabel("Groundedness (%)")
    ax.set_title("Groundedness per Query", fontweight="bold")
    ax.set_ylim(0, 105); ax.set_xticks(idx); ax.legend(fontsize=8)
    ax.spines["top"].set_visible(False); ax.spines["right"].set_visible(False)
    fig.tight_layout(); return fig


def make_query_type_chart():
    fig, ax = plt.subplots(figsize=(4, 3))
    if not METRICS["t"]:
        ax.text(0.5, 0.5, "No data yet", ha="center", va="center", transform=ax.transAxes)
        return fig
    counts  = {}
    for t in METRICS["t"]:
        counts[t] = counts.get(t, 0) + 1
    palette = ["#3498db", "#9b59b6", "#1abc9c", "#e67e22"][:len(counts)]
    wedges, texts, autotexts = ax.pie(
        counts.values(), labels=counts.keys(), autopct="%1.0f%%",
        colors=palette, startangle=90,
        wedgeprops={"edgecolor": "white", "linewidth": 1.5})
    for at in autotexts: at.set_fontsize(9)
    ax.set_title("Query Types", fontweight="bold")
    fig.tight_layout(); return fig


def make_latency_chart():
    fig, ax = plt.subplots(figsize=(7, 3))
    if not METRICS["lat"]:
        ax.text(0.5, 0.5, "No data yet", ha="center", va="center", transform=ax.transAxes)
        return fig
    idx    = list(range(1, len(METRICS["lat"]) + 1))
    colors = ["#f39c12" if c else "#3498db" for c in METRICS["cached"]]
    ax.bar(idx, METRICS["lat"], color=colors, width=0.6, edgecolor="none")
    ax.set_xlabel("Query #"); ax.set_ylabel("Latency (s)")
    ax.set_title("Latency per Query  (orange = cache hit)", fontweight="bold")
    ax.set_xticks(idx)
    ax.spines["top"].set_visible(False); ax.spines["right"].set_visible(False)
    fig.tight_layout(); return fig


def make_cost_chart():
    fig, ax = plt.subplots(figsize=(7, 3))
    if not METRICS["cost"]:
        ax.text(0.5, 0.5, "No data yet", ha="center", va="center", transform=ax.transAxes)
        return fig
    running, total = [], 0
    for c in METRICS["cost"]:
        total += c; running.append(round(total, 6))
    idx = list(range(1, len(running) + 1))
    ax.plot(idx, running, marker="o", color="#9b59b6", lw=2, markersize=5)
    ax.fill_between(idx, running, alpha=0.15, color="#9b59b6")
    ax.set_xlabel("Query #"); ax.set_ylabel("Cumulative cost (USD)")
    ax.set_title("Running Cost Tracker", fontweight="bold")
    ax.set_xticks(idx)
    ax.spines["top"].set_visible(False); ax.spines["right"].set_visible(False)
    fig.tight_layout(); return fig


def update_dashboard():
    if not METRICS["q"]:
        empty = plt.figure()
        return [], 0, 0, 0, 0, 0, 0, empty, empty, empty, empty
    rows = list(zip(
        range(1, len(METRICS["q"]) + 1), METRICS["q"],
        METRICS["lat"], METRICS["tok"], METRICS["g"], METRICS["r"],
        METRICS["c"],   METRICS["t"],
        ["✅" if v else "—" for v in METRICS["review"]],
        ["🛡️" if v else "—" for v in METRICS["guardrail"]],
        ["⚡" if v else "—" for v in METRICS["cached"]],
        METRICS["cost"],
    ))
    n = len(METRICS["g"])
    return (
        rows,
        round(sum(METRICS["g"])            / n, 2),
        round(sum(METRICS["lat"])           / n, 2),
        round(sum(METRICS["tok"])           / n, 2),
        round(sum(METRICS["review"])        / n * 100, 1),
        round(sum(METRICS["guardrail"])     / n * 100, 1),
        round(sum(METRICS["cost"]), 6),
        make_groundedness_chart(), make_query_type_chart(),
        make_latency_chart(),      make_cost_chart(),
    )


# ═══════════════════════════════════════════════════════
#  UI
# ═══════════════════════════════════════════════════════

with gr.Blocks(title="Kubernetes & Docker RAG Assistant") as app:
    gr.Markdown(
        "# ☸️ Kubernetes & Docker RAG Assistant\n"
        "> Hybrid RAG · Layered Guardrails · Conversation Memory · "
        "Follow-up Suggestions · Retrieval Transparency"
    )

    # ── Tab 1: Chat ───────────────────────────────────
    with gr.Tab("💬 Chat"):
        chat    = gr.Chatbot(height=460)
        user_in = gr.Textbox(
            label="Ask about Kubernetes or Docker",
            placeholder="e.g. How do I set resource limits for a Pod?",
        )
        with gr.Row():
            submit_btn = gr.Button("Send", variant="primary")
            clear_btn  = gr.Button("Clear conversation")

        # ── Feature 1: Follow-up question buttons ─────
        gr.Markdown("**💡 Suggested follow-ups** — click to ask:")
        with gr.Row():
            fq1_btn = gr.Button("", visible=True, size="sm")
            fq2_btn = gr.Button("", visible=True, size="sm")
            fq3_btn = gr.Button("", visible=True, size="sm")

        # ── Feature 2: Retrieval details accordion ────
        with gr.Accordion("🔍 Why this answer? (Retrieval Details)", open=False):
            retrieval_panel = gr.Markdown(
                "_Ask a question to see which documentation chunks were retrieved, "
                "their reranker scores, and which ones were cited._"
            )

        # Wire up main send
        all_outputs = [chat, user_in, fq1_btn, fq2_btn, fq3_btn, retrieval_panel]
        user_in.submit(answer_question, [user_in, chat], all_outputs)
        submit_btn.click(answer_question, [user_in, chat], all_outputs)
        clear_btn.click(
            lambda: ([], "", "", "", "", "_Ask a question to see retrieval details._"),
            None, all_outputs
        )

        # Wire up follow-up buttons — clicking one runs it as a new query
        fq1_btn.click(use_followup, [fq1_btn, chat], all_outputs)
        fq2_btn.click(use_followup, [fq2_btn, chat], all_outputs)
        fq3_btn.click(use_followup, [fq3_btn, chat], all_outputs)

    # ── Tab 2: Guardrails ─────────────────────────────
    with gr.Tab("🛡️ Guardrails"):
        gr.Markdown("""
### Guardrails Architecture

**Before-Agent (input filtering)**

| Layer | Type | What it checks |
|---|---|---|
| Rate Limiting | Deterministic | Max 10 req / min per user |
| PII Detection | Deterministic | Email, phone, SSN, credit card, IP |
| Prompt Injection | Deterministic | Regex patterns on known attack phrases |
| Safety Classifier | **Model-based** | LLM classifies Safe / Unsafe with reason |

**After-Agent (output safety)**

| Layer | Type | What it checks |
|---|---|---|
| PII Redaction | Deterministic | Scrubs leaked PII from LLM output |
| Citation Verification | Deterministic | Checks [N] refs exist and are valid |
| Groundedness Score | Deterministic | Citation ratio × reranker score — not LLM self-report |
| Confidence Threshold | Deterministic | Flags for human review if score < 0.35 or ground < 40% |

**Human-in-the-Loop**

When confidence is low the assistant shows a _Human Review Recommended_ banner.
The reviewer sees groundedness %, reranker scores, and source URLs in the Retrieval Details panel.
        """)

    # ── Tab 3: Analytics ──────────────────────────────
    with gr.Tab("📊 Analytics"):
        gr.Markdown("### Query Analytics Dashboard")
        refresh_btn = gr.Button("🔄 Refresh Dashboard")
        with gr.Row():
            avgG_box         = gr.Number(label="Avg Groundedness (%)")
            avgL_box         = gr.Number(label="Avg Latency (s)")
            avgT_box         = gr.Number(label="Avg Tokens")
            pctReview_box    = gr.Number(label="% Human Review")
            pctGuardrail_box = gr.Number(label="% Blocked")
            totalCost_box    = gr.Number(label="Total Cost (USD)")
        with gr.Row():
            ground_chart  = gr.Plot(label="Groundedness over time")
            latency_chart = gr.Plot(label="Latency (orange = cache hit)")
        with gr.Row():
            type_chart = gr.Plot(label="Query types")
            cost_chart = gr.Plot(label="Running cost tracker")
        gr.Markdown("#### Raw query log")
        table = gr.DataFrame(
            headers=["ID", "Query", "Latency (s)", "Tokens", "Groundedness (%)",
                     "Rerank Score", "Citations", "Type",
                     "Human Review", "Guardrail", "Cached", "Cost (USD)"],
            interactive=False,
        )
        refresh_btn.click(
            update_dashboard, None,
            [table, avgG_box, avgL_box, avgT_box,
             pctReview_box, pctGuardrail_box, totalCost_box,
             ground_chart, type_chart, latency_chart, cost_chart],
        )

app.launch()