from __future__ import annotations

try:
    import spaces
except ImportError:
    spaces = None

import html
import os
import threading
import uuid
from pathlib import Path
from typing import Any
import gradio as gr

from config import (
    APP_TITLE,
    DEFAULT_TOP_K,
    EMBEDDING_MODEL_ID,
    FAR_DATASET_REPO,
    HF_INFERENCE_MODEL,
    HF_TOKEN,
    LOCAL_MODEL_ID,
    MAX_UPLOAD_MB,
)

from document_loader import load_file

from rag_engine import (
    KnowledgeBase,
    answer_question,
    load_far_dataset,
    render_sources,
)


SESSIONS: dict[str, KnowledgeBase] = {}
SESSIONS_LOCK = threading.Lock()


def _session(sid: str | None) -> tuple[str, KnowledgeBase]:
    sid = sid or uuid.uuid4().hex
    with SESSIONS_LOCK:
        if sid not in SESSIONS:
            SESSIONS[sid] = KnowledgeBase()
        return sid, SESSIONS[sid]


def _status(kb: KnowledgeBase, message: str = "Ready") -> str:
    return f"""
    <div class="status-card">
      <strong>{html.escape(message)}</strong><br>
      {kb.document_count} source document(s) · {len(kb.chunks):,} indexed chunks
    </div>
    """


def index_uploads(files: list[Any] | None, sid: str | None):
    sid, kb = _session(sid)
    if not files:
        return sid, _status(kb, "No files selected"), {"documents": kb.document_count, "chunks": len(kb.chunks)}

    added = 0
    names = []
    for uploaded in files:
        path = Path(uploaded.name if hasattr(uploaded, "name") else str(uploaded))
        size_mb = path.stat().st_size / (1024 * 1024)
        if size_mb > MAX_UPLOAD_MB:
            raise gr.Error(f"{path.name} is {size_mb:.1f} MB. Limit is {MAX_UPLOAD_MB} MB per file.")
        try:
            chunks = load_file(path)
            added += kb.add(chunks)
            names.append(path.name)
        except Exception as exc:
            raise gr.Error(f"Could not index {path.name}: {exc}") from exc

    message = f"Indexed {added:,} chunks from {len(names)} file(s)"
    diagnostics = {
        "documents": kb.document_count,
        "chunks": len(kb.chunks),
        "added_chunks": added,
        "files": names,
        "embedding_model": EMBEDDING_MODEL_ID,
    }
    return sid, _status(kb, message), diagnostics


def _build_far_suggestions(chunks, limit: int = 8) -> list[str]:
    """Build question suggestions from topics actually present in loaded FAR chunks."""
    corpus = "\n".join(chunk.text.lower() for chunk in chunks)

    topic_questions = [
        (("market research",), "What does the FAR require for market research?"),
        (("acquisition planning",), "What does the FAR require for acquisition planning?"),
        (("full and open competition", "competition requirements"), "What does the FAR require for full and open competition?"),
        (("commercial products", "commercial services"), "What FAR requirements apply to commercial products and commercial services?"),
        (("contracting officer",), "What responsibilities and authorities does the FAR assign to contracting officers?"),
        (("small business",), "What does the FAR require agencies to consider regarding small business participation?"),
        (("past performance",), "How does the FAR address contractor past performance?"),
        (("sole source", "other than full and open competition"), "When does the FAR permit an acquisition without full and open competition?"),
        (("contract modification", "modifications"), "What does the FAR say about contract modifications?"),
        (("termination for convenience", "termination for default"), "What does the FAR say about contract termination?"),
    ]

    suggestions: list[str] = []
    for keywords, question in topic_questions:
        if any(keyword in corpus for keyword in keywords):
            suggestions.append(question)
        if len(suggestions) >= limit:
            break

    # Fallback: derive section-specific questions from section metadata if needed
    if len(suggestions) < limit:
        seen_sections: set[str] = set()
        for chunk in chunks:
            section = (chunk.section or "").strip()
            if not section or section in seen_sections:
                continue
            seen_sections.add(section)
            suggestions.append(f"What requirements are described in {section}?")
            if len(suggestions) >= limit:
                break

    return suggestions[:limit]


