"""
Ammu-1.0 FactCheck Classifier - Gradio Demo for Hugging Face Spaces
Author: Kalyan KS (https://huggingface.co/kalyan-ks)
"""

import os
import pandas as pd
import gradio as gr
import torch
from transformers import pipeline

# Support Hugging Face ZeroGPU Spaces with local fallback
try:
    import spaces
except ImportError:
    class spaces:
        @staticmethod
        def GPU(func):
            return func

############################################################
# Model Catalog & Pipeline Cache
############################################################

MODEL_OPTIONS = {
    "Ammu-1.0-Base (150M)": "kalyan-ks/Ammu-1.0-FactCheck-Need-base",
    "Ammu-1.0-Small (68M)": "kalyan-ks/Ammu-1.0-FactCheck-Need-small",
    "Ammu-1.0-Tiny (33M)": "kalyan-ks/Ammu-1.0-FactCheck-Need-tiny",
    "Ammu-1.0-Nano (17M)": "kalyan-ks/Ammu-1.0-FactCheck-Need-nano",
}

MODEL_URLS = {
    "Ammu-1.0-Base (150M)": "https://huggingface.co/kalyan-ks/Ammu-1.0-FactCheck-Need-base",
    "Ammu-1.0-Small (68M)": "https://huggingface.co/kalyan-ks/Ammu-1.0-FactCheck-Need-small",
    "Ammu-1.0-Tiny (33M)": "https://huggingface.co/kalyan-ks/Ammu-1.0-FactCheck-Need-tiny",
    "Ammu-1.0-Nano (17M)": "https://huggingface.co/kalyan-ks/Ammu-1.0-FactCheck-Need-nano",
}

MODEL_PERFORMANCE = {
    "Ammu-1.0-Base (150M)": "96.56% Macro-F1",
    "Ammu-1.0-Small (68M)": "96.13% Macro-F1",
    "Ammu-1.0-Tiny (33M)": "96.35% Macro-F1",
    "Ammu-1.0-Nano (17M)": "95.66% Macro-F1",
}

# In-memory pipeline cache to prevent reloading overhead
MODELS = {}

def get_factcheck_pipeline(model_name):
    if model_name not in MODELS:
        device = 0 if torch.cuda.is_available() else -1
        MODELS[model_name] = pipeline(
            task="text-classification",
            model=model_name,
            device=device,
            top_k=None  # Returns probabilities for all classes
        )
    return MODELS[model_name]


############################################################
# Benchmark Examples
############################################################

EXAMPLES = {
    "Example 1 (Factual Question)": "India is the most populous country in the world. Do you agree?",
    "Example 2 (Complex Age Calculation)": "Assume that I was born 16 years 4 months 3 days exactly after Sachin Tendulkar was born. How old I am today?",
    "Example 3 (Short Age Query)": "I was born 16 years 4 months 3 days exactly after Sachin Tendulkar was born. My age?",
    "Example 4 (Input-Text Entity Extraction)": "Kalyan KS is from India. Extract named entities from the given sentence.",
    "Example 5 (Task Instruction with Input)": "Extract named entities: Kalyan KS is from India.",
}


############################################################
# UI & Helper Functions
############################################################

def update_model_info(selected_model_key):
    url = MODEL_URLS[selected_model_key]
    model_id = MODEL_OPTIONS[selected_model_key]
    perf = MODEL_PERFORMANCE[selected_model_key]
    
    info = f"""
- **Model:** [{model_id}]({url})
- **Architecture:** ModernBERT Encoder
- **Parameters:** {selected_model_key.split('(')[-1].replace(')', '')}
- **Task:** Fact-Check Need Classification (Single-Label)
- **Classes:** `FACT_CHECK_NEEDED`, `NO_FACT_CHECK_NEEDED`
- **Performance:** {perf}
"""
    footer = f"""
### 🔗 Resources & Links
- **Hugging Face Model:** [{model_id}]({url})
- **Creator:** [Kalyan KS](https://huggingface.co/kalyan-ks)
- **License:** Apache 2.0

---
Built with ❤️ using **Gradio** and **Hugging Face Transformers**.
"""
    return info, footer


