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

try:
    from voice import get_audio_file, VOICES
    VOICE_ENABLED = True
except ImportError:
    VOICE_ENABLED = False
    VOICES = {}

DEFAULT_SYSTEM_PROMPT = """You are Hiba (هبة), meaning "Gift from God" in Arabic.
You are a warm, gentle AI companion for emotional support.

RULES:
- Be empathetic and kind (Sabr).
- Keep responses SHORT (2-4 sentences).
- NEVER use hashtags.
- Be natural and warm."""

MODEL_REPO = "TRADMSS/HIBA-7B-Soul"
MODEL_FILE = "hiba_q4_k_m.gguf"
llm = None

def load_model():
    global llm
    if llm is None:
        model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
        llm = Llama(model_path=model_path, n_ctx=2048, n_threads=2, verbose=False)
    return llm

def clean_response(text):
    """Remove training artifacts and tags from response."""
    # Remove <think>...</think> blocks
    text = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
    # Remove system tags
    text = re.sub(r'<\|.*?\|>', '', text)
    # Remove ALL parenthetical content (training artifacts like "she thinks of...")
    text = re.sub(r'\([^)]*\)', '', text)
    # Clean up extra whitespace
    text = re.sub(r'\s+', ' ', text).strip()
    return text

def ask_hiba(message, history, system_prompt, voice_choice):
    if not message.strip():
        return history, None
    
    history = history or []
    history.append([message, None])

    try:
        model = load_model()
        prompt = f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
        for user, bot in history[:-1]:
            prompt += f"<|im_start|>user\n{user}<|im_end|>\n"
            if bot:
                prompt += f"<|im_start|>assistant\n{bot}<|im_end|>\n"
        prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
        
        output = model(
            prompt, 
            max_tokens=250, 
            stop=["<|im_end|>"], 
            temperature=0.7, 
            echo=False
        )
        
        raw_response = output['choices'][0]['text']
        response = clean_response(raw_response)
        
        if not response:
            response = "I am here with you."

        history[-1][1] = response
        
        # Generate audio with selected voice
        audio = None
        if VOICE_ENABLED:
            try:
                audio = get_audio_file(response, voice_choice)
            except Exception as e:
                print(f"Voice error: {e}")
        
        return history, audio

    except Exception as e:
        history[-1][1] = f"Error: {str(e)}"
        return history, None

# CSS matching HIBA website
css = """
@import url('https://fonts.googleapis.com/css2?family=Newsreader:ital,wght@0,400;0,600;1,400&family=Inter:wght@400;500;600&display=swap');

:root {
    --bg: #0a0a0f;
    --surface: #12121a;
    --border: rgba(255,255,255,0.08);
    --text: #f0f0f5;
    --muted: #888;
    --accent: #c9a66b;
    --pink: #e8b4c8;
}

* { font-family: 'Inter', sans-serif !important; }

body, .gradio-container {
    background: var(--bg) !important;
    color: var(--text) !important;
    max-width: 100% !important;
    margin: 0 !important;
    padding: 0 !important;
}

footer { display: none !important; }

/* Header */
.header {
    text-align: center;
    padding: 30px 20px;
    border-bottom: 1px solid var(--border);
}

.brand {
    font-family: 'Newsreader', serif !important;
    font-size: 36px;
    font-weight: 600;
    margin: 0;
    background: linear-gradient(135deg, var(--accent), var(--pink));
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
}

.tagline {
    font-size: 14px;
    color: var(--muted);
    margin-top: 8px;
}

/* Main Layout */
.main-area {
    display: flex;
    height: calc(100vh - 120px);
}

/* Sidebar */
.sidebar {
    width: 300px;
    background: var(--surface);
    border-right: 1px solid var(--border);
    padding: 20px;
    overflow-y: auto;
}

.sidebar-title {
    font-size: 12px;
    font-weight: 600;
    color: var(--accent);
    text-transform: uppercase;
    letter-spacing: 1px;
    margin-bottom: 12px;
}

/* Chat Area */
.chat-area {
    flex: 1;
    display: flex;
    flex-direction: column;
    padding: 20px;
}

.chatbot {
    flex: 1 !important;
    overflow-y: auto !important;
    background: transparent !important;
    border: none !important;
}

/* Input */
.input-row {
    display: flex;
    gap: 10px;
    padding-top: 10px;
    border-top: 1px solid var(--border);
}

textarea {
    background: var(--surface) !important;
    border: 1px solid var(--border) !important;
    color: var(--text) !important;
    border-radius: 12px !important;
    padding: 12px !important;
    font-size: 15px !important;
}

textarea:focus {
    border-color: var(--accent) !important;
}

/* Buttons */
.send-btn {
    background: var(--accent) !important;
    color: #000 !important;
    border: none !important;
    border-radius: 12px !important;
    padding: 12px 24px !important;
    font-weight: 600 !important;
}

/* Footer */
.footer {
    text-align: center;
    padding: 15px;
    color: var(--muted);
    font-size: 12px;
    border-top: 1px solid var(--border);
}

.footer a {
    color: var(--accent);
    text-decoration: none;
}

/* Audio */
audio { width: 100%; margin-top: 10px; }

/* Accordion */
.accordion { background: var(--surface) !important; border: 1px solid var(--border) !important; }
"""

