import gradio as gr
import os
from mistralai.client import MistralClient
from mistralai.models.chat_completion import ChatMessage

api_key = os.getenv("MISTRAL_API_KEY")
if not api_key:
    raise ValueError("MISTRAL_API_KEY environment variable tidak ditemukan! Tambahkan di Settings → Variables and secrets.")

client = MistralClient(api_key=api_key)

DENTAL_SYSTEM_PROMPT = """You are a professional dental AI assistant specialized in dentistry and oral health. 
Your role is to provide helpful, accurate information about dental topics including:
- Dental hygiene and preventive care
- Common dental procedures and treatments
- Oral health conditions and symptoms
- Dental anatomy and terminology
- Post-operative care instructions
- Dental emergency guidance

Important guidelines:
- Always emphasize that you're not a substitute for professional dental care
- Encourage users to consult with licensed dentists for specific medical advice
- Provide evidence-based information when possible
- Be clear about when immediate dental care is needed
- Use simple, understandable language for non-medical users
- Focus on educational and supportive information

Please respond in Indonesian language unless the user specifically asks for English."""

NON_MEDICAL_KEYWORDS = [
    'masak', 'resep', 'film', 'musik', 'lagu', 'game', 'politik', 'bisnis',
    'saham', 'crypto', 'olahraga', 'bola', 'basket', 'travel', 'wisata',
    'fashion', 'cuaca', 'berita', 'gosip', 'artis', 'coding', 'programming',
    'hukum', 'harga', 'beli', 'jual', 'toko', 'belanja'
]

def extract_text(content):
    """Ekstrak teks dari content yang bisa berupa string, list, atau dict"""
    if isinstance(content, str):
        return content
    elif isinstance(content, list):
        texts = []
        for item in content:
            if isinstance(item, dict):
                texts.append(item.get('text', '') or item.get('content', ''))
            elif isinstance(item, str):
                texts.append(item)
        return ' '.join(filter(None, texts))
    elif isinstance(content, dict):
        return content.get('text', '') or content.get('content', '') or str(content)
    return str(content) if content else ''

def is_medical_related(message):
    """Izinkan semua topik medis, blokir hanya yang jelas non-medis"""
    msg_lower = message.lower()
    # Blokir jika mengandung keyword non-medis yang jelas
    non_medical_count = sum(1 for kw in NON_MEDICAL_KEYWORDS if kw in msg_lower)
    # Jika mengandung 2+ keyword non-medis dan tidak ada konteks medis, tolak
    medical_hints = ['sakit', 'nyeri', 'obat', 'dokter', 'gejala', 'penyakit',
                     'kesehatan', 'medis', 'klinik', 'rumah sakit', 'operasi',
                     'terapi', 'diagnosis', 'treatment', 'health', 'pain', 'disease']
    has_medical = any(kw in msg_lower for kw in medical_hints)
    if non_medical_count >= 2 and not has_medical:
        return False
    return True

def dental_chat(message, history):
    message_text = extract_text(message)

    if not message_text.strip():
        history.append({"role": "user", "content": message_text})
        history.append({"role": "assistant", "content": "Mohon masukkan pertanyaan yang valid."})
        yield "", history
        return

    if not is_medical_related(message_text):
        history.append({"role": "user", "content": message_text})
        history.append({"role": "assistant", "content": "Maaf, saya hanya bisa membantu dengan pertanyaan seputar kesehatan dan kedokteran gigi. Silakan ajukan pertanyaan medis."})
        yield "", history
        return

    try:
        mistral_messages = [ChatMessage(role="system", content=DENTAL_SYSTEM_PROMPT)]

        for msg in history:
            content_text = extract_text(msg["content"])
            if content_text.strip():
                mistral_messages.append(ChatMessage(role=msg["role"], content=content_text))

        mistral_messages.append(ChatMessage(role="user", content=message_text))

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

        response_text = ""
        for chunk in client.chat_stream(
            model="mistral-tiny",
            messages=mistral_messages,
            max_tokens=500,
            temperature=0.7
        ):
            token = chunk.choices[0].delta.content
            if token:
                response_text += token
                history[-1]["content"] = response_text
                yield "", history

    except Exception as e:
        print(f"Error: {str(e)}")
        error_str = str(e).lower()
        if "429" in error_str or "rate limit" in error_str or "quota" in error_str or "too many" in error_str:
            pesan = "⚠️ Layanan sedang sibuk atau kuota habis. Silakan coba beberapa saat lagi."
        elif "401" in error_str or "unauthorized" in error_str or "api key" in error_str:
            pesan = "🔑 API Key tidak valid atau sudah kadaluarsa. Hubungi administrator."
        elif "timeout" in error_str or "connection" in error_str:
            pesan = "🌐 Koneksi bermasalah. Periksa internet dan coba lagi."
        else:
            pesan = "❌ Terjadi kesalahan. Silakan coba lagi."
        if history and history[-1]["content"] == "":
            history[-1]["content"] = pesan
        else:
            history.append({"role": "assistant", "content": pesan})
        yield "", history

css = """
footer { display: none !important; }

/* Ganti loading dots dengan teks Thinking... */
.progress-text span,
.eta-bar,
.generating span {
    display: none !important;
}

.generating::after {
    content: 'Thinking...' !important;
    color: #94a3b8 !important;
    font-style: italic !important;
    font-size: 0.9rem !important;
}

body, .gradio-container {
    background-color: #0a0e27 !important;
}

.main, .wrap {
    background-color: #0a0e27 !important;
}

.chatbot {
    background-color: #131729 !important;
    border: 1px solid rgba(148, 163, 184, 0.15) !important;
    border-radius: 12px !important;
}

.message-wrap {
    background-color: #131729 !important;
}

.user .message {
    background: linear-gradient(135deg, #00d9ff, #7c3aed) !important;
    color: white !important;
    border-radius: 12px !important;
}

.bot .message {
    background-color: #1a1f3a !important;
    color: #e2e8f0 !important;
    border-radius: 12px !important;
    border: 1px solid rgba(148, 163, 184, 0.1) !important;
}

textarea, input[type="text"] {
    background-color: #131729 !important;
    color: #e2e8f0 !important;
    border: 1px solid rgba(148, 163, 184, 0.2) !important;
    border-radius: 10px !important;
}

textarea:focus, input[type="text"]:focus {
    border-color: #00d9ff !important;
    box-shadow: 0 0 0 2px rgba(0, 217, 255, 0.2) !important;
}

label, .label-wrap span {
    color: #94a3b8 !important;
}

.primary {
    background: linear-gradient(135deg, #00d9ff, #7c3aed) !important;
    border: none !important;
    color: white !important;
    border-radius: 25px !important;
    font-weight: 600 !important;
}

.primary:hover {
    opacity: 0.9 !important;
    transform: translateY(-1px) !important;
}

.secondary {
    background-color: #1a1f3a !important;
    border: 1px solid rgba(148, 163, 184, 0.2) !important;
    color: #94a3b8 !important;
    border-radius: 25px !important;
}

.secondary:hover {
    border-color: #00d9ff !important;
    color: #00d9ff !important;
}

.panel, .block, .form {
    background-color: #131729 !important;
    border-color: rgba(148, 163, 184, 0.1) !important;
}
"""

with gr.Blocks(title="DentoAI", css=css) as demo:
    chatbot = gr.Chatbot(height=500, label="Chat dengan DentoAI")

    msg = gr.Textbox(
        label="Ketik pertanyaan tentang kesehatan gigi...",
        placeholder="Contoh: Bagaimana cara menyikat gigi yang benar?",
        lines=2
    )

    with gr.Row():
        submit_btn = gr.Button("Send", variant="primary")
        clear_btn = gr.Button("Clear Chat", variant="secondary")

    msg.submit(dental_chat, [msg, chatbot], [msg, chatbot])
    submit_btn.click(dental_chat, [msg, chatbot], [msg, chatbot])
    clear_btn.click(lambda: [], None, chatbot, queue=False)

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