def update_stats(text):
    if not text:
        return "Characters: 0", "Approx. Tokens: 0"
    char_count = len(text)
    approx_tokens = int(len(text.split()) * 1.3)
    return f"Characters: {char_count}", f"Approx. Tokens: {approx_tokens}"


def load_example_text(example_key):
    return EXAMPLES.get(example_key, "")


def clear_inputs():
    empty_df = pd.DataFrame(columns=["Class Label", "Probability", "Confidence Bar"])
    return (
        "",                          # input_text
        "Characters: 0",             # char_caption
        "Approx. Tokens: 0",         # token_caption
        "",                          # status_output
        "",                          # result_card_html
        empty_df,                    # prob_table
        "",                          # metric_prediction
        "",                          # metric_confidence
        "",                          # metric_routing
    )


@spaces.GPU
def process_classification(input_text, selected_model_key):
    if not input_text or not input_text.strip():
        empty_df = pd.DataFrame(columns=["Class Label", "Probability", "Confidence Bar"])
        return (
            "⚠️ Please enter a prompt to classify.",
            "",
            empty_df,
            "-",
            "-",
            "-",
        )

    model_id = MODEL_OPTIONS[selected_model_key]
    classifier = get_factcheck_pipeline(model_id)

    # Run inference
    raw_results = classifier(input_text.strip())

    if isinstance(raw_results, list) and len(raw_results) > 0 and isinstance(raw_results[0], list):
        scores_list = raw_results[0]
    elif isinstance(raw_results, list):
        scores_list = raw_results
    else:
        scores_list = [raw_results]

    scores_dict = {item["label"]: float(item["score"]) for item in scores_list}
    
    top_label = max(scores_dict, key=scores_dict.get)
    top_score = scores_dict[top_label]

    # Styling based on prediction
    if top_label == "FACT_CHECK_NEEDED":
        badge_bg = "#FFEBEE"
        badge_border = "#EF5350"
        badge_color = "#C62828"
        icon = "⚠️"
        action_title = "FACT CHECK NEEDED"
        routing_desc = "Route to Retrieval (RAG) / Search / Knowledge Base before generating response."
        routing_short = "Trigger RAG / Web Search"
    else:
        badge_bg = "#E8F5E9"
        badge_border = "#66BB6A"
        badge_color = "#2E7D32"
        icon = "✅"
        action_title = "NO FACT CHECK NEEDED"
        routing_desc = "Safe for direct LLM generation (Creative writing, code, math logic, or input-grounded task)."
        routing_short = "Direct LLM Generation"

    # HTML Result Card
    styled_html = f"""
    <div style="
        border: 2px solid {badge_border};
        background: {badge_bg};
        border-radius: 12px;
        padding: 20px;
        margin-top: 10px;
        box-shadow: 0 2px 8px rgba(0,0,0,0.05);">
        <div style="display: flex; align-items: center; justify-content: space-between; flex-wrap: wrap; gap: 10px;">
            <div style="display: flex; align-items: center; gap: 12px;">
                <span style="font-size: 28px;">{icon}</span>
                <div>
                    <h3 style="margin: 0; color: {badge_color}; font-size: 20px; font-weight: 700; letter-spacing: 0.5px;">
                        {action_title}
                    </h3>
                    <p style="margin: 4px 0 0 0; color: #424242; font-size: 14px;">
                        {routing_desc}
                    </p>
                </div>
            </div>
            <div style="text-align: right; min-width: 120px;">
                <span style="font-size: 26px; font-weight: 800; color: {badge_color};">
                    {top_score * 100:.2f}%
                </span>
                <div style="font-size: 11px; text-transform: uppercase; color: #616161; font-weight: 600;">
                    Confidence
                </div>
            </div>
        </div>
    </div>
    """

    # Probability Table
    prob_rows = []
    for label, score in sorted(scores_dict.items(), key=lambda x: x[1], reverse=True):
        bar_len = int(score * 20)
        prob_rows.append({
            "Class Label": label,
            "Probability (%)": f"{score * 100:.2f}%",
            "Confidence Bar": "█" * bar_len + "░" * (20 - bar_len),
        })
    prob_df = pd.DataFrame(prob_rows)

    status_msg = f"✅ Classification completed using **{selected_model_key}**."

    return (
        status_msg,
        styled_html,
        prob_df,
        top_label,
        f"{top_score * 100:.2f}%",
        routing_short,
    )


