# main.py
import gradio as gr
from gradio.themes import Soft

from question_runner import run_tool
from config import MODEL_PRIORITY, SYNTAX_DOC_URL, MORPHOLOGY_DOC_URL, SEMANTICS_DOC_URL
from doc_utils import get_questions_from_doc

# Estimate runtime based on # of questions
def estimate_runtime(passage, doc_type):
    if not passage or not doc_type:
        return ""
    if doc_type.lower() == "syntax":
        doc_url = SYNTAX_DOC_URL
    elif doc_type.lower() == "morphology":
        doc_url = MORPHOLOGY_DOC_URL
    else:  # semantics
        doc_url = SEMANTICS_DOC_URL
    questions = get_questions_from_doc(doc_url)
    if not questions or (isinstance(questions, list) and questions and str(questions[0]).startswith("Error")):
        return "Unable to load questions."
    est_seconds = round(len(questions) * 2.5, 1)
    return f"Estimated generation time: ~{est_seconds} seconds"

# Build the Gradio interface at module level (required for HF Spaces)
with gr.Blocks(theme=Soft()) as demo:
    gr.Markdown("""
    ## **Classical Language Query Assistant**
    Submit a Latin or Greek passage and select the question type.
    Answers are generated using a rotating chain of hosted AI models via OpenRouter.
    """)

    with gr.Row():
        passage_input = gr.Textbox(label="Latin or Greek Passage", lines=4)
        question_type = gr.Radio(["Syntax", "Morphology", "Semantics"], label="Question Type")

    top_model = MODEL_PRIORITY[0]
    full_model_list = "\n".join(f"- `{m}`" for m in MODEL_PRIORITY)
    gr.Markdown(f"""
**Currently prioritized model:** `{top_model}`  
**Model fallback chain (if needed):**  
{full_model_list}
""")

    estimated_time_box = gr.Textbox(label="Estimated Time", interactive=False)

    with gr.Row():
        output_text = gr.Textbox(label="Generated Answers", lines=25, interactive=False)
        output_file = gr.File(label="Download Answers (.txt)", interactive=False)

    passage_input.change(fn=estimate_runtime, inputs=[passage_input, question_type], outputs=estimated_time_box)
    question_type.change(fn=estimate_runtime, inputs=[passage_input, question_type], outputs=estimated_time_box)

    submit_button = gr.Button("Generate Answers")

    # Connect the button to run_tool function directly
    submit_button.click(
        fn=run_tool,
        inputs=[passage_input, question_type],
        outputs=[output_text, output_file, estimated_time_box],
        api_name="generate"
    )

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