# app.py — Updated for TinyLlama automatic download

import gradio as gr
import os
import tempfile
import textwrap
import multiprocessing

from datetime import datetime
from pathlib import Path
from typing import List, Dict, Any, Optional

from huggingface_hub import hf_hub_download
from llama_cpp import Llama


import os

token=os.environ["HF_TOKEN"]

from src.conversation import ConversationMemory
from src.chatbot import LocalChatbot
from src.model_loader import load_local_model


memory = ConversationMemory()     # <<< ADD THIS
llm = load_local_model()
bot = LocalChatbot(llm, memory)


# ----------------------
# MODEL DOWNLOAD (UPDATED)
# ----------------------
def load_remote_model():
    print("Downloading Phi-1.5 model (GGUF) from Hugging Face...")

    model_path = hf_hub_download(
        repo_id="tensorblock/phi-1_5-GGUF",
        filename="phi-1_5-Q4_K_M.gguf"
    )

    print(f"Model downloaded to: {model_path}")

    return Llama(
        model_path=model_path,
        n_ctx=4096,        # Phi-1.5 benefits from longer context
        n_threads=6,       # Adjust to CPU cores
        n_gpu_layers=0,    # CPU inference on HF Space
        verbose=False
    )

INTENT_TEMPLATES = {
    "math": "You are a math solver. Solve step-by-step only.",
    "code": "You are a coding expert. Provide clean, working code.",
    "civics": "Explain clearly like a Class 10 SST teacher.",
    "exam": "Prepare concise exam-focused notes and important questions."
}


def now_ts():
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")







# ----------------------
# CUSTOM CSS
# ----------------------
CUSTOM_CSS = """
/* GLOBAL */
.gradio-container { 
    background: linear-gradient(135deg, #f6f7f9 0%, #e9ecf1 100%);
    font-family: Inter, system-ui;
}
#main_card h3 {
    font-size: 28px !important;
    font-weight: 700 !important;
    background: linear-gradient(90deg, #0ea5e9, #06b6d4);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    text-align: center;
    border-bottom: 2px solid rgba(0,0,0,0.15);
}

/* MAIN CARD */
#main_card { 
    background: #ffffff;
    border: 1px solid #e3e8ef;
    border-radius: 16px;
    padding: 16px;
    box-shadow: 0 4px 16px rgba(0,0,0,0.05);
}

/* Chat UI */
.gradio-chatbot .assistant {
    background: #4f46e5 !important;
    color: white !important;
    border-radius: 14px;
    padding: 12px;
}
.gradio-chatbot .user {
    background: #f1f5f9 !important;
    border-radius: 14px;
    padding: 12px;
}


/* Input box */
#message-box textarea {
    background: #e0e7ff !important;
    border-radius: 12px !important;
    font-size: 24px !important;
}


.tool-btn {
    width: 100%;
    background: #ffffff;
    color: #374151;
    border-radius: 12px;
    border: 1px solid #e5e7eb;
    font-weight: 600;
    transition: 0.2s ease;
}

.tool-btn:hover {
    background: #f3f4f6;
    transform: scale(1.02);
}

/* Send button */
.send-btn {
    background: #4f46e5 !important;
    color: white !important;
    transition: background 0.2s ease, transform 0.2s ease;
}

.send-btn:hover {
    background: #4338ca !important;
    transform: scale(1.05);
}

/* Mic button */
.icon-btn {
    background: #f1f5f9 !important;
    transition: background 0.2s ease, transform 0.2s ease;
}

.icon-btn:hover {
    background: #e2e8f0 !important;
    transform: scale(1.05);
}
"""


# ----------------------
# Voice JS
# ----------------------
PAGE_JS = """
<script>
(function(){
  window.startVoiceRecognition = function(elem_id){
    const wrapper = document.getElementById(elem_id);
    if(!wrapper) return;
    const textarea = wrapper.querySelector('textarea');
    if(!textarea) return;

    const SR = window.SpeechRecognition || window.webkitSpeechRecognition;
    if(!SR) return alert('Speech recognition not supported');

    const recog = new SR();
    recog.lang = 'en-US';
    recog.interimResults = false;

    recog.onresult = function(e){
      textarea.value = e.results[0][0].transcript;

      textarea.style.background = "#e7f5ff";
      setTimeout(() => { textarea.style.background = ""; }, 400);
    };

    recog.start();
  };
})();
</script>
"""


# ----------------------
# EXPORT TXT/PDF  (FULLY FIXED VERSION FOR GRADIO)
# ----------------------

EXPORT_DIR = "exports"
os.makedirs(EXPORT_DIR, exist_ok=True)

def convert_history_to_pairs(history: List[Dict[str, str]]):
    """
    Converts Gradio Chatbot dict-history → list of (user, assistant) pairs.
    Gradio format:
        [{"role": "user", "content": "..."},
         {"role": "assistant", "content": "..."}]

    Output:
        [("user message", "assistant message"), ...]
    """
    pairs = []
    buffer = {}

    for msg in history:
        if msg["role"] == "user":
            buffer["user"] = msg["content"]
        elif msg["role"] == "assistant":
            buffer["assistant"] = msg["content"]

        if "user" in buffer and "assistant" in buffer:
            pairs.append((buffer["user"], buffer["assistant"]))
            buffer = {}

    return pairs