############################################################
# Gradio Interface
############################################################

css = """
#root, body, html, .gradio-container {
    overflow-y: auto !important;
    height: 100vh !important;
}
textarea { font-size: 15px !important; }
"""

with gr.Blocks(title="🔎 Fact-Check Classifier Demo", css=css) as demo:

    gr.Markdown("# 🔎 Fact-Check Need Classification using Ammu-1.0")
    gr.Markdown(
        "Determine whether a prompt requires factual verification (**RAG / Web Search**) or can be processed directly "
        "(**Creative Writing, Reasoning, Code, or Input-Grounded Tasks**) using the fine-tuned **Ammu-1.0 FactCheck** encoder models."
    )

    with gr.Sidebar():
        gr.Markdown("## ⚙️ Settings")
        selected_model = gr.Dropdown(
            choices=list(MODEL_OPTIONS.keys()),
            value="Ammu-1.0-Base (150M)",
            label="Select Model"
        )

        gr.Markdown("---")
        gr.Markdown("### 📌 Model Information")
        model_info_md = gr.Markdown()

    with gr.Column():
        gr.Markdown("### 📝 Input Prompt")

        with gr.Row():
            example_selector = gr.Dropdown(
                choices=list(EXAMPLES.keys()),
                value="Example 1 (Factual Question)",
                label="Choose Example"
            )
            load_example_btn = gr.Button("📄 Load Example")

        input_text = gr.Textbox(
            value=EXAMPLES["Example 1 (Factual Question)"],
            lines=5,
            placeholder="Type or paste a user prompt to classify...",
            label="Enter Prompt"
        )

        with gr.Row():
            char_caption = gr.Markdown("Characters: 0")
            token_caption = gr.Markdown("Approx. Tokens: 0")

        with gr.Row():
            classify_btn = gr.Button("🔍 Classify Prompt", variant="primary")
            clear_btn = gr.Button("🗑️ Clear")

        status_output = gr.Markdown()

    with gr.Column():
        gr.Markdown("---")
        gr.Markdown("### 🎯 Classification Result")
        result_card_html = gr.HTML()

        with gr.Row():
            metric_prediction = gr.Textbox(label="Predicted Class", interactive=False)
            metric_confidence = gr.Textbox(label="Confidence", interactive=False)
            metric_routing = gr.Textbox(label="Routing Recommendation", interactive=False)

        gr.Markdown("### 📊 Class Probability Breakdown")
        prob_table = gr.Dataframe(interactive=False)

    gr.Markdown("---")
    footer_md = gr.Markdown()

    # Event Connections
    demo.load(
        fn=update_model_info,
        inputs=[selected_model],
        outputs=[model_info_md, footer_md]
    )

    demo.load(
        fn=update_stats,
        inputs=[input_text],
        outputs=[char_caption, token_caption]
    )

    selected_model.change(
        fn=update_model_info,
        inputs=[selected_model],
        outputs=[model_info_md, footer_md]
    )

    input_text.change(
        fn=update_stats,
        inputs=[input_text],
        outputs=[char_caption, token_caption]
    )

    load_example_btn.click(
        fn=load_example_text,
        inputs=[example_selector],
        outputs=[input_text]
    )

    clear_btn.click(
        fn=clear_inputs,
        outputs=[
            input_text,
            char_caption,
            token_caption,
            status_output,
            result_card_html,
            prob_table,
            metric_prediction,
            metric_confidence,
            metric_routing,
        ]
    )

    classify_btn.click(
        fn=process_classification,
        inputs=[input_text, selected_model],
        outputs=[
            status_output,
            result_card_html,
            prob_table,
            metric_prediction,
            metric_confidence,
            metric_routing,
        ]
    )

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