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 Comprehensive Doctor Response
# =======================================================
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

    # Enhanced Medical Prompt - Comprehensive Doctor Approach
    prompt = f"""You are an experienced and compassionate medical doctor providing a comprehensive consultation.

Based on the patient's concern, provide a detailed response that includes:

1. **Assessment**: Acknowledge their symptoms and provide initial medical assessment
2. **Possible Causes**: Explain potential causes or conditions
3. **Medications**: Recommend appropriate over-the-counter or prescription medications (with dosages when relevant)
4. **Nutrition & Diet**: Suggest specific foods, nutrients, or dietary changes that can help
5. **Lifestyle Modifications**: Recommend lifestyle changes, exercises, rest, or habits to adopt
6. **Follow-up**: Advise when to see a doctor or what warning signs to watch for

Guidelines:
- Do NOT use labels like "Doctor:" or "Patient:" in your response
- Be professional, empathetic, and thorough like a real doctor
- Provide specific, actionable recommendations
- Use medical terminology but explain it simply
- Structure your response clearly with the categories above
- End with: "⚕️ *Please consult a healthcare provider for proper diagnosis and personalized treatment plan.*"

Patient's concern: {user_message}

Comprehensive Medical Response:"""

    # 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.92,
        top_k=50,
        do_sample=True,
        max_new_tokens=600,
        pad_token_id=tokenizer.pad_token_id,
        eos_token_id=tokenizer.eos_token_id,
        repetition_penalty=1.18,
        no_repeat_ngram_size=3
    )

    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 response
    response = clean_medical_response(response)

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

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


def clean_medical_response(response: str) -> str:
    """Clean and format the medical response."""
    # Remove common prefixes
    prefixes = ["assistant:", "doctor:", "response:", "comprehensive medical response:", "medical response:"]
    response_lower = response.lower()
    for prefix in prefixes:
        if response_lower.startswith(prefix):
            response = response[len(prefix):].strip()
            break
    
    # Remove any remaining role labels
    lines = response.split('\n')
    cleaned_lines = []
    for line in lines:
        if not line.lower().strip().startswith(('doctor:', 'assistant:', 'patient:')):
            cleaned_lines.append(line)
    response = '\n'.join(cleaned_lines)
    
    # Ensure proper ending
    if response and response[-1] not in '.!?':
        response += '.'
    
    # Add disclaimer if not present
    if '⚕️' not in response and 'consult' not in response.lower():
        response += '\n\n⚕️ *Please consult a healthcare provider for proper diagnosis and personalized treatment plan.*'
    
    # Fallback for very short responses
    if len(response.strip()) < 30:
        response = "I understand your concern. To provide you with comprehensive medical guidance including medications, diet, and lifestyle recommendations, could you please describe your symptoms in more detail? For example, when did they start, how severe are they, and have you noticed any triggers?"
    
    return response.strip()


# =======================================================
# Gradio Interface
# =======================================================
with gr.Blocks(theme=gr.themes.Soft(), css="""
    .medical-header {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        padding: 20px;
        border-radius: 10px;
        color: white;
        text-align: center;
        margin-bottom: 20px;
    }
""") as demo:
    
    gr.HTML("""
        <div class="medical-header">
            <h1>🏥 AI Medical Doctor Consultation</h1>
            <p>Comprehensive Medical Guidance • Medications • Nutrition • Lifestyle</p>
        </div>
    """)

    chatbot = gr.Chatbot(
        label="💬 Doctor-Patient Consultation",
        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 in detail (e.g., 'I have fever, headache, and body pain for 3 days')...",
            label="🧍 Describe Your Symptoms",
            lines=3,
            scale=4
        )

    with gr.Row():
        send_btn = gr.Button("💬 Consult Doctor", variant="primary", scale=1, size="lg")
        clear_btn = gr.Button("🧹 New Consultation", scale=1, size="lg")

    gr.Markdown("### 💡 Example Consultations")
    gr.Examples(
        examples=[
            "I have a fever of 102°F, headache, and body aches for 2 days. What should I do?",
            "I've been having persistent headaches and feeling tired. Need advice on diet and lifestyle.",
            "I have acidity and stomach pain after eating. What medications and diet should I follow?",
            "I'm feeling stressed and anxious. Suggest lifestyle changes and natural remedies.",
            "I have high blood pressure. What diet and lifestyle changes should I make?",
            "I caught a cold with sore throat and cough. What treatment do you recommend?",
        ],
        inputs=user_input,
    )

    gr.Markdown("""
    ---
    ### 🩺 What This AI Doctor Provides:
    
    ✅ **Medical Assessment** - Initial diagnosis and condition evaluation  
    ✅ **Medication Recommendations** - Appropriate medicines with dosage guidance  
    ✅ **Nutrition & Diet Plans** - Foods and nutrients to help recovery  
    ✅ **Lifestyle Modifications** - Exercise, sleep, stress management tips  
    ✅ **Follow-up Advice** - When to see a doctor and warning signs  
    
    ⚠️ **Important Medical Disclaimer:**  
    This AI provides general medical information for educational purposes only. It is NOT a substitute for professional medical advice, diagnosis, or treatment. Always seek the advice of your physician or qualified healthcare provider with any questions about a medical condition. Never disregard professional medical advice or delay seeking it because of something you have read here.
    
    🚨 **Emergency:** If you are experiencing a medical emergency, call emergency services immediately.
    """)

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

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

        # Model sees only current message (stateless for consistent behavior)
        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 & Input 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("🏥 AI Medical Doctor Consultation System Starting...")
    print("="*60)
    demo.queue(max_size=20)
    demo.launch(
        share=True,
        show_error=True,
        server_name="0.0.0.0"
    )