def load_far(sid: str | None, row_limit: int):
    sid, kb = _session(sid)
    try:
        chunks = load_far_dataset(max_rows=int(row_limit))
        added = kb.add(chunks)
        suggestions = _build_far_suggestions(chunks)
    except Exception as exc:
        raise gr.Error(f"FAR starter load failed: {exc}") from exc

    diagnostics = {
        "documents": kb.document_count,
        "chunks": len(kb.chunks),
        "added_far_chunks": added,
        "dataset": FAR_DATASET_REPO,
        "embedding_model": EMBEDDING_MODEL_ID,
        "suggested_questions": len(suggestions),
    }

    suggestion_component = gr.Radio(
        choices=suggestions,
        value=None,
        label="Suggested questions from the loaded FAR corpus",
        info="These appear only after FAR chunks are loaded and are selected from topics found in the indexed corpus.",
        interactive=True,
        visible=True,
        elem_id="far-suggestions",
    )

    return (
        sid,
        _status(kb, f"Loaded FAR starter corpus: {added:,} chunks"),
        diagnostics,
        suggestion_component,
    )


def clear_kb(sid: str | None):
    sid, kb = _session(sid)
    kb.clear()
    hidden_suggestions = gr.Radio(
        choices=[],
        value=None,
        label="Suggested questions from the loaded FAR corpus",
        interactive=True,
        visible=False,
        elem_id="far-suggestions",
    )
    return (
        sid,
        _status(kb, "Knowledge base cleared"),
        {},
        "",
        "<div class='status-card'>No sources retrieved.</div>",
        hidden_suggestions,
    )


def ask(question: str, sid: str | None, top_k: int, temperature: float, max_new_tokens: int):
    clean = (question or "").strip()
    if not clean:
        return "Enter a policy question.", "<div class='status-card'>No sources retrieved.</div>", {}

    sid, kb = _session(sid)
    if not kb.chunks:
        return (
            "Index uploaded documents or load the FAR starter corpus first.",
            "<div class='status-card'>No sources retrieved.</div>",
            {"documents": 0, "chunks": 0},
        )

    hits = kb.search(clean, top_k=int(top_k))
    answer, diagnostics = answer_question(
        clean,
        hits,
        temperature=float(temperature),
        max_new_tokens=int(max_new_tokens),
    )
    diagnostics.update({"documents": kb.document_count, "indexed_chunks": len(kb.chunks)})
    return answer, render_sources(hits), diagnostics


