"""Gradio UI for ValluvarAI."""

# Patch Jinja2 LRUCache to handle unhashable keys (gradio 4.44.0 bug).
try:
    from jinja2.utils import LRUCache

    def _make_hashable(key):
        if isinstance(key, dict):
            return tuple(sorted((k, _make_hashable(v)) for k, v in key.items()))
        if isinstance(key, list):
            return tuple(_make_hashable(i) for i in key)
        return key

    _original_getitem = LRUCache.__getitem__
    _original_get = LRUCache.get
    _original_setitem = LRUCache.__setitem__

    def _safe_getitem(self, key):
        return _original_getitem(self, _make_hashable(key))

    def _safe_get(self, key, default=None):
        return _original_get(self, _make_hashable(key), default)

    def _safe_setitem(self, key, value):
        return _original_setitem(self, _make_hashable(key), value)

    LRUCache.__getitem__ = _safe_getitem
    LRUCache.get = _safe_get
    LRUCache.__setitem__ = _safe_setitem
except Exception:
    pass

import huggingface_hub

if not hasattr(huggingface_hub, "HfFolder"):
    class _HfFolder:
        path = None

        @staticmethod
        def get_token():
            return None

    huggingface_hub.HfFolder = _HfFolder

import gradio as gr

from model_engine import model, stoi, itos
from kural_engine import generate_kural
from game_engine import new_round, check_guess

with gr.Blocks(title="ValluvarAI") as demo:
    gr.Markdown("# 🕉️ ValluvarAI")
    gr.Markdown(
        "An AI that writes new Thirukkurals in the style of Thiruvalluvar. "
        "Enter a Tamil theme to generate bilingual wisdom."
    )

    with gr.Row():
        with gr.Column(scale=1):
            pass
        with gr.Column(scale=1, min_width=150):
            gr.Image(
                "image/thiruvalluvar_stamp.jpeg",
                label="Thiruvalluvar",
                show_label=False,
                height=150,
                width=150,
            )
        with gr.Column(scale=1):
            pass

    with gr.Accordion("Model Card", open=False):
        gr.Markdown(
            f"""
            **Architecture:** GPT ({model.config.n_layer}L/{model.config.n_head}H/{model.config.n_embd}D)  
            **Parameters:** {sum(p.numel() for p in model.parameters()) / 1e6:.1f}M  
            **Vocabulary:** {len(stoi)} characters (Tamil + English)  
            **Tokenization:** Character-level  
            **Training Steps:** 10,000  
            **Device:** Apple MPS (Mac Mini)  
            **Training Time:** ~5 hours  
            **Final Loss:** ~1.5
            """
        )

    with gr.Accordion("Training Details", open=False):
        gr.Markdown(
            f"""
            **Steps:** 10,000   
            **Device:** Apple MPS (Mac Mini)  
            **Time:** ~5 hours  
            **Final Loss:** ~1.5  
            """
        )

    with gr.Tab("✨ Generate Kural"):
        with gr.Row():
            with gr.Column():
                prompt = gr.Textbox(
                    label="Theme (Tamil)",
                    placeholder="e.g., கடவுள் வாழ்த்து, நட்பு, அரசியல்",
                    value="கடவுள் வாழ்த்து",
                )
                temperature = gr.Slider(
                    minimum=0.1,
                    maximum=2.0,
                    value=0.8,
                    step=0.1,
                    label="Temperature (Creativity)",
                )
                max_tokens = gr.Slider(
                    minimum=50,
                    maximum=400,
                    value=200,
                    step=50,
                    label="Max Tokens",
                )
                generate_btn = gr.Button("Generate", variant="primary")

            with gr.Column():
                output = gr.Textbox(label="Generated Kural", lines=10)
                source = gr.Textbox(label="Source")

        generate_btn.click(
            fn=generate_kural,
            inputs=[prompt, temperature, max_tokens],
            outputs=[output, source],
        )

        gr.Markdown("### Quick Themes")
        with gr.Row():
            themes = [
                "கடவுள் வாழ்த்து",
                "வான் சிறப்பு",
                "நட்பு",
                "அரசியல்",
                "அறன் வலியுறுத்தல்",
                "கல்வி",
                "காதல்",
                "பொருள்",
                "அறம்",
                "வீரம்",
                "வாய்மை",
                "அன்பு",
            ]
            for theme in themes:
                btn = gr.Button(theme)
                btn.click(lambda t=theme: t, outputs=prompt)

    with gr.Tab("🎯 Valluvar or AI?"):
        gr.Markdown(
            "### Can you tell the difference?\n"
            "Read the Tamil couplet and guess whether it was written by "
            "Thiruvalluvar or generated by the AI."
        )

        couplet_display = gr.Textbox(
            label="குறள்",
            lines=3,
            interactive=False,
            value="Click **Next Couplet** to begin!",
        )
        game_state = gr.State(value=None)

        with gr.Row():
            valluvar_btn = gr.Button("📖 Valluvar", variant="secondary", scale=1)
            ai_btn = gr.Button("🤖 AI", variant="secondary", scale=1)

        reveal_display = gr.Markdown()
        next_btn = gr.Button("Next Couplet", variant="primary")

        def safe_new_round():
            try:
                return new_round()
            except Exception as e:
                print(f"[ERROR] Round generation failed: {e}")
                fallback = "அகர முதல் எழுத்தெல்லாம் ஆதி\nபகவன் முதற்றே உலகு"
                return fallback, None, "⚠️ Error loading round. Try again."

        next_btn.click(
            fn=safe_new_round,
            outputs=[couplet_display, game_state, reveal_display],
        )

        valluvar_btn.click(
            fn=lambda s: check_guess("valluvar", s),
            inputs=[game_state],
            outputs=[reveal_display],
        )

        ai_btn.click(
            fn=lambda s: check_guess("ai", s),
            inputs=[game_state],
            outputs=[reveal_display],
        )

    with gr.Tab("📊 About"):
        gr.Markdown(
            f"""
            ValluvarAI is a GPT model trained on the **only** Thirukkural dataset.

            ## Capabilities

            - ✅ Generate authentic Tamil couplets (2 lines × 4 words)
            - ✅ Produce coherent English translations
            - ✅ Handle traditional themes (virtue, politics, love)
            - ❌ Modern topics (science, technology) - not in training data
            - ❌ Contemporary language - only classical Tamil
            - ⚠️ May contain inaccuracies or repetitions

            ## Examples of AI vs Original

            The model sometimes generates exact memorized kurals from the 1330,
            and sometimes creates entirely new ones in Thiruvalluvar's style.

            ## About Me

            **NaveenKumar Namachivayam**

            - [QAInsights.com](https://www.qainsights.com)
            - [Dosa.Dev](https://dosa.dev)
            - [JMeter.AI](https://jmeter.ai)
            - [LinkedIn](https://www.linkedin.com/in/naveenkumarn/)

            ## Acknowledgements

            - Special thanks to [Project Madurai](https://www.projectmadurai.org/pm_etexts/utf8/pmuni0153.html) for making the Thirukkural text and translations available.

            - Built with Tamil ❤️ using PyTorch and Gradio.

            """
        )

import os

port = int(os.environ.get("PORT", "7860"))
demo.launch(server_name="0.0.0.0", server_port=port)