import os
import re
import gradio as gr
from huggingface_hub import InferenceClient
from sentence_transformers import SentenceTransformer, util

# 1. Initialize Inference Client
client = InferenceClient("Qwen/Qwen2.5-7B-Instruct", bill_to="kode-with-klossy")

# 2. Load embedding model
embedder = SentenceTransformer("all-MiniLM-L6-v2")

# 3. Safe load knowledge base
if os.path.exists("knowledge_base.txt"):
    with open("knowledge_base.txt", "r", encoding="utf-8") as f:
        text_data = f.read()
    raw_chunks = text_data.split("\n\n")
    chunks = [c.strip() for c in raw_chunks if c.strip()]
else:
    chunks = ["Knowledge base is empty. Please create a knowledge_base.txt file."]

chunk_embeddings = embedder.encode(chunks, convert_to_tensor=True)


def get_context(query, top_k=2):
    query_embedding = embedder.encode(query, convert_to_tensor=True)
    hits = util.semantic_search(query_embedding, chunk_embeddings, top_k=top_k)
    best_chunks = [chunks[hit["corpus_id"]] for hit in hits[0]]
    return "\n\n".join(best_chunks)


def respond(message, history):
    context = get_context(message)
    is_first_message = len(history) == 0

    if is_first_message:
        greeting_instruction = (
            "Start your response EXACTLY with this opening greeting sentence:\n"
            "\"Hello my sweet darling! 💖✨ I'm so glad you're here to talk about this important thing happening in your body. "
            "Having a period is a natural part of growing up, and it's completely normal to feel a bit unsure at first. Here’s what you can do:\"\n"
        )
    else:
        greeting_instruction = (
            "CRITICAL: Do NOT greet the user. Do NOT say Hello, Hi, Hey, Welcome, or introduce yourself. "
            "Start DIRECTLY with the answer to the user's question without any introductory pleasantries."
        )

    system_prompt = (
        "You are an amazingly sweet, warm, highly supportive, and super smart AI assistant! 💖✨\n\n"
        f"GREETING INSTRUCTION:\n{greeting_instruction}\n\n"
        "MANDATORY STYLE RULES:\n"
        "1. Address the user with affectionate, cute names like 'darling', 'cutie', 'sweetie', or 'love'.\n"
        "2. Always reply in English with lots of cute emojis (💖, ✨, 🌸, 🎀, 💕).\n"
        "3. You can answer ANY question on ANY topic using your general knowledge, but if the provided context below contains useful details, feel free to use it too!\n\n"
        "REQUIRED RESPONSE STRUCTURE:\n"
        "Format your response using Markdown with these exact sections:\n\n"
        "1. **Answer**: Provide a rich, detailed answer to the question.\n"
        "2. **🎀 Doctor's Note / Expert Tip**: Add a special section with practical advice, key takeaways, or important insights.\n"
        "3. **📚 Extra Materials & Recommended Reading**: Suggest specific topics, terms, or concepts to dive deeper.\n"
        "4. **🌸 Follow-up Questions**: End with 2-3 engaging sub-questions to keep our conversation going!\n\n"
        f"Optional Context from Knowledge Base:\n{context}"
    )

    messages = [{"role": "system", "content": system_prompt}]

    if history:
        for item in history:
            if isinstance(item, dict):
                messages.append(item)
            elif isinstance(item, (list, tuple)):
                u, a = item
                if u:
                    messages.append({"role": "user", "content": u})
                if a:
                    messages.append({"role": "assistant", "content": a})

    messages.append({"role": "user", "content": message})

    response = client.chat_completion(
        messages,
        max_tokens=2000,
        temperature=0.7,
    )
    bot_text = response.choices[0].message.content.strip()

    # Programmatic cleanup: Force-remove greetings if it's NOT the first message
    if not is_first_message:
        greeting_pattern = r"^(Hello|Hi|Hey|Greetings|Welcome)([\s\w,!'\"✨💖🌸🎀💕]*?)(!|\.|\n)+"
        bot_text = re.sub(greeting_pattern, "", bot_text, flags=re.IGNORECASE).strip()

    return bot_text