CSS = r"""
:root {
  --bg: #05070a;
  --surface: rgba(10, 14, 20, .94);
  --surface-2: rgba(16, 23, 34, .92);
  --surface-3: rgba(255,255,255,.035);
  --line: rgba(255,255,255,.12);
  --line-strong: rgba(255,255,255,.22);
  --muted: #aab3c2;
  --text: #f7f9fc;
  --navy: #12345b;
  --navy-2: #0a2747;
  --red: #b3263f;
  --red-2: #8f1830;
  --white: #ffffff;
  --success: #8fd6ac;
  --danger: #ff9eaa;
}

html, body {
  background: var(--bg) !important;
}

.gradio-container {
  max-width: 1500px !important;
  margin: 0 auto !important;
  color: var(--text) !important;
  background:
    radial-gradient(circle at 12% 4%, rgba(18,52,91,.34), transparent 31%),
    radial-gradient(circle at 88% 10%, rgba(179,38,63,.13), transparent 26%),
    linear-gradient(145deg,#040609 0%,#090d13 46%,#0b111a 100%) !important;
  min-height: 100vh;
}

.main-shell {
  padding: 28px 24px 44px;
}

.hero {
  position: relative;
  overflow: hidden;
  border: 1px solid var(--line);
  background:
    linear-gradient(135deg,rgba(11,20,32,.98),rgba(5,8,13,.96));
  border-radius: 22px;
  padding: 32px 34px 34px;
  box-shadow: 0 28px 80px rgba(0,0,0,.42);
  margin-bottom: 20px;
}

.hero::before {
  content: "";
  position: absolute;
  inset: 0 auto auto 0;
  width: 100%;
  height: 4px;
  background: linear-gradient(90deg,#ffffff 0 34%,#12345b 34% 67%,#b3263f 67% 100%);
  opacity: .95;
}

.hero::after {
  content: "";
  position: absolute;
  width: 430px;
  height: 430px;
  right: -210px;
  top: -210px;
  border: 1px solid rgba(255,255,255,.07);
  border-radius: 50%;
  box-shadow:
    0 0 0 54px rgba(18,52,91,.08),
    0 0 0 108px rgba(179,38,63,.035);
}

.eyebrow {
  position: relative;
  z-index: 1;
  color: #c5d6ea;
  font-size: 12px;
  font-weight: 900;
  letter-spacing: .17em;
  text-transform: uppercase;
}

.hero h1 {
  position: relative;
  z-index: 1;
  margin: 9px 0 8px;
  font-size: clamp(34px,5vw,58px);
  line-height: 1.02;
  letter-spacing: -.045em;
  color: #ffffff;
}

.gradient-word {
  background: linear-gradient(105deg,#ffffff 0%,#d7e4f2 45%,#8aa8c8 100%);
  -webkit-background-clip: text;
  background-clip: text;
  color: transparent;
}

.hero p {
  position: relative;
  z-index: 1;
  max-width: 940px;
  color: #b6c0cf;
  font-size: 16px;
  line-height: 1.65;
  margin: 0;
}

.badges {
  position: relative;
  z-index: 1;
  display: flex;
  flex-wrap: wrap;
  gap: 9px;
  margin-top: 19px;
}

.badge {
  border: 1px solid rgba(255,255,255,.13);
  background: rgba(255,255,255,.04);
  padding: 7px 11px;
  border-radius: 999px;
  color: #bcc7d6;
  font-size: 12px;
  backdrop-filter: blur(8px);
}

.badge strong {
  color: #ffffff;
  margin-right: 4px;
}

.app-panel {
  background: var(--surface) !important;
  border: 1px solid var(--line) !important;
  border-radius: 20px !important;
  box-shadow: 0 20px 55px rgba(0,0,0,.30);
  overflow: hidden;
  padding-top: 14px !important;
}

.sidebar-card {
  background: var(--surface-2);
  border: 1px solid var(--line);
  border-radius: 17px;
  padding: 18px;
  margin-bottom: 14px;
  box-shadow: inset 0 1px 0 rgba(255,255,255,.025);
}

.sidebar-card h3 {
  margin: 0 0 8px;
  color: #ffffff;
  font-size: 14px;
}

.sidebar-card p,
.sidebar-card li {
  color: var(--muted);
  font-size: 13px;
  line-height: 1.58;
}

.sidebar-card ol {
  margin: 9px 0 0;
  padding-left: 20px;
}

.status-card {
  border: 1px solid var(--line);
  border-left: 3px solid #587da5;
  background: rgba(255,255,255,.025);
  border-radius: 13px;
  padding: 12px 14px;
  color: #b5c0cf;
  font-size: 12px;
  line-height: 1.55;
}

.status-card strong {
  color: #ffffff;
}

#ask-button,
#index-button,
#far-button {
  font-weight: 900;
}

button.primary,
#ask-button,
#index-button,
#far-button {
  background: linear-gradient(135deg,#173d68,#0b2949) !important;
  color: #ffffff !important;
  border: 1px solid rgba(157,190,225,.30) !important;
  box-shadow: 0 8px 24px rgba(5,18,34,.30) !important;
}

button.primary:hover,
#ask-button:hover,
#index-button:hover,
#far-button:hover {
  filter: brightness(1.12);
  border-color: rgba(255,255,255,.30) !important;
}

/* Federal-style tab navigation. Build KB is the left/default tab. */
#main-tabs [role="tablist"] {
  padding: 14px 22px 16px !important;
  gap: 12px !important;
  border-bottom: 1px solid var(--line) !important;
  background: rgba(255,255,255,.012) !important;
}

#main-tabs button[role="tab"] {
  margin: 0 !important;
  padding: 11px 18px !important;
  min-height: 44px !important;
  border: 1px solid rgba(255,255,255,.09) !important;
  border-radius: 11px !important;
  color: #c0cad8 !important;
  background: rgba(255,255,255,.025) !important;
  font-weight: 800 !important;
}

#main-tabs button[role="tab"][aria-selected="true"] {
  color: #ffffff !important;
  background: linear-gradient(135deg,#15375f,#0b2949) !important;
  border-color: rgba(158,191,226,.28) !important;
  box-shadow: inset 0 -2px 0 #b3263f, 0 8px 18px rgba(0,0,0,.18) !important;
}

.tab-body {
  padding: 24px 22px 26px !important;
}

/* Keep main Ask action compact and centered. */
.ask-action-row {
  justify-content: center !important;
  align-items: center !important;
  gap: 12px !important;
  margin: 10px 0 5px !important;
}

#ask-button {
  flex: 0 0 auto !important;
  width: 240px !important;
  min-width: 210px !important;
  max-width: 260px !important;
}

.clear-answer-button {
  flex: 0 0 auto !important;
  width: 128px !important;
  min-width: 118px !important;
  max-width: 140px !important;
}

#far-suggestions {
  margin-top: 16px !important;
  padding-top: 5px !important;
  border-top: 1px solid rgba(255,255,255,.08) !important;
}

.answer-panel {
  min-height: 300px;
}

.sources-grid {
  display: grid;
  grid-template-columns: repeat(auto-fit,minmax(260px,1fr));
  gap: 10px;
  margin-top: 4px;
}

.source-card {
  border: 1px solid var(--line);
  background: rgba(255,255,255,.025);
  border-radius: 14px;
  padding: 13px;
}

.source-top {
  display: flex;
  justify-content: space-between;
  align-items: center;
  gap: 8px;
  margin-bottom: 8px;
}

.source-id {
  display: inline-flex;
  min-width: 34px;
  justify-content: center;
  border: 1px solid rgba(255,255,255,.20);
  border-radius: 999px;
  padding: 4px 8px;
  color: #ffffff;
  background: #173d68;
  font-weight: 900;
  font-size: 11px;
}

.score {
  color: #8996a8;
  font-size: 11px;
}

.source-title {
  color: #f4f7fb;
  font-weight: 800;
  font-size: 12px;
  line-height: 1.45;
  margin-bottom: 7px;
}

.source-excerpt {
  color: #aeb9c8;
  font-size: 12px;
  line-height: 1.55;
}

.source-link {
  margin-top: 9px;
  font-size: 11px;
}

.source-link a {
  color: #9fc0e2 !important;
}

.footer-note {
  color: #7f8b9c;
  font-size: 11px;
  text-align: center;
  margin-top: 17px;
}

.warning-note {
  margin-top: 14px;
  padding: 11px 13px;
  border: 1px solid rgba(179,38,63,.30);
  border-left: 3px solid var(--red);
  border-radius: 10px;
  background: rgba(179,38,63,.06);
  color: #c9cfd8;
  font-size: 11px;
  line-height: 1.55;
}

@media(max-width:800px) {
  .main-shell { padding: 14px 10px 28px; }
  .hero { padding: 24px 20px 25px; border-radius: 18px; }
  .hero h1 { font-size: 37px; }
  #main-tabs [role="tablist"] { padding: 12px 14px 14px !important; }
  .tab-body { padding: 20px 16px 22px !important; }
}
"""

