import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
import torch
import time

# =======================================================
# Load Model
# =======================================================
model_name = "augtoma/qCammel-13"
print("Loading tokenizer and model...")

tokenizer = AutoTokenizer.from_pretrained(model_name)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    torch_dtype=torch.float16,
    trust_remote_code=True,
    low_cpu_mem_usage=True
)
model.eval()

print("✅ Model loaded successfully!")
print(f"Device map: {model.hf_device_map}")
print(f"Model device: {next(model.parameters()).device}")


# =======================================================
# Generate Doctor Response (Refined for natural tone)
# =======================================================
def generate_doctor_response(history):
    user_message = history[-1]["content"]

    if not user_message.strip():
        history.append({"role": "assistant", "content": "⚠️ Please describe your symptoms or ask a question."})
        yield history
        return

    # 🩺 Refined, Doctor-Like Prompt
    prompt = f"""
You are Dr. Aiden, a compassionate, calm, and experienced medical doctor. 
You speak naturally, like in a real consultation, providing medical reasoning and empathy. 
You should:
- Greet the patient kindly and acknowledge their concern.
- Offer a likely cause in simple medical terms.
- Suggest possible medicines (with safe dosage and common over-the-counter names).
- Recommend home remedies, foods, and hydration advice.
- Share short lifestyle or rest tips to aid recovery.
- End with reassurance and a disclaimer.

Keep your tone friendly yet professional — like an experienced doctor talking directly to the patient.
Avoid using headings, bullet points, or medical jargon unless necessary.
Keep your response under 180 words.

Patient says: "{user_message}"

Dr. Aiden:
"""

    # Tokenize input
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048).to(model.device)

    gen_config = GenerationConfig(
        temperature=0.7,
        top_p=0.9,
        do_sample=True,
        max_new_tokens=600,
        pad_token_id=tokenizer.pad_token_id,
        eos_token_id=tokenizer.eos_token_id,
        repetition_penalty=1.15,
    )

    input_len = inputs["input_ids"].shape[1]
    with torch.no_grad():
        output_ids = model.generate(**inputs, generation_config=gen_config)

    generated_ids = output_ids[0][input_len:]
    response = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()

    # Clean up the response
    response = clean_medical_response(response)

    # Stream response (simulated)
    history.append({"role": "assistant", "content": ""})
    for i in range(0, len(response), 5):
        history[-1]["content"] = response[:i + 5] + "▌"
        yield history.copy()
        time.sleep(0.01)

    history[-1]["content"] = response
    yield history


# =======================================================
# Clean the response
# =======================================================
def clean_medical_response(response: str) -> str:
    remove_prefixes = ["assistant:", "doctor:", "dr. aiden:", "response:", "patient:"]
    for p in remove_prefixes:
        if response.lower().startswith(p):
            response = response[len(p):].strip()

    response = response.replace("Dr. Aiden:", "").strip()

    # Ensure punctuation
    if response and response[-1] not in ".!?":
        response += "."

    # Add disclaimer if missing
    if "⚕️" not in response and "consult" not in response.lower():
        response += "\n\n⚕️ *Please note: This is AI-generated medical guidance, not a substitute for a licensed healthcare provider. Always consult a doctor for personal medical care.*"

    return response.strip()


# =======================================================
# Gradio UI
# =======================================================
with gr.Blocks(theme=gr.themes.Soft(), css="""
    .medical-header {
        background: linear-gradient(135deg, #2c3e50 0%, #3498db 100%);
        padding: 20px;
        border-radius: 12px;
        color: white;
        text-align: center;
        margin-bottom: 20px;
        box-shadow: 0 4px 12px rgba(0,0,0,0.15);
    }
""") as demo:
    
    gr.HTML("""
        <div class="medical-header">
            <h1>🏥 Dr. Aiden – AI Medical Consultation</h1>
            <p>Friendly • Professional • Science-Backed Guidance</p>
        </div>
    """)

    chatbot = gr.Chatbot(
        label="💬 Your Consultation with Dr. Aiden",
        type='messages',
        avatar_images=(
            "https://cdn-icons-png.flaticon.com/512/706/706830.png",  # Patient
            "https://cdn-icons-png.flaticon.com/512/3774/3774299.png"   # Doctor
        ),
        height=550,
        show_copy_button=True
    )

    with gr.Row():
        user_input = gr.Textbox(
            placeholder="Describe your symptoms or ask a question (e.g., 'I have a fever and sore throat for two days')...",
            label="🧍 Describe Your Symptoms",
            lines=3,
            scale=4
        )

    with gr.Row():
        send_btn = gr.Button("💬 Ask Dr. Aiden", variant="primary", size="lg")
        clear_btn = gr.Button("🧹 New Consultation", size="lg")

    gr.Markdown("### 💡 Example Questions")
    gr.Examples(
        examples=[
            "I have a fever and headache for two days. What should I take?",
            "I feel tired all day and have trouble sleeping. What could be wrong?",
            "I have mild chest tightness when I exercise. Should I worry?",
            "I'm feeling anxious and stressed. Any natural remedies?",
            "I have stomach pain after eating. What can I do?",
            "I caught a cold and sore throat. What treatment do you recommend?",
        ],
        inputs=user_input,
    )

    # =======================================================
    # Respond Function (stateless model, persistent chat)
    # =======================================================
    def respond(message, history):
        user_message = message.strip()
        if not user_message:
            return "", history

        # Show user input
        history.append({"role": "user", "content": user_message})

        # Model sees only current input (no memory)
        temp_history = [{"role": "user", "content": user_message}]

        for updated_history in generate_doctor_response(temp_history):
            if len(history) == 0 or history[-1]["role"] != "assistant":
                history.append({"role": "assistant", "content": updated_history[-1]["content"]})
            else:
                history[-1]["content"] = updated_history[-1]["content"]
            yield "", history

    # =======================================================
    # Button Bindings
    # =======================================================
    send_btn.click(respond, [user_input, chatbot], [user_input, chatbot])
    user_input.submit(respond, [user_input, chatbot], [user_input, chatbot])
    clear_btn.click(lambda: [], None, chatbot, queue=False)


# =======================================================
# Launch App
# =======================================================
if __name__ == "__main__":
    print("="*60)
    print("🏥 Dr. Aiden – AI Medical Doctor is starting...")
    print("="*60)
    demo.queue(max_size=20)
    demo.launch(
        share=True,
        show_error=True,
        server_name="0.0.0.0"
    )