import gradio as gr
import openai
import math
import os
from pypdf import PdfReader

# ── CONFIG ───────────────────────────────────────────────────────────────────
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
client         = openai.OpenAI(api_key=OPENAI_API_KEY)

EMBED_MODEL = "text-embedding-3-small"
CHAT_MODEL  = "gpt-4o-mini"
CHUNK_SIZE  = 100
OVERLAP     = 60
TOP_K       = 2

# ── GLOBALS ──────────────────────────────────────────────────────────────────
stored_chunks     = []
stored_embeddings = []
doc_name          = ""

# ── CORE LOGIC ───────────────────────────────────────────────────────────────
def split_chunks(text, size=CHUNK_SIZE, overlap=OVERLAP):
    words, chunks, start = text.split(), [], 0
    while start < len(words):
        chunks.append(" ".join(words[start:start + size]))
        start += size - overlap
    return chunks

def embed(text):
    return client.embeddings.create(model=EMBED_MODEL, input=text).data[0].embedding

def cosine_similarity(a, b):
    dot   = sum(x * y for x, y in zip(a, b))
    normA = math.sqrt(sum(x * x for x in a))
    normB = math.sqrt(sum(x * x for x in b))
    return dot / (normA * normB)

def get_top_chunks(question, k=TOP_K):
    q_vec  = embed(question)
    scores = [(cosine_similarity(q_vec, cv), c)
              for cv, c in zip(stored_embeddings, stored_chunks)]
    scores.sort(reverse=True)
    return [c for _, c in scores[:k]]

def upload_pdf(file):
    global stored_chunks, stored_embeddings, doc_name
    if file is None:
        return (
            "— No file selected.",
            gr.update(interactive=False)
        )

    reader    = PdfReader(file.name)
    full_text = "".join(page.extract_text() or "" for page in reader.pages)
    doc_name  = file.name.split("/")[-1]

    stored_chunks     = split_chunks(full_text)
    stored_embeddings = [embed(c) for c in stored_chunks]

    return (
        f"✦  {doc_name}\n→  {len(stored_chunks)} chunks indexed and ready.",
        gr.update(interactive=True)
    )

def chat(question, history):
    if not stored_chunks:
        history.append({"role": "user", "content": question})
        history.append({"role": "assistant", "content": "Please index a document in the Upload tab first."})
        return history
    
    context = "\n\n".join([f"[{i+1}] {c}" for i, c in enumerate(get_top_chunks(question))])
    prompt  = (
        "Answer strictly from the document text below.\n"
        "If not found, reply: \"This information is not in the document.\"\n\n"
        f"TEXT:\n{context}\n\nQUESTION: {question}\n\nANSWER:"
    )
    response = client.chat.completions.create(
        model=CHAT_MODEL,
        messages=[
            {"role": "system", "content":
                "Strict document assistant. Only use provided text. "
                "If answer absent, say: This information is not in the document."},
            {"role": "user", "content": prompt}
        ]
    )
    history.append({"role": "user", "content": question})
    history.append({"role": "assistant", "content": response.choices[0].message.content})
    return history