HEAD = """
<meta name="theme-color" content="#07101c">
<meta name="description" content="Citation-grounded RAG for FAR, DFARS, agency policy, directives, manuals, and uploaded federal documents.">
"""


def build_app() -> gr.Blocks:
    backend_label = f"HF Inference: {HF_INFERENCE_MODEL}" if HF_TOKEN else f"Local fallback: {LOCAL_MODEL_ID}"
    hero = f"""
    <div class="hero">
      <div class="eyebrow">RAG · Federal Acquisition</div>
      <h1>Federal Policy <span class="gradient-word">RAG Assistant</span></h1>
      <p>Upload agency policies, directives, manuals, FAR/DFARS material, and related guidance. Retrieve the most relevant passages and answer with traceable source markers mapped to exact source cards.</p>
      <div class="badges">
        <span class="badge"><strong>Generator</strong> {html.escape(HF_INFERENCE_MODEL)}</span>
        <span class="badge"><strong>Embeddings</strong> {html.escape(EMBEDDING_MODEL_ID)}</span>
        <span class="badge"><strong>Index</strong> FAISS</span>
        <span class="badge"><strong>Citations</strong> [S1] source mapping</span>
      </div>
    </div>
    """

    with gr.Blocks(title=APP_TITLE) as demo:
        sid = gr.State("")
        with gr.Column(elem_classes=["main-shell"]):
            gr.HTML(hero)
            with gr.Row(equal_height=False):
                with gr.Column(scale=8, min_width=560, elem_classes=["app-panel"]):
                    with gr.Tabs(selected="build-kb", elem_id="main-tabs"):
                        # build first so the app opens in the ingestion workflow
                        with gr.Tab("Build Knowledge Base", id="build-kb"):
                            with gr.Column(elem_classes=["tab-body"]):
                                uploads = gr.File(
                                    label="Upload policy documents",
                                    file_count="multiple",
                                    file_types=[".pdf", ".docx", ".txt", ".md", ".html", ".htm", ".csv", ".json", ".jsonl", ".yaml", ".yml"],
                                )
                                with gr.Row():
                                    index_button = gr.Button("Index Uploaded Files", variant="primary", elem_id="index-button")
                                    clear_kb_button = gr.Button("Clear Knowledge Base")

                                gr.Markdown("### Subjective Federal Corpus:")
                                far_rows = gr.Slider(250, 6000, value=2500, step=250, label="FAR starter chunks (amount of FAR text segments that load into vdb)")
                                far_button = gr.Button(
                                    "Load FAR corpus (via Hugging Face)",
                                    variant="primary",
                                    elem_id="far-button",
                                )

                                suggested_questions = gr.Radio(
                                    choices=[],
                                    value=None,
                                    label="Suggested questions from the loaded FAR corpus",
                                    info="This list appears after indexing. Select a question to populate the Ask policy input.",
                                    interactive=True,
                                    visible=False,
                                    elem_id="far-suggestions",
                                )

                                gr.HTML(
                                    "<div class='warning-note'>Public Space warning: do not upload classified, CUI, source-selection-sensitive, procurement-sensitive, proprietary, or personal data to a public deployment.</div>"
                                )

                        # ask second so it appears on the right of Build knowledge base
                        with gr.Tab("Ask Policy", id="ask-policy"):
                            with gr.Column(elem_classes=["tab-body"]):
                                question = gr.Textbox(
                                    label="Policy question",
                                    placeholder="Example: Under the indexed FAR material, when is market research required and what should it address?",
                                    lines=3,
                                    max_lines=7,
                                )

                                with gr.Row(elem_classes=["ask-action-row"]):
                                    ask_button = gr.Button(
                                        "Ask with Citations",
                                        variant="primary",
                                        elem_id="ask-button",
                                        scale=0,
                                        min_width=210,
                                    )
                                    clear_answer = gr.Button(
                                        "Clear Answer",
                                        scale=0,
                                        min_width=118,
                                        elem_classes=["clear-answer-button"],
                                    )

                                answer = gr.Markdown("", label="Grounded answer", elem_classes=["answer-panel"])
                                gr.Markdown("### Retrieved sources")
                                sources = gr.HTML("<div class='status-card'>No sources retrieved.</div>")

                with gr.Column(scale=5, min_width=360):
                    kb_status = gr.HTML("<div class='status-card'><strong>Ready</strong><br>0 source documents · 0 indexed chunks</div>")
                    gr.HTML(
                        f"""
                        <div class="sidebar-card">
                          <h3>Model runtime</h3>
                          <p><strong>Active path:</strong><br>{html.escape(backend_label)}</p>
                          <p><strong>Local fallback:</strong><br>{html.escape(LOCAL_MODEL_ID)}</p>
                        </div>
                        <div class="sidebar-card">
                          <h3>How it works</h3>
                          <ol>
                            <li>Extract text and source metadata from uploaded files.</li>
                            <li>Chunk and embed passages, then index them in FAISS.</li>
                            <li>Retrieve the closest passages for the question.</li>
                            <li>Generate an answer constrained to those passages with [S#] citations.</li>
                          </ol>
                        </div>
                        """
                    )
                    with gr.Accordion("Retrieval controls", open=True):
                        top_k = gr.Slider(2, 12, value=DEFAULT_TOP_K, step=1, label="Retrieved passages")
                        temperature = gr.Slider(0.0, 0.8, value=0.15, step=0.05, label="Temperature")
                        max_new_tokens = gr.Slider(256, 1400, value=700, step=64, label="Maximum answer tokens")
                    with gr.Accordion("Diagnostics", open=False):
                        diagnostics = gr.JSON(value={}, label="RAG diagnostics")
                    gr.HTML(
                        """
                        <div class="sidebar-card">
                          <h3>Scope</h3>
                          <p>This app is designed for source-grounded research.</p>
                        </div>
                        """
                    )

            gr.HTML("<div class='footer-note'>Qwen · Sentence Transformers · FAISS · Hugging Face · Gradio</div>")

        index_button.click(index_uploads, inputs=[uploads, sid], outputs=[sid, kb_status, diagnostics])
        far_button.click(load_far, inputs=[sid, far_rows], outputs=[sid, kb_status, diagnostics, suggested_questions])
        suggested_questions.change(
            lambda selected: selected or "",
            inputs=[suggested_questions],
            outputs=[question],
        )
        clear_kb_button.click(clear_kb, inputs=[sid], outputs=[sid, kb_status, diagnostics, answer, sources, suggested_questions])
        ask_button.click(ask, inputs=[question, sid, top_k, temperature, max_new_tokens], outputs=[answer, sources, diagnostics])
        question.submit(ask, inputs=[question, sid, top_k, temperature, max_new_tokens], outputs=[answer, sources, diagnostics])
        clear_answer.click(lambda: ("", "<div class='status-card'>No sources retrieved.</div>", {}), outputs=[answer, sources, diagnostics])

    return demo


if __name__ == "__main__":
    app = build_app()
    app.queue(default_concurrency_limit=2).launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", "7860")),
        show_error=True,
        ssr_mode=False,
        theme=gr.themes.Base(),
        css=CSS,
        head=HEAD,
    )