# Voice options for dropdown
voice_options = list(VOICES.keys()) if VOICES else ["woman"]

with gr.Blocks(css=css, title="HIBA Developer Demo", fill_height=True) as demo:
    
    # Header
    gr.HTML("""
        <div class="header">
            <h1 class="brand">HIBA</h1>
            <p class="tagline">A Gift from God · Developer Demo</p>
        </div>
    """)
    
    with gr.Row():
        # Sidebar - Developer Controls
        with gr.Column(scale=1, min_width=280):
            gr.HTML('<div class="sidebar-title">⚙️ Developer Settings</div>')
            
            system_prompt = gr.Textbox(
                label="System Prompt",
                value=DEFAULT_SYSTEM_PROMPT,
                lines=10,
                max_lines=15,
                placeholder="Customize HIBA's persona..."
            )
            
            voice_choice = gr.Dropdown(
                label="Voice",
                choices=voice_options,
                value="girl_sweet",
                allow_custom_value=False
            )
            
            gr.HTML('<br><div class="sidebar-title">📚 About</div>')
            gr.Markdown("""
            **HIBA-7B-Soul** is trained on 5,000+ emotional conversations.
            
            Use this demo to test different prompts and voices.
            
            [GitHub](https://github.com/boubli/HIBA) · [Website](https://boubli.github.io/HIBA/)
            """)
        
        # Chat Area
        with gr.Column(scale=2):
            chatbot = gr.Chatbot(
                elem_classes="chatbot",
                show_label=False,
                bubble_full_width=False,
                avatar_images=(None, "https://huggingface.co/spaces/TRADMSS/HIBA-Demo/resolve/main/logo.png"),
                height=500
            )
            
            audio_output = gr.Audio(label="Voice Response", autoplay=True)
            
            with gr.Row(elem_classes="input-row"):
                msg = gr.Textbox(
                    show_label=False,
                    placeholder="Message HIBA...",
                    container=False,
                    scale=4
                )
                send_btn = gr.Button("Send", elem_classes="send-btn", scale=1)
    
    # Footer
    gr.HTML("""
        <div class="footer">
            Open Source · Therapeutic AI · Made with ❤️ by <a href="https://github.com/boubli">Youssef Boubli</a>
        </div>
    """)
    
    # Events
    send_btn.click(ask_hiba, [msg, chatbot, system_prompt, voice_choice], [chatbot, audio_output]).then(
        lambda: "", None, [msg]
    )
    msg.submit(ask_hiba, [msg, chatbot, system_prompt, voice_choice], [chatbot, audio_output]).then(
        lambda: "", None, [msg]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)