import huggingface_hub
if not hasattr(huggingface_hub, 'HfFolder'):
    class HfFolder:
        @staticmethod
        def path(): return None
        @staticmethod
        def get_token(): return None
    huggingface_hub.HfFolder = HfFolder

import gradio_client.utils as _gcu
_json_schema_to_python_type_orig = _gcu._json_schema_to_python_type
def _json_schema_to_python_type_patched(schema, defs):
    if not isinstance(schema, dict):
        return str(schema)
    return _json_schema_to_python_type_orig(schema, defs)
_gcu._json_schema_to_python_type = _json_schema_to_python_type_patched

import gradio as gr
from llama_cpp import Llama
from huggingface_hub import hf_hub_download

MODEL_REPO = "eulogik/Bharat-Tiny-LLM-GGUF"
MODEL_FILE = "bharat-tiny-llm-q4_k_m.gguf"
IM_END = 151645

DESCRIPTION = """
<div style="text-align: center; max-width: 600px; margin: 0 auto;">
  <h1 style="font-size: 2.5rem; margin-bottom: 0.25rem;">🇮🇳 Bharat-Tiny-LLM</h1>
  <p style="font-size: 1.05rem; color: var(--body-text-color-subdued); margin-top: 0;">
    Hinglish AI — <strong>trained on a Mac Mini M4</strong>, runs on ₹8,000 phones
  </p>
  <p style="font-size: 0.85rem; margin-top: 0.5rem;">
    Built by <a href="https://eulogik.com" target="_blank"><strong>eulogik</strong></a>
    · <a href="https://huggingface.co/eulogik/Bharat-Tiny-LLM" target="_blank">Download Model</a>
    · <a href="https://github.com/eulogik/Bharat-Tiny-LLM" target="_blank">GitHub</a>
    · <a href="https://pypi.org/project/bharat-tiny-llm/" target="_blank">PyPI</a>
  </p>
</div>
"""

EXAMPLES = [
    ["Chai peete hain?", 0.3, 128],
    ["Kal interview hai, nervous ho raha hoon", 0.3, 128],
    ["Biryani kaise banate hain?", 0.3, 128],
    ["Delhi me rehne ke liye kya karna padega?", 0.3, 128],
    ["Mujhe Hindi seekhni hai", 0.3, 128],
    ["Weekend pe kya karein?", 0.3, 128],
]

CUSTOM_CSS = """
#title { text-align: center; margin-top: 2rem; }
#title h1 { font-weight: 700; letter-spacing: -0.02em; }
#chatbot { min-height: 400px; }
#chatbot .message { border-radius: 18px !important; }
#chatbot .user { background: var(--primary-400) !important; }
.footer { text-align: center; font-size: 0.8rem; color: var(--body-text-color-subdued); padding: 1rem 0; }
.footer a { color: var(--primary-500); text-decoration: none; }
.footer a:hover { text-decoration: underline; }
.examples-section { margin: 1rem 0; }
"""

print("Loading GGUF model...")
gguf_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
llm = Llama(model_path=gguf_path, n_ctx=1024, n_gpu_layers=0, verbose=False)
print("Model loaded!")


def generate(message, history, temperature, max_tokens):
    messages = [{"role": "user", "content": message}]
    if history:
        messages = [m for m in history] + messages
    out = llm.create_chat_completion(
        messages=messages,
        temperature=temperature,
        top_p=0.85,
        repeat_penalty=1.25,
        max_tokens=max_tokens,
    )
    return out["choices"][0]["message"]["content"]


def format_chat(message, history):
    history = history or []
    history.append({"role": "user", "content": message})
    yield history, ""
    response = generate(message, history, temperature_slider.value, max_tokens_slider.value)
    history.append({"role": "assistant", "content": response})
    yield history, ""


with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Soft(primary_hue="blue", neutral_hue="gray")) as demo:
    gr.HTML(DESCRIPTION)

    chatbot = gr.Chatbot(
        type="messages",
        label="Chat",
        elem_id="chatbot",
        height=450,
        bubble_full_width=False,
    )

    with gr.Row():
        msg = gr.Textbox(
            placeholder="Hinglish mein kuch bhi poochhein...",
            show_label=False,
            container=False,
            scale=8,
        )
        send = gr.Button("Send", scale=1, variant="primary")
        clear = gr.ClearButton([msg, chatbot], scale=1)

    with gr.Accordion("Settings", open=False):
        with gr.Row():
            temperature_slider = gr.Slider(0.1, 1.0, value=0.3, step=0.05, label="Temperature")
            max_tokens_slider = gr.Slider(8, 256, value=128, step=8, label="Max Tokens")

    gr.Examples(
        examples=EXAMPLES,
        inputs=[msg, temperature_slider, max_tokens_slider],
        label="Try these examples",
    )

    gr.HTML(
        '<div class="footer">'
        'Built with ❤️ by <a href="https://eulogik.com" target="_blank"><strong>eulogik</strong></a>'
        ' · Model: <a href="https://huggingface.co/eulogik/Bharat-Tiny-LLM" target="_blank">Bharat-Tiny-LLM</a>'
        ' · <a href="https://github.com/eulogik/Bharat-Tiny-LLM" target="_blank">Open Source (Apache 2.0)</a>'
        '</div>'
    )

    msg.submit(format_chat, [msg, chatbot], [chatbot, msg])
    send.click(format_chat, [msg, chatbot], [chatbot, msg])

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