from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
from reportlab.lib.styles import getSampleStyleSheet
from reportlab.lib.pagesizes import letter


def export_chat_files(pairs):
    import os, textwrap
    from datetime import datetime
    from reportlab.lib.pagesizes import A4
    from reportlab.pdfgen import canvas

    # ensure folder
    folder = "exports"
    os.makedirs(folder, exist_ok=True)

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

    txt_path = os.path.join(folder, f"chat_{timestamp}.txt")
    pdf_path = os.path.join(folder, f"chat_{timestamp}.pdf")

    # ---- WRITE TXT ----
    with open(txt_path, "w", encoding="utf-8") as f:
        for user_msg, bot_msg in pairs:
            f.write("👤 User:\n" + str(user_msg) + "\n\n")
            f.write("🤖 Assistant:\n" + str(bot_msg) + "\n")
            f.write("-" * 60 + "\n\n")

    # ---- WRITE PDF ----
    c = canvas.Canvas(pdf_path, pagesize=A4)
    width, height = A4
    margin = 40
    text = c.beginText(margin, height - margin)
    text.setFont("Helvetica", 10)

    with open(txt_path, "r", encoding="utf-8") as fh:
        for line in fh:
            for wrapped in textwrap.wrap(line.rstrip(), 90):
                text.textLine(wrapped)

    c.drawText(text)
    c.showPage()
    c.save()

    # return {"txt": txt_path, "pdf": pdf_path}
    return {
        "txt": txt_path if txt_path is not None else None,
        "pdf": pdf_path if pdf_path is not None else None,
    }



# ----------------------
# Core chat function
# ----------------------
def generate_reply(user_msg: str, history: List[Dict[str, Any]]):
    if history is None:
        history = []

    if not user_msg.strip():
        return history

    # ---------------- Your Existing Logic ----------------
    intent = None
    low = user_msg.lower()
    for key in INTENT_TEMPLATES:
        if low.startswith(key):
            intent = key
            user_msg = user_msg[len(key):].strip()
            break

    system_prefix = INTENT_TEMPLATES.get(intent, None)
    if system_prefix:
        prompt = f"{system_prefix}\nUser: {user_msg}"
    else:
        prompt = f"User: {user_msg}"

    bot_reply = bot.ask(prompt)
    ts = now_ts()
    bot_reply_ts = f"{bot_reply}\n\n🕒 {ts}"

    history.append({"role": "user", "content": user_msg})
    history.append({"role": "assistant", "content": bot_reply_ts})

    try:
        memory.add(user_msg, bot_reply)
    except:
        pass

    return history


    
def export_handler(history):
    import time

    if history is None:
        history = []

    pairs = convert_history_to_pairs(history)
    files = export_chat_files(pairs)

    time.sleep(0.15)  # Required for HF Spaces

    return (
        gr.update(value=files["txt"], visible=True),
        gr.update(value=files["pdf"], visible=True),
    )


# ----------------------
# UI
# ----------------------
with gr.Blocks(title="Tayyab — Chatbot") as demo:

    gr.HTML(f"""
    <style>{CUSTOM_CSS}</style>
    <script>{PAGE_JS}</script>
    """)

    with gr.Row():
        with gr.Column(scale=1, min_width=220):
            gr.Markdown("### ⚡ Tools & Export")

            new_chat_btn = gr.Button("➕ New Chat", elem_classes="tool-btn")
            export_btn = gr.Button("📥 Export TXT/PDF", elem_classes="tool-btn")

            with gr.Accordion("📄 Exported Files", open=False):
                file_txt = gr.File(label="Download TXT", visible=False)
                file_pdf = gr.File(label="Download PDF", visible=False)

        with gr.Column(scale=3, elem_id="main_card"):
            gr.Markdown("<h3>Smart Learning Assistant - Tayyab</h3>")

            chatbot = gr.Chatbot(
                height=480,
                elem_id="chatbot_box",
                avatar_images=(None, None)   # works fine in Gradio 3.50.2
            )

            with gr.Row():
                msg = gr.Textbox(
                    placeholder="Type a message or use the mic",
                    elem_id="message-box",
                    show_label=False,
                    lines=3,
                )
                send_btn = gr.Button("Send", elem_classes="send-btn")
                mic_btn = gr.Button("🎤 Voice input", elem_classes="icon-btn")
                mic_btn.click(None, None, None, js='startVoiceRecognition("message-box")')

    # ----------- Actions / Logic ----------------

    send_btn.click(generate_reply, [msg, chatbot], [chatbot])
    msg.submit(generate_reply, [msg, chatbot], [chatbot])

    def new_chat():
        memory.clear()
        return []

    new_chat_btn.click(new_chat, outputs=[chatbot])

    export_btn.click(
        export_handler,
        inputs=[chatbot],
        outputs=[file_txt, file_pdf]
    )

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