import gradio as gr

from config import (
    APP_SUBTITLE,
    APP_TITLE,
    DEFAULT_CONCURRENCY_LIMIT,
    DEFAULT_JURISDICTION,
    JURISDICTIONS,
    MAX_RETRIEVAL_LOOPS,
    TOP_K_RETRIEVAL,
    UPLOAD_SESSION_ROOT,
)
from document_service import (
    clear_document_session,
    process_uploaded_document,
    validate_document_session,
)
from pipeline import run_document_rag, run_rag_agent
from job_service import (
    forget_rag_job,
    get_rag_job_status,
    start_document_processing_job,
    start_document_rag_job,
    start_rag_job,
)
from services import check_system_status


CUSTOM_CSS = """
.gradio-container {max-width:1280px!important;margin:0 auto!important;
font-family:"Segoe UI",Tahoma,Geneva,Verdana,sans-serif!important}
.app-header {text-align:center;padding:24px;
background:linear-gradient(135deg,#5f9eea 0%,#7bc5ef 100%);
color:white;border-radius:10px;margin-bottom:18px}
.app-header h1 {margin:0;font-size:2.35rem;font-weight:650}
.app-header p {margin:8px 0 0;font-size:1.05rem;opacity:.95}
.app-section {background:white;padding:18px;border-radius:10px;
box-shadow:0 3px 12px rgba(0,0,0,.07);margin-bottom:16px}
.answer-panel {border-left:4px solid #5f9eea}
.process-panel {border-left:4px solid #7a8ca5}
.source-panel {border-left:4px solid #65a97b}
.footer {text-align:center;padding:14px;margin-top:24px;color:#666;
font-size:.9rem;border-top:1px solid #e6e6e6}
"""

JURISDICTION_CHOICES = [
    (str(details["name"]), code)
    for code, details in JURISDICTIONS.items()
]


def clear_interface():
    return "", "", "", "", "", "Status: Ready"