# ── CSS ───────────────────────────────────────────────────────────────────────
css = """
@import url('https://fonts.googleapis.com/css2?family=DM+Serif+Display:ital@0;1&family=DM+Mono:wght@300;400;500&family=DM+Sans:ital,opsz,wght@0,9..40,300;0,9..40,400;0,9..40,500;1,9..40,300&display=swap');

*, *::before, *::after { box-sizing: border-box; margin: 0; padding: 0; }

:root {
    --bg:          #0b0b0d;
    --surface:     #131318;
    --surface-2:   #18181f;
    --border:      rgba(255,255,255,0.06);
    --border-2:    rgba(255,255,255,0.10);
    --accent:      #b8a98a;
    --accent-glow: rgba(184,169,138,0.09);
    --green:       #52a97a;
    --green-glow:  rgba(82,169,122,0.10);
    --text-1:      #eeeae2;
    --text-2:      #75757f;
    --text-3:      #3e3e48;
    --font-serif:  'DM Serif Display', Georgia, serif;
    --font-sans:   'DM Sans', sans-serif;
    --font-mono:   'DM Mono', monospace;
    --r:           10px;
    --r-lg:        18px;
    --ease:        cubic-bezier(0.4,0,0.2,1);
}

html, body, .gradio-container, #root {
    font-family: var(--font-sans) !important;
    background:  var(--bg) !important;
    color:       var(--text-1) !important;
}

.gradio-container {
    max-width: 1080px !important;
    margin: 0 auto !important;
    padding: 44px 28px 68px !important;
}

::-webkit-scrollbar { width: 3px; }
::-webkit-scrollbar-track { background: transparent; }
::-webkit-scrollbar-thumb { background: var(--border-2); border-radius: 3px; }

/* ── MASTHEAD ── */
.mast {
    display: flex;
    align-items: flex-start;
    justify-content: space-between;
    padding-bottom: 26px;
    border-bottom: 1px solid var(--border);
    margin-bottom: 32px;
}
.mast-left {}
.mast-title {
    font-family: var(--font-serif) !important;
    font-size: 2.6rem;
    font-weight: 400;
    color: var(--text-1);
    letter-spacing: -0.025em;
    line-height: 1.05;
}
.mast-title em { font-style: italic; color: var(--accent); }
.mast-sub {
    font-family: var(--font-mono);
    font-size: 0.6rem;
    font-weight: 300;
    color: var(--text-3);
    letter-spacing: 0.14em;
    text-transform: uppercase;
    margin-top: 9px;
}

/* ── PORTFOLIO BADGE ── */
.portfolio-badge {
    display: flex;
    flex-direction: column;
    align-items: flex-end;
    gap: 6px;
}
.badge-inner {
    display: flex;
    align-items: center;
    gap: 8px;
    background: var(--surface);
    border: 1px solid var(--border-2);
    border-radius: 8px;
    padding: 9px 14px 9px 10px;
}
.badge-dot {
    width: 7px;
    height: 7px;
    border-radius: 50%;
    background: var(--green);
    box-shadow: 0 0 6px var(--green-glow);
    animation: pulse 2.4s ease-in-out infinite;
}
@keyframes pulse {
    0%, 100% { opacity: 1; box-shadow: 0 0 6px rgba(82,169,122,0.4); }
    50%       { opacity: 0.5; box-shadow: 0 0 12px rgba(82,169,122,0.2); }
}
.badge-text {
    font-family: var(--font-mono);
    font-size: 0.65rem;
    font-weight: 500;
    color: var(--text-1);
    letter-spacing: 0.04em;
}
.badge-sub {
    font-family: var(--font-mono);
    font-size: 0.58rem;
    color: var(--text-3);
    letter-spacing: 0.1em;
    text-transform: uppercase;
}
.stack-chips { display: flex; gap: 5px; }
.chip {
    font-family: var(--font-mono);
    font-size: 0.56rem;
    font-weight: 400;
    color: var(--text-3);
    background: var(--surface);
    border: 1px solid var(--border);
    border-radius: 4px;
    padding: 3px 7px;
    letter-spacing: 0.08em;
    text-transform: uppercase;
}

/* ── TABS ── */
.tabs > .tab-nav {
    border-bottom: 1px solid var(--border) !important;
    background: transparent !important;
    margin-bottom: 28px !important;
    gap: 0 !important;
}
.tabs > .tab-nav button {
    font-family: var(--font-mono) !important;
    font-size: 0.65rem !important;
    font-weight: 500 !important;
    letter-spacing: 0.12em !important;
    text-transform: uppercase !important;
    color: var(--text-3) !important;
    background: transparent !important;
    border: none !important;
    border-bottom: 2px solid transparent !important;
    padding: 10px 20px !important;
    border-radius: 0 !important;
    transition: color 0.2s var(--ease), border-color 0.2s var(--ease) !important;
    margin-bottom: -1px !important;
}
.tabs > .tab-nav button:hover {
    color: var(--text-2) !important;
}
.tabs > .tab-nav button.selected {
    color: var(--accent) !important;
    border-bottom-color: var(--accent) !important;
    background: transparent !important;
}

/* ── PANEL ── */
.panel {
    background: var(--surface);
    border: 1px solid var(--border);
    border-radius: var(--r-lg);
    padding: 28px 26px;
}

/* ── EYEBROW ── */
.eyebrow {
    font-family: var(--font-mono);
    font-size: 0.56rem;
    font-weight: 500;
    letter-spacing: 0.18em;
    text-transform: uppercase;
    color: var(--text-3);
    margin-bottom: 16px;
    display: flex;
    align-items: center;
    gap: 10px;
}
.eyebrow::after {
    content: '';
    flex: 1;
    height: 1px;
    background: var(--border);
}

/* ── UPLOAD TAB LAYOUT ── */
.upload-grid {
    display: grid;
    grid-template-columns: 1fr 1fr;
    gap: 20px;
    align-items: start;
}
.upload-right {
    display: flex;
    flex-direction: column;
    gap: 16px;
}

/* ── SPEC GRID ── */
.specs {
    display: grid;
    grid-template-columns: 1fr 1fr;
    gap: 7px;
}
.spec {
    background: var(--bg);
    border: 1px solid var(--border);
    border-radius: 8px;
    padding: 11px 13px;
}
.spec-k {
    font-family: var(--font-mono);
    font-size: 0.54rem;
    letter-spacing: 0.12em;
    text-transform: uppercase;
    color: var(--text-3);
    margin-bottom: 5px;
}
.spec-v {
    font-family: var(--font-mono);
    font-size: 0.76rem;
    color: var(--accent);
    font-weight: 400;
}
.spec.wide { grid-column: span 2; }

/* ── ARCHITECTURE BOX ── */
.arch-box {
    background: var(--bg);
    border: 1px solid var(--border);
    border-radius: 10px;
    padding: 16px;
}
.arch-box .arch-title {
    font-family: var(--font-mono);
    font-size: 0.56rem;
    letter-spacing: 0.16em;
    text-transform: uppercase;
    color: var(--text-3);
    margin-bottom: 12px;
}
.arch-step {
    display: flex;
    align-items: flex-start;
    gap: 12px;
    margin-bottom: 10px;
}
.arch-step:last-child { margin-bottom: 0; }
.arch-num {
    font-family: var(--font-mono);
    font-size: 0.62rem;
    color: var(--accent);
    min-width: 16px;
    margin-top: 1px;
}
.arch-desc {
    font-family: var(--font-sans);
    font-size: 0.78rem;
    font-weight: 300;
    color: var(--text-2);
    line-height: 1.5;
}

/* ── FILE INPUT ── */
.gr-file {
    border: 1px dashed var(--border-2) !important;
    border-radius: var(--r) !important;
    background: var(--surface-2) !important;
    transition: border-color 0.2s var(--ease), background 0.2s var(--ease) !important;
}
.gr-file:hover {
    border-color: var(--accent) !important;
    background: var(--accent-glow) !important;
}

/* ── INPUTS ── */
textarea, input[type="text"] {
    font-family: var(--font-sans) !important;
    font-size: 0.88rem !important;
    font-weight: 300 !important;
    color: var(--text-1) !important;
    background: var(--surface-2) !important;
    border: 1px solid var(--border) !important;
    border-radius: var(--r) !important;
    padding: 11px 15px !important;
    line-height: 1.65 !important;
    transition: border-color 0.2s var(--ease), box-shadow 0.2s var(--ease) !important;
    caret-color: var(--accent) !important;
}
textarea:focus, input[type="text"]:focus {
    border-color: var(--border-2) !important;
    box-shadow: 0 0 0 3px var(--accent-glow) !important;
    outline: none !important;
    background: var(--surface) !important;
}
textarea::placeholder, input::placeholder {
    color: var(--text-3) !important;
    font-weight: 300 !important;
}

.status-box textarea {
    font-family: var(--font-mono) !important;
    font-size: 0.73rem !important;
    color: var(--text-2) !important;
    background: var(--bg) !important;
    border: 1px solid var(--border) !important;
    line-height: 1.9 !important;
}

/* ── BUTTONS ── */
button { font-family: var(--font-sans) !important; cursor: pointer !important; }
button.primary {
    background: var(--accent) !important;
    color: #0b0b0d !important;
    border: none !important;
    border-radius: 7px !important;
    padding: 10px 22px !important;
    font-size: 0.8rem !important;
    font-weight: 500 !important;
    letter-spacing: 0.02em !important;
    transition: all 0.18s var(--ease) !important;
}
button.primary:hover {
    background: #ccbea0 !important;
    transform: translateY(-1px) !important;
    box-shadow: 0 6px 20px rgba(184,169,138,0.2) !important;
}
button.primary:active { transform: translateY(0) !important; box-shadow: none !important; }
button.primary:disabled {
    background: var(--surface-2) !important;
    color: var(--text-3) !important;
    box-shadow: none !important;
    cursor: not-allowed !important;
    transform: none !important;
}

button.secondary {
    background: transparent !important;
    color: var(--text-2) !important;
    border: 1px solid var(--border-2) !important;
    border-radius: 7px !important;
    padding: 10px 18px !important;
    font-size: 0.8rem !important;
    font-weight: 400 !important;
    transition: all 0.18s var(--ease) !important;
}
button.secondary:hover {
    color: var(--text-1) !important;
    border-color: var(--text-3) !important;
}

/* ── CHATBOT ── */
.chatbot-wrap > div {
    background: var(--bg) !important;
    border: 1px solid var(--border) !important;
    border-radius: var(--r) !important;
    padding: 10px !important;
}
.message-wrap .user-row .message {
    background: var(--surface-2) !important;
    color: var(--text-1) !important;
    border: 1px solid var(--border-2) !important;
    border-radius: 14px 14px 3px 14px !important;
    font-size: 0.85rem !important;
    font-weight: 300 !important;
    line-height: 1.7 !important;
    padding: 11px 15px !important;
    max-width: 76% !important;
    margin-left: auto !important;
}
.message-wrap .bot-row .message {
    background: var(--surface) !important;
    color: var(--text-1) !important;
    border: 1px solid var(--border) !important;
    border-radius: 14px 14px 14px 3px !important;
    font-size: 0.85rem !important;
    font-weight: 300 !important;
    line-height: 1.7 !important;
    padding: 11px 15px !important;
    max-width: 76% !important;
}

/* ── DIVIDER ── */
.div { border: none; border-top: 1px solid var(--border); margin: 22px 0; }

/* ── LABELS ── */
label > span, .gr-form label span {
    font-family: var(--font-mono) !important;
    font-size: 0.56rem !important;
    font-weight: 500 !important;
    letter-spacing: 0.12em !important;
    text-transform: uppercase !important;
    color: var(--text-3) !important;
}

/* ── FOOTER ── */
.footer {
    text-align: center;
    margin-top: 44px;
    padding-top: 20px;
    border-top: 1px solid var(--border);
    display: flex;
    align-items: center;
    justify-content: center;
    gap: 16px;
}
.footer p {
    font-family: var(--font-mono);
    font-size: 0.57rem;
    color: var(--text-3);
    letter-spacing: 0.1em;
    text-transform: uppercase;
}
.footer-sep { color: var(--text-3); font-size: 0.5rem; }
"""

# ── UI ────────────────────────────────────────────────────────────────────────
with gr.Blocks(css=css, title="Document Intelligence — Jason Kishore") as demo:

    # ── Masthead ─────────────────────────────────────────────────────────────
    gr.HTML("""
    <div class="mast">
        <div class="mast-left">
            <div class="mast-title">Document <em>Intelligence</em></div>
            <div class="mast-sub">Retrieval-Augmented Generation · Semantic Search · GPT-4o Mini</div>
        </div>
        <div class="portfolio-badge">
            <div class="badge-inner">
                <div class="badge-dot"></div>
                <div>
                    <div class="badge-text">Jason Kishore</div>
                    <div class="badge-sub" style="margin-top:2px">Portfolio Project</div>
                </div>
            </div>
            <div class="stack-chips">
                <span class="chip">Python</span>
                <span class="chip">OpenAI</span>
                <span class="chip">Gradio</span>
                <span class="chip">RAG</span>
            </div>
        </div>
    </div>
    """)

    # ── Tabs ─────────────────────────────────────────────────────────────────
    with gr.Tabs(elem_classes=["tabs"]):

        # ── Tab 1 · Upload ───────────────────────────────────────────────────
        with gr.TabItem("01 · Upload & Index"):
            with gr.Column(elem_classes=["panel"]):

                gr.HTML('<div class="eyebrow">Document Ingestion</div>')

                gr.HTML('<div class="upload-grid">')

                with gr.Column():
                    pdf_input = gr.File(label="Upload PDF", file_types=[".pdf"])
                    upload_btn = gr.Button("Index Document", variant="primary")
                    upload_status = gr.Textbox(
                        label="Status",
                        lines=2,
                        interactive=False,
                        value="— Awaiting document",
                        elem_classes=["status-box"]
                    )

                with gr.Column():
                    gr.HTML("""
                    <div class="upload-right">
                        <div>
                            <div class="eyebrow">Configuration</div>
                            <div class="specs">
                                <div class="spec">
                                    <div class="spec-k">Embed Model</div>
                                    <div class="spec-v">3-small</div>
                                </div>
                                <div class="spec">
                                    <div class="spec-k">Chat Model</div>
                                    <div class="spec-v">4o-mini</div>
                                </div>
                                <div class="spec">
                                    <div class="spec-k">Chunk Size</div>
                                    <div class="spec-v">100 words</div>
                                </div>
                                <div class="spec">
                                    <div class="spec-k">Overlap</div>
                                    <div class="spec-v">40 words</div>
                                </div>
                                <div class="spec wide">
                                    <div class="spec-k">Top-K Retrieval</div>
                                    <div class="spec-v">2 chunks / query</div>
                                </div>
                            </div>
                        </div>
                        <div class="arch-box">
                            <div class="arch-title">Pipeline</div>
                            <div class="arch-step">
                                <div class="arch-num">01</div>
                                <div class="arch-desc">PDF text extracted page-by-page via pypdf</div>
                            </div>
                            <div class="arch-step">
                                <div class="arch-num">02</div>
                                <div class="arch-desc">Text split into overlapping 100-word chunks</div>
                            </div>
                            <div class="arch-step">
                                <div class="arch-num">03</div>
                                <div class="arch-desc">Each chunk embedded with text-embedding-3-small</div>
                            </div>
                            <div class="arch-step">
                                <div class="arch-num">04</div>
                                <div class="arch-desc">Query embedded → top-2 chunks retrieved by cosine similarity</div>
                            </div>
                            <div class="arch-step">
                                <div class="arch-num">05</div>
                                <div class="arch-desc">GPT-4o Mini answers strictly from retrieved context</div>
                            </div>
                        </div>
                    </div>
                    """)

                gr.HTML('</div>')  # close upload-grid

        # ── Tab 2 · Chat ─────────────────────────────────────────────────────
        with gr.TabItem("02 · Chat"):
            with gr.Column(elem_classes=["panel"]):

                gr.HTML('<div class="eyebrow">Conversation</div>')

                chatbot = gr.Chatbot(
                    height=440,
                    show_label=False,
                    elem_classes=["chatbot-wrap"],
                    #type="messages",          # Already recommended earlier — ensures dict format
                    render_markdown=False,    # ← Disable Markdown if you don't need bold/italics/links
                    line_breaks=False,        # ← Prevents GitHub-style single-\n breaks
                    #bubble_full_width=False,  # Optional: was removed in Gradio 6, but set if using older version
                    placeholder=...
                )

                # chatbot = gr.Chatbot(
                #     height=440,
                #     show_label=False,
                #     elem_classes=["chatbot-wrap"],
                #     placeholder=(
                #         "<div style='font-family:DM Mono,monospace;font-size:0.62rem;"
                #         "color:#3e3e48;text-align:center;padding:80px 0;"
                #         "letter-spacing:0.14em;text-transform:uppercase;line-height:2'>"
                #         "Index a document in the Upload tab<br>to begin your session</div>"
                #     )
                # )

                msg_box = gr.Textbox(
                    placeholder="Ask a question about your document…",
                    label="Query",
                    lines=1,
                    max_lines=5,
                )

                with gr.Row():
                    send_btn  = gr.Button("Send Query", variant="primary",   size="sm")
                    clear_btn = gr.Button("Clear",      variant="secondary", size="sm")

    # ── Footer ───────────────────────────────────────────────────────────────
    gr.HTML("""
    <div class="footer">
        <p>Jason Kishore · Portfolio Project</p>
        <span class="footer-sep">◆</span>
        <p>Gradio · OpenAI · pypdf</p>
        <span class="footer-sep">◆</span>
        <p>Answers grounded strictly in uploaded document</p>
    </div>
    """)

    # ── Events ───────────────────────────────────────────────────────────────
    upload_btn.click(
        fn=upload_pdf,
        inputs=pdf_input,
        outputs=[upload_status, send_btn]
    )

    send_btn.click(
        fn=chat, inputs=[msg_box, chatbot], outputs=chatbot
    ).then(lambda: "", outputs=msg_box)

    msg_box.submit(
        fn=chat, inputs=[msg_box, chatbot], outputs=chatbot
    ).then(lambda: "", outputs=msg_box)

    clear_btn.click(lambda: [], outputs=chatbot)

demo.launch()