import gradio as gr
from utils.template_discovery import discover_templates
from utils.data_generator import generate_examples
from llm_client import ask_llm
from utils.evaluation import compute_similarity, validate_answer
import asyncio
import os


print("Scanning for all available Engchain templates...")
ALL_TEMPLATES = discover_templates()
print("...Scan complete.")


def format_name(name):
    return name.replace("_", " ").replace(".py", "").replace("template ", "").title()


def call_llm(question):
    if not question:
        return "⚠️ No question provided."
    return ask_llm(question)


with gr.Blocks(theme=gr.themes.Soft(), css="""
    .example-card {
        border: 2px solid var(--border-color-primary);
        border-radius: 12px;
        padding: 15px;
        margin-bottom: 18px;
        background: var(--block-background-fill);
        box-shadow: 0 4px 8px rgba(0,0,0,0.07);
    }

    .question-box {
        background: var(--panel-background-fill);
        padding: 10px;
        border-radius: 8px;
        margin-bottom: 10px;
        color: var(--body-text-color);
    }

    .solution-box {
        background: var(--panel-background-fill);
        padding: 10px;
        border-radius: 8px;
        margin-bottom: 10px;
        color: var(--body-text-color);
    }

    .eval-box {
        display: flex;
        gap: 12px;
        margin-bottom: 10px;
    }
               
    .badge {
        padding: 6px 12px;
        border-radius: 8px;
        font-weight: bold;
        font-size: 0.9em;
        display: inline-block;
    }
               
    .badge-blue {
        background: #e3f2fd;
        color: #1565c0;
    }
               
    .badge-green {
        background: #e8f5e9;
        color: #2e7d32;
    }
               
    .badge-red {
        background: #ffebee;
        color: #c62828;
    }
               
    .badge span {
        color: black; /* text label always black */
    }

    .llm-response {
        background: var(--panel-background-fill);
        padding: 14px;
        border-radius: 8px;
        border: 1px solid var(--border-color-primary);
        margin-top: 5px;
    }

    .compare-btn {
        background: linear-gradient(90deg, #6a11cb, #2575fc) !important;
        color: white !important;
        font-weight: bold !important;
        border-radius: 8px !important;
        padding: 10px 14px !important;
        border: none !important;
        box-shadow: 0 2px 6px rgba(0,0,0,0.15) !important;
        transition: transform 0.1s ease-in-out;
    }
    .compare-btn:hover {
        transform: scale(1.05);
    }

    .dropdown-box {
        border: 1px solid var(--border-color-primary);
        border-radius: 10px;
        padding: 12px;
        margin-bottom: 16px;
        background: var(--panel-background-fill);
    }
""") as demo:

    gr.Markdown(
        """
        # 📈 Engchain Template Playground  
        Select a branch, domain, file, and template to generate 10 unique Q&A examples.  
        Compare your solutions against a live LLM and see similarity + correctness checks.
        """
    )

    with gr.Group(elem_classes="dropdown-box"):
        with gr.Row():
            branch_dropdown = gr.Dropdown(
                label="1. Select Branch",
                choices=[(format_name(b), b) for b in sorted(list(ALL_TEMPLATES.keys()))],
            )
            domain_dropdown = gr.Dropdown(label="2. Select Engineering Domain", interactive=False)
            file_dropdown = gr.Dropdown(label="3. Select Specific Area / File", interactive=False)
            template_dropdown = gr.Dropdown(label="4. Select Specific Tool / Template", interactive=False)

    generate_button = gr.Button("Generate 10 Examples", variant="primary")

    # Pre-create a slot for template heading
    template_heading = gr.Markdown("", visible=False, elem_classes="template-heading")

    # Pre-create 10 example slots
    example_blocks = []
    for i in range(10):
        with gr.Group(visible=False, elem_classes="example-card") as grp:
            q_md = gr.Markdown(elem_classes="question-box")
            s_md = gr.Markdown(elem_classes="solution-box")
            btn = gr.Button("⚖️ See an LLM's Take on this Question", elem_classes="compare-btn")

            with gr.Column(visible=False) as llm_group:
                eval_md = gr.HTML(elem_classes="eval-box")
                llm_response_md = gr.Markdown(elem_classes="llm-response")

            example_blocks.append((grp, q_md, s_md, btn, llm_group, eval_md, llm_response_md))

            async def reveal_llm(q, s, llm_group=llm_group, eval_md=eval_md, llm_response_md=llm_response_md):
                await asyncio.sleep(0)
                yield gr.update(visible=True), gr.update(value=""), gr.update(value="⏳ Fetching LLM response...")
                response = await asyncio.to_thread(call_llm, q)
                sim = compute_similarity(s, response)
                match_flag = validate_answer(s, response)
                eval_html = (
                    f"<div class='badge badge-blue'><span>Textual Similarity:</span> {sim}%</div>"
                    f"<div class='badge {'badge-green' if match_flag else 'badge-red'}'>"
                    f"<span>Final Answer Validation:</span> {'✅ Yes' if match_flag else '❌ No'}</div>"
                )
                yield gr.update(visible=True), gr.update(value=eval_html), gr.update(value=response)

            btn.click(reveal_llm, inputs=[q_md, s_md], outputs=[llm_group, eval_md, llm_response_md])

    # Dropdown update logic
    def update_domain_dropdown(branch):
        if not branch:
            return gr.update(choices=[], value=None, interactive=False), \
                   gr.update(choices=[], value=None, interactive=False), \
                   gr.update(choices=[], value=None, interactive=False)
        domains = sorted(list(ALL_TEMPLATES.get(branch, {}).keys()))
        formatted_domains = [(format_name(d), d) for d in domains]
        return gr.update(choices=formatted_domains, value=None, interactive=True), \
               gr.update(choices=[], value=None, interactive=False), \
               gr.update(choices=[], value=None, interactive=False)

        
    def update_file_dropdown(branch, domain):
        if not branch or not domain:
            return (
                gr.update(choices=[], value=None, interactive=False),
                gr.update(choices=[], value=None, interactive=False),
            )
        # get all files for this branch+domain
        files_dict = ALL_TEMPLATES.get(branch, {}).get(domain, {})
        files = sorted(list(files_dict.keys()))

        # Use only the last part of the path as label
        formatted_files = [(format_name(os.path.basename(f)), f) for f in files]

        return (
            gr.update(choices=formatted_files, value=None, interactive=True),
            gr.update(choices=[], value=None, interactive=False),
        )


    def update_template_dropdown(branch, domain, filename):
        if not branch or not domain or not filename:
            return gr.update(choices=[], value=None, interactive=False)
        templates = sorted(ALL_TEMPLATES.get(branch, {}).get(domain, {}).get(filename, []))
        formatted_templates = [(format_name(t), t) for t in templates]
        return gr.update(choices=formatted_templates, value=None, interactive=True)

    branch_dropdown.change(update_domain_dropdown, branch_dropdown, [domain_dropdown, file_dropdown, template_dropdown])
    domain_dropdown.change(update_file_dropdown, [branch_dropdown, domain_dropdown], [file_dropdown, template_dropdown])
    file_dropdown.change(update_template_dropdown, [branch_dropdown, domain_dropdown, file_dropdown], template_dropdown)

    # Populate example slots
    def render_examples(branch, domain, filename, template_name):
        if not branch or not domain or not filename or not template_name:
            heading_text = "<div style='margin-top:24px; font-size: 20px; color: red;'>⚠️ Please select a branch, domain, file, and template/tool before generating examples.</div>"
            template_heading_update = gr.update(value=heading_text, visible=True)

            updates = []
            for grp, q_md, s_md, btn, llm_group, eval_md, llm_response_md in example_blocks:
                updates.extend([
                    gr.update(visible=False),
                    gr.update(value=""),
                    gr.update(value=""),
                    gr.update(visible=False),
                    gr.update(visible=False),
                    gr.update(value=""),
                    gr.update(value=""),
                ])
            return [template_heading_update] + updates

        examples = generate_examples(branch, domain, filename, template_name)
        heading_text = f"<div style='margin-top:24px; font-size: 24px;'>Generated Examples for Template: <b>{format_name(template_name)}</b></div>"
        template_heading_update = gr.update(value=heading_text, visible=True)

        updates = []
        for i, (grp, q_md, s_md, btn, llm_group, eval_md, llm_response_md) in enumerate(example_blocks):
            if i < len(examples):
                ex = examples[i]
                updates.extend([
                    gr.update(visible=True),
                    gr.update(value=f"**❓ Question:** {ex['question']}"),
                    gr.update(value=f"**✅ Solution:** {ex['solution']}"),
                    gr.update(visible=True),
                    gr.update(visible=False),
                    gr.update(value=""),
                    gr.update(value=""),
                ])
            else:
                updates.extend([
                    gr.update(visible=False),
                    gr.update(value=""),
                    gr.update(value=""),
                    gr.update(visible=False),
                    gr.update(visible=False),
                    gr.update(value=""),
                    gr.update(value=""),
                ])
        return [template_heading_update] + updates

    generate_button.click(
        render_examples,
        [branch_dropdown, domain_dropdown, file_dropdown, template_dropdown],
        [template_heading] + sum(([grp, q_md, s_md, btn, llm_group, eval_md, llm_response_md]
             for (grp, q_md, s_md, btn, llm_group, eval_md, llm_response_md) in example_blocks), [])
    )


if __name__ == "__main__":
    print("Launching Gradio app...")
    demo.launch()