with gr.Blocks(title=APP_TITLE, fill_width=True) as demo:
    with gr.Column(elem_classes="app-header"):
        gr.Markdown(f"# {APP_TITLE}")
        gr.Markdown(APP_SUBTITLE)

    with gr.Column(elem_classes="app-section"):
        gr.Markdown(
            "Select a jurisdiction, then enter a nutrient-policy, "
            "wastewater-management, or regulatory-compliance question. "
            "The agent will search the corresponding regulatory vector "
            "database and generate a grounded answer."
        )

        with gr.Row():
            jurisdiction_dropdown = gr.Dropdown(
                choices=JURISDICTION_CHOICES,
                value=DEFAULT_JURISDICTION,
                label="Jurisdiction",
                interactive=True,
            )

        question_box = gr.Textbox(
            label="Question",
            placeholder=(
                "Example: What phosphorus-related compliance requirements "
                "apply to municipal wastewater treatment facilities?"
            ),
            lines=4,
        )

        with gr.Row():
            top_k_slider = gr.Slider(
                1,
                15,
                value=TOP_K_RETRIEVAL,
                step=1,
                label="Passages retrieved per search",
            )
            max_loops_slider = gr.Slider(
                1,
                10,
                value=min(MAX_RETRIEVAL_LOOPS, 10),
                step=1,
                label="Maximum retrieval steps",
            )

        with gr.Row():
            run_button = gr.Button(
                "Run RAG Agent",
                variant="primary",
            )
            clear_button = gr.Button("Clear")
            status_button = gr.Button("Check System")

        system_status = gr.Textbox(
            value="Status: Ready",
            label="System status",
            lines=20,
            interactive=False,
        )


    with gr.Column(elem_classes="app-section"):
        gr.Markdown("## Uploaded document mode")
        gr.Markdown(
            "Process one document independently from the Florida and New Jersey "
            "regulatory databases. Supported formats are PDF, DOC, DOCX, PPTX, "
            "HTML, and TXT. The temporary session expires after 30 minutes "
            "of inactivity."
        )

        uploaded_document = gr.File(
            label="Document",
            file_types=[".pdf", ".doc", ".docx", ".pptx", ".html", ".txt"],
            type="filepath",
        )
        process_document_button = gr.Button(
            "Process document",
            variant="secondary",
        )
        document_requested_session_id = gr.Textbox(
            label="Requested document session ID",
            visible=False,
        )
        document_session_id = gr.Textbox(
            label="Document session ID",
            visible=False,
        )
        document_ready = gr.Checkbox(
            label="Document database ready",
            value=False,
            visible=False,
        )
        document_session_valid = gr.Checkbox(
            label="Document session valid",
            value=False,
            visible=False,
        )
        document_session_validation_message = gr.Textbox(
            label="Document session validation message",
            visible=False,
        )
        validate_document_session_button = gr.Button(
            "Validate Uploaded Document Session",
            visible=False,
        )
        document_status = gr.Markdown(
            "No document has been processed."
        )
        vector_download = gr.File(
            label="Download temporary vector database",
            interactive=False,
        )

        document_question = gr.Textbox(
            label="Question about the uploaded document",
            lines=3,
        )
        with gr.Row():
            run_document_button = gr.Button(
                "Ask Uploaded Document",
                variant="primary",
            )
            clear_document_button = gr.Button(
                "Clear Uploaded Document Session"
            )

    # Hidden components used by the resumable job API.
    resumable_requested_job_id = gr.Textbox(
        visible=False,
        label="Requested resumable job ID",
    )
    resumable_job_id = gr.Textbox(
        visible=False,
        label="Resumable job ID",
    )
    resumable_job_status = gr.Textbox(
        visible=False,
        label="Resumable job status JSON",
    )
    resumable_start_button = gr.Button(
        "Start resumable state RAG job",
        visible=False,
    )
    resumable_start_document_button = gr.Button(
        "Start resumable document RAG job",
        visible=False,
    )
    resumable_start_document_processing_button = gr.Button(
        "Start resumable document-processing job",
        visible=False,
    )
    resumable_status_button = gr.Button(
        "Get resumable RAG job status",
        visible=False,
    )
    resumable_forget_button = gr.Button(
        "Forget resumable RAG job",
        visible=False,
    )

    with gr.Column(
        elem_classes=["app-section", "answer-panel"]
    ):
        gr.Markdown("## Grounded answer")
        answer_output = gr.Markdown(
            "The final grounded answer will appear here."
        )

    with gr.Accordion(
        "Retrieval process",
        open=True,
        elem_classes=["app-section", "process-panel"],
    ):
        process_output = gr.Markdown(
            "No retrieval has been run."
        )

    with gr.Accordion(
        "Retrieved source passages",
        open=False,
        elem_classes=["app-section", "source-panel"],
    ):
        sources_output = gr.Markdown(
            "No sources have been retrieved."
        )

    with gr.Accordion("Run diagnostics", open=False):
        diagnostics_output = gr.Textbox(
            label="Diagnostics",
            lines=12,
            interactive=False,
        )

    with gr.Column(elem_classes="footer"):
        gr.Markdown(
            "WNP-ARAG | Built with Gradio, FAISS, "
            "LangChain, and Llama 3.1"
        )

    inputs = [
        question_box,
        jurisdiction_dropdown,
        top_k_slider,
        max_loops_slider,
    ]
    outputs = [
        answer_output,
        process_output,
        sources_output,
        diagnostics_output,
    ]

    run_button.click(
        run_rag_agent,
        inputs,
        outputs,
        api_name="run_rag_agent",
        api_visibility="public",
    )

    resumable_start_button.click(
        start_rag_job,
        inputs=[
            question_box,
            jurisdiction_dropdown,
            top_k_slider,
            max_loops_slider,
            resumable_requested_job_id,
        ],
        outputs=resumable_job_id,
        api_name="start_rag_job",
        api_visibility="public",
        queue=False,
    )

    resumable_start_document_processing_button.click(
        start_document_processing_job,
        inputs=[
            uploaded_document,
            document_requested_session_id,
            resumable_requested_job_id,
        ],
        outputs=resumable_job_id,
        api_name="start_document_processing_job",
        api_visibility="public",
        queue=False,
    )

    resumable_start_document_button.click(
        start_document_rag_job,
        inputs=[
            document_question,
            document_session_id,
            top_k_slider,
            max_loops_slider,
            resumable_requested_job_id,
        ],
        outputs=resumable_job_id,
        api_name="start_document_rag_job",
        api_visibility="public",
        queue=False,
    )

    resumable_status_button.click(
        get_rag_job_status,
        inputs=resumable_job_id,
        outputs=resumable_job_status,
        api_name="get_rag_job_status",
        api_visibility="public",
        queue=False,
    )

    resumable_forget_button.click(
        forget_rag_job,
        inputs=resumable_job_id,
        outputs=resumable_job_status,
        api_name="forget_rag_job",
        api_visibility="public",
        queue=False,
    )

    question_box.submit(
        run_rag_agent,
        inputs,
        outputs,
        api_visibility="private",
    )


    process_document_button.click(
        process_uploaded_document,
        inputs=[
            uploaded_document,
            document_requested_session_id,
        ],
        outputs=[
            document_session_id,
            document_status,
            vector_download,
            document_ready,
        ],
        api_name="process_uploaded_document",
        api_visibility="public",
    )

    run_document_button.click(
        run_document_rag,
        inputs=[
            document_question,
            document_session_id,
            top_k_slider,
            max_loops_slider,
        ],
        outputs=outputs,
        api_name="run_document_rag",
        api_visibility="public",
    )

    validate_document_session_button.click(
        validate_document_session,
        inputs=document_session_id,
        outputs=[
            document_session_valid,
            document_session_validation_message,
        ],
        api_name="validate_document_session",
        api_visibility="public",
        queue=False,
    )

    clear_document_button.click(
        clear_document_session,
        inputs=document_session_id,
        outputs=document_status,
        api_name="clear_document_session",
        api_visibility="public",
        queue=False,
    )

    clear_button.click(
        clear_interface,
        outputs=[
            question_box,
            answer_output,
            process_output,
            sources_output,
            diagnostics_output,
            system_status,
        ],
        api_name="clear_interface",
        queue=False,
    )

    status_button.click(
        check_system_status,
        outputs=system_status,
        api_name="check_system_status",
        queue=False,
    ) 


if __name__ == "__main__":
    demo.queue(
        default_concurrency_limit=DEFAULT_CONCURRENCY_LIMIT
    ).launch(
        debug=True,
        show_error=True,
        theme=gr.themes.Glass(primary_hue="blue"),
        css=CUSTOM_CSS,
        allowed_paths=[str(UPLOAD_SESSION_ROOT.resolve())],
        ssr_mode=False,
    )