import base64

import gradio as gr
from mistralai import Mistral
from openai import OpenAI

# ─── Helpers ──────────────────────────────────────────────────────────────────


def _run_ocr(api_key: str, pdf_bytes: bytes | None = None, url: str | None = None):
    """Call Mistral OCR and return the raw response."""
    client = Mistral(api_key=api_key)

    if pdf_bytes:
        upload = client.files.upload(
            file={"file_name": "document.pdf", "content": pdf_bytes},
            purpose="ocr",
        )
        signed = client.files.get_signed_url(file_id=upload.id)
        doc = {"type": "document_url", "document_url": signed.url}
    else:
        doc = {"type": "document_url", "document_url": url}

    return client.ocr.process(
        model="mistral-ocr-latest",
        document=doc,
        include_image_base64=True,
    )


def _combine_markdown(ocr_response) -> str:
    """Concatenate all pages' markdown with page dividers."""
    parts = []
    for i, page in enumerate(ocr_response.pages):
        if len(ocr_response.pages) > 1:
            parts.append(f"\n\n---\n*— Page {i + 1} —*\n\n")
        parts.append(page.markdown or "")
    return "".join(parts)


def _run_chat(
    api_key: str,
    markdown_content: str,
    question: str,
    history: list[tuple[str, str]],
) -> str:
    """Send a question to GPT-4o with the OCR markdown as context."""
    client = OpenAI(api_key=api_key)

    system_msg = (
        "You are a helpful assistant that answers questions about a document that was "
        "processed with an OCR model. Base your answers solely on the document content "
        "provided below.\n\n"
        "--- DOCUMENT CONTENT (OCR) ---\n"
        f"{markdown_content}\n"
        "--- END OF DOCUMENT ---"
    )
    messages = [{"role": "system", "content": system_msg}]
    for q, a in history:
        messages.append({"role": "user", "content": q})
        messages.append({"role": "assistant", "content": a})
    messages.append({"role": "user", "content": question})

    resp = client.chat.completions.create(model="gpt-4o", messages=messages)
    return resp.choices[0].message.content


def _pdf_preview_html(pdf_file: str | None, pdf_url: str) -> str:
    """Build an inline iframe preview for either an uploaded file or a URL."""
    if pdf_file:
        with open(pdf_file, "rb") as f:
            b64 = base64.b64encode(f.read()).decode()
        src = f"data:application/pdf;base64,{b64}"
    elif pdf_url.strip():
        src = pdf_url.strip()
    else:
        return ""
    return f'<iframe src="{src}" width="100%" height="600px" style="border:none;"></iframe>'


def _has_document(pdf_file, pdf_url) -> bool:
    return bool(pdf_file) or bool((pdf_url or "").strip())


# ─── Callbacks ────────────────────────────────────────────────────────────────


def on_document_change(pdf_file, pdf_url):
    """Reset OCR/chat state whenever the document changes, and refresh preview."""
    preview = _pdf_preview_html(pdf_file, pdf_url)
    preview_visible = gr.update(visible=bool(preview), value=preview)
    chat_visible = gr.update(visible=_has_document(pdf_file, pdf_url))
    return preview_visible, chat_visible, None, False, []


def clear_chat():
    return None, False, []


def ask(
    mistral_key,
    openai_key,
    pdf_file,
    pdf_url,
    question,
    chat_display,
    ocr_markdown,
    ocr_done,
    history,
):
    question = (question or "").strip()
    if not question:
        return chat_display, ocr_markdown, ocr_done, history, ""
    if not mistral_key or not openai_key:
        raise gr.Error("Add your Mistral and OpenAI API keys in the sidebar.")
    if not _has_document(pdf_file, pdf_url):
        raise gr.Error("Select a lab report PDF in the sidebar.")

    try:
        if not ocr_done:
            pdf_bytes = None
            if pdf_file:
                with open(pdf_file, "rb") as f:
                    pdf_bytes = f.read()
            url = pdf_url.strip() if not pdf_bytes else None

            result = _run_ocr(mistral_key, pdf_bytes=pdf_bytes, url=url)
            ocr_markdown = _combine_markdown(result)
            ocr_done = True

        answer = _run_chat(openai_key, ocr_markdown, question, history)
        history = history + [(question, answer)]
        chat_display = chat_display + [
            {"role": "user", "content": question},
            {"role": "assistant", "content": answer},
        ]
    except Exception as exc:
        raise gr.Error(f"Request failed: {exc}")

    return chat_display, ocr_markdown, ocr_done, history, ""


# ─── UI ───────────────────────────────────────────────────────────────────────

with gr.Blocks(title="Lab Results OCR Extraction") as demo:
    ocr_markdown_state = gr.State(None)
    ocr_done_state = gr.State(False)
    history_state = gr.State([])

    gr.Markdown("# Lab Results Information Extraction")
    gr.Markdown(
        "Extract structured information from lab result PDFs using Mistral OCR, "
        "then ask GPT-4o questions about the extracted findings."
    )

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("## Lab Results OCR")
            gr.Markdown("Configure your keys, then upload a lab report PDF.")

            gr.Markdown("### API Keys")
            mistral_key = gr.Textbox(
                label="Mistral API Key",
                type="password",
                placeholder="Enter your Mistral API key…",
                info="Required for OCR. Get yours at console.mistral.ai",
            )
            openai_key = gr.Textbox(
                label="OpenAI API Key",
                type="password",
                placeholder="Enter your OpenAI API key…",
                info="Required for the chat section. Get yours at platform.openai.com",
            )

            gr.Markdown("### Lab Report")
            input_method = gr.Radio(
                ["Upload PDF", "Enter URL"],
                value="Upload PDF",
                show_label=False,
            )
            pdf_file = gr.File(
                label="Upload PDF",
                file_types=[".pdf"],
                type="filepath",
                visible=True,
            )
            pdf_url = gr.Textbox(
                label="PDF URL",
                placeholder="https://example.com/lab-results.pdf",
                show_label=False,
                visible=False,
            )

        with gr.Column(scale=2):
            with gr.Accordion("View uploaded PDF", open=False, visible=False) as pdf_accordion:
                pdf_preview = gr.HTML()

            chat_col = gr.Column(visible=False)
            with chat_col:
                with gr.Row():
                    gr.Markdown("### Describe what information to extract")
                    clear_btn = gr.Button("Clear chat", size="sm", scale=0)

                chatbot = gr.Chatbot(show_label=False)
                question_box = gr.Textbox(
                    placeholder="e.g. Which lab values are outside the reference range?",
                    show_label=False,
                    lines=3,
                )
                submit_btn = gr.Button("Submit", variant="primary")

    # Toggle upload-vs-url inputs
    input_method.change(
        lambda m: (gr.update(visible=m == "Upload PDF"), gr.update(visible=m == "Enter URL")),
        inputs=input_method,
        outputs=[pdf_file, pdf_url],
    )

    # Reset OCR/chat state and refresh preview when the document changes
    for comp in (pdf_file, pdf_url):
        comp.change(
            on_document_change,
            inputs=[pdf_file, pdf_url],
            outputs=[pdf_preview, chat_col, ocr_markdown_state, ocr_done_state, history_state],
        ).then(lambda: [], outputs=chatbot)
        comp.change(
            lambda f, u: gr.update(visible=_has_document(f, u)),
            inputs=[pdf_file, pdf_url],
            outputs=pdf_accordion,
        )

    clear_btn.click(
        clear_chat,
        outputs=[ocr_markdown_state, ocr_done_state, history_state],
    ).then(lambda: [], outputs=chatbot)

    ask_inputs = [
        mistral_key,
        openai_key,
        pdf_file,
        pdf_url,
        question_box,
        chatbot,
        ocr_markdown_state,
        ocr_done_state,
        history_state,
    ]
    ask_outputs = [chatbot, ocr_markdown_state, ocr_done_state, history_state, question_box]

    submit_btn.click(ask, inputs=ask_inputs, outputs=ask_outputs)
    question_box.submit(ask, inputs=ask_inputs, outputs=ask_outputs)


FORCE_LIGHT_THEME_HEAD = """
<script>
(function () {
    var params = new URLSearchParams(window.location.search);
    if (!params.has("__theme")) {
        params.set("__theme", "light");
        window.location.replace(
            window.location.pathname + "?" + params.toString() + window.location.hash
        );
    }
})();
</script>
"""

if __name__ == "__main__":
    demo.launch(head=FORCE_LIGHT_THEME_HEAD)