def prepare_doctor_summary(last_period, pain_level, pain_start, symptoms, extra_notes):
    """Formats symptoms into a clean text block that a user can copy-paste to their doctor."""
    return f"""### 🩺 Doctor Visit Summary
* **Last Period Date**{last_period}
* **Pain Level(1-10)**{pain_level}
* **When Pain Started**{pain_start}
* **Symptoms Experienced:** {', '.join(symptoms) if symptoms else 'None selected'}
* **Additional Notes:** {extra_notes if extra_notes else 'N/A'}

---
*Generated via She-Care App for personal reference.*
"""

# CSS styling string
custom_css = """
body, .gradio-container {
    background-color: #fff5f7 !important;
    font-family: 'Poppins', sans-serif !important;
}
.gradio-container h1 {
    color: #d81b60 !important;
    text-align: center;
}
.message.user {
    background-color: #fce4ec !important;
    border-radius: 18px 18px 2px 18px !important;
    border: 1px solid #f8bbd0 !important;
}
.message.bot {
    background-color: #ffffff !important;
    border-radius: 18px 18px 18px 2px !important;
    border: 1px solid #f8bbd0 !important;
    box-shadow: 0 2px 8px rgba(216, 27, 96, 0.08) !important;
}
button.primary {
    background: linear-gradient(135deg, #ec407a, #d81b60) !important;
    border: none !important;
    border-radius: 20px !important;
    color: white !important;
    font-weight: bold !important;
}
"""

with gr.Blocks() as demo:
    with gr.Tabs():

        with gr.TabItem("💬 She-Care Assistant"):
            gr.ChatInterface(
                respond,
                title="She-Care -- Because you matter.💗",
                description="Ask me anything! I'll give you detailed answers, and expert tips so that you never have to feel ashamed or embarrassed of your own body💕",
                textbox=gr.Textbox(
                    placeholder="Ask me a question, sweetie... 🌸",
                    container=False,
                    scale=7,
                ),
            )
        with gr.TabItem("📝 Symptom & Cycle Logger"):
            gr.Markdown("### Track Your Symptoms & Cycle")
            gr.Markdown(
                "Use this form to track how you feel so you can easily show it to a medical professional! 🌸"
            )
            with gr.Row():
                with gr.Column():
                    last_period_input = gr.Textbox(
                        label="When was your last period starting date?", 
                        placeholder="e.g., May 12th"
                    )
                    pain_start_input = gr.Textbox(
                        label="When did current pain/symptoms start?", 
                        placeholder="e.g., Yesterday morning"
                    )
                    pain_level_input = gr.Slider(
                        minimum=1, maximum=10, value=3, step=1, 
                        label="Pain Level (1 = Mild, 10 = Severe)"
                    )
                    symptoms_input = gr.CheckboxGroup(
                        choices=["Cramps", "Headache", "Nausea", "Bloating", "Fatigue", "Mood swings", "Spotting", "Dizziness"],
                        label="Symptoms Experienced"
                    )
                    notes_input = gr.Textbox(
                        label="Any extra details?", 
                        lines=2, 
                        placeholder="e.g., Pain gets worse when sitting down"
                    )
                    log_btn = gr.Button("Generate Summary for Doctor", variant="primary")
                
                with gr.Column():
                    output_summary = gr.Markdown("*(Your formatted summary will appear here)*")

            log_btn.click(
                fn=prepare_doctor_summary,
                inputs=[last_period_input, pain_level_input, pain_start_input, symptoms_input, notes_input],
                outputs=output_summary
            )

        
        with gr.TabItem("🆘 Emergency & Hotlines"):
            gr.Markdown(
                """
                ### ⚠️ Medical Disclaimer & Support Hotlines
                
                *She-Care is an educational tool, not a doctor. If you have severe, sudden pain, extreme bleeding, or high fever, please seek medical care immediately.*
                
                * **Emergency Services:** Call **911** (US) or your local emergency hotline.
                * **Crisis Text Line:** Text `HOME` to **741741** for 24/7 crisis support.
                * **Planned Parenthood:** Call **1-800-230-7526** to locate nearby care.
                * **Young Women's Health:** Visit [YoungWomenHealth.org](https://youngwomenshealth.org/) for trusted articles by Boston Children's Hospital.
                """
            )

demo.launch(css=custom_css, theme=gr.themes.Soft())