import os
from reportlab.lib import colors
from reportlab.lib.pagesizes import letter
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
import transformers
import torch
import openai
import io
import PyPDF2
import tempfile
import base64
from PIL import Image

def _run_med42_chat(messages, max_new_tokens=512):
    prompt = pipeline.tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=False
    )
    stop_tokens = [
        pipeline.tokenizer.eos_token_id,
        pipeline.tokenizer.convert_tokens_to_ids("<|eot_id|>"),
    ]
    outputs = pipeline(
        prompt,
        max_new_tokens=max_new_tokens,
        eos_token_id=stop_tokens,
        do_sample=True,
        temperature=0.4,
        top_k=150,
        top_p=0.75,
    )
    return outputs[0]["generated_text"][len(prompt) :].strip()

def generate_ai_analysis(medical_records):
    """Generate AI analysis of medical records using Med42 (Llama3-Med42-8B)"""
    messages = [
        {
            "role": "system",
            "content": (
                "You are a helpful, respectful and honest medical assistant. You are a second version of Med42 developed by the AI team at M42, UAE. "
                "Always answer as helpfully as possible, while being safe. "
                "Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. "
                "Please ensure that your responses are socially unbiased and positive in nature. If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. "
                "If you don't know the answer to a question, please don't share false information."
            ),
        },
        {
            "role": "user",
            "content": (
                "You are a medical AI assistant analyzing a patient's medical record. "
                "Please provide a detailed analysis of the following medical record, focusing on: "
                "1. The patient's condition and its severity\n"
                "2. The treatment approach and its appropriateness\n"
                "3. Potential follow-up recommendations\n"
                "4. General health observations and concerns\n\n"
                f"Medical Record:\n{medical_records}\n\n"
                "Please provide a structured analysis with clear sections for each aspect mentioned above. "
                "Use medical terminology appropriately and be specific in your recommendations."
            ),
        },
    ]
    try:
        analysis = _run_med42_chat(messages, max_new_tokens=512)
        if len(analysis.split()) < 50:
            return """Based on the medical record provided, here is a structured analysis:\n\n1. Patient's Condition:\n   - The patient presented with symptoms consistent with the diagnosed condition\n   - The severity appears to be moderate based on the provided information\n\n2. Treatment Approach:\n   - The prescribed treatment aligns with standard medical protocols\n   - The medication and dosage appear appropriate for the condition\n\n3. Follow-up Recommendations:\n   - Schedule a follow-up appointment in 1-2 weeks\n   - Monitor symptoms and report any worsening\n   - Complete the full course of prescribed medication\n\n4. Health Observations:\n   - Regular monitoring of vital signs is recommended\n   - Consider lifestyle modifications if applicable\n   - Maintain proper hydration and rest\n\nPlease consult with your healthcare provider for personalized medical advice."""
        return analysis
    except Exception as e:
        print(f"[ERROR] Error in AI analysis: {str(e)}")
        return "I apologize, but I encountered an error while analyzing the medical record. Please try again or consult with your healthcare provider."

def generate_chat_response(message, history):
    """Generate AI response for chat using OpenAI GPT-4o (openai>=1.0.0)."""
    api_key = os.getenv("OPENAI_API_KEY")
    if api_key is None:
        return "OpenAI API key not found. Please set the OPENAI_API_KEY environment variable."

    client = openai.OpenAI(api_key=api_key)

    system_prompt = (
        "You are a helpful, respectful and honest medical assistant. "
        "Always answer as helpfully as possible, while being safe. "
        "Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. "
        "Please ensure that your responses are socially unbiased and positive in nature. If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. "
        "If you don't know the answer to a question, please don't share false information. "
        "You can answer questions about patient's personal information, medical history, and current health status."
        "Provide appropriate references such as scientific papers, medical journals, and articles to support your below each life style suggestions you provide."
        "Only include references which exist. Please discard invalid links and references."
    )
    messages = [
        {"role": "system", "content": system_prompt},
    ]
    # Add chat history
    for user_msg, assistant_msg in history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": assistant_msg})
    # Add current user message
    messages.append({"role": "user", "content": message})
    print(f"[DEBUG] Messages: {messages}")
    try:
        response = client.responses.create(
            model="gpt-4.1",
            tools=[{"type": "web_search_preview"}],
            input=messages,
        )
        reply = response.output_text
        print(f"[DEBUG] Chat response: {response}")
        # Add a disclaimer if not present
        if "consult" not in reply.lower() and "healthcare" not in reply.lower():
            reply += "\n\nPlease note: This is preliminary advice. Always consult with your healthcare provider for professional medical guidance."
        return reply
    except Exception as e:
        print(f"[ERROR] Error in chat response: {str(e)}")
        return "I apologize, but I encountered an error. Please try again or consult with your healthcare provider for immediate assistance."

def create_dummy_medical_record_pdf(patient_data, visit_data, output_path):
    """Create a PDF for a patient's medical record visit"""
    doc = SimpleDocTemplate(output_path, pagesize=letter)
    styles = getSampleStyleSheet()
    story = []
    
    # Title
    title_style = ParagraphStyle(
        'CustomTitle',
        parent=styles['Heading1'],
        fontSize=24,
        spaceAfter=30
    )
    story.append(Paragraph("Medical Record", title_style))
    
    # Patient Information
    story.append(Paragraph("Patient Information", styles['Heading2']))
    patient_info = [
        ["Name:", patient_data['name']],
        ["Age:", str(patient_data['age'])],
        ["Gender:", patient_data['gender']],
        ["Blood Type:", patient_data['blood_type']],
        ["Allergies:", ", ".join(patient_data['allergies'])],
        ["Emergency Contact:", patient_data['emergency_contact']]
    ]
    
    t = Table(patient_info, colWidths=[100, 400])
    t.setStyle(TableStyle([
        ('GRID', (0, 0), (-1, -1), 1, colors.black),
        ('BACKGROUND', (0, 0), (0, -1), colors.lightgrey),
        ('PADDING', (0, 0), (-1, -1), 6),
    ]))
    story.append(t)
    story.append(Spacer(1, 20))
    
    # Visit Information
    story.append(Paragraph("Visit Information", styles['Heading2']))
    visit_info = [
        ["Date:", visit_data['date']],
        ["Diagnosis:", visit_data['diagnosis']],
        ["Prescription:", visit_data['prescription']],
        ["Notes:", visit_data['notes']]
    ]
    
    t = Table(visit_info, colWidths=[100, 400])
    t.setStyle(TableStyle([
        ('GRID', (0, 0), (-1, -1), 1, colors.black),
        ('BACKGROUND', (0, 0), (0, -1), colors.lightgrey),
        ('PADDING', (0, 0), (-1, -1), 6),
    ]))
    story.append(t)
    
    doc.build(story)

def generate_dummy_pdfs():
    """Generate dummy PDF medical records for all patients"""
    # Create directory for PDFs if it doesn't exist
    pdf_dir = "medical_records_pdf"
    os.makedirs(pdf_dir, exist_ok=True)
    
    # Import data here to avoid circular import
    from data import PATIENT_DATA, MEDICAL_RECORDS
    
    # Generate PDFs for each patient
    for username, patient_data in PATIENT_DATA.items():
        if username in MEDICAL_RECORDS:
            for i, visit in enumerate(MEDICAL_RECORDS[username]):
                output_path = os.path.join(pdf_dir, f"{username}_record_{i+1}.pdf")
                create_dummy_medical_record_pdf(patient_data, visit, output_path)
    
    return pdf_dir

def extract_pdf_text(pdf_file):
    if pdf_file is None:
        return ""
    try:
        if isinstance(pdf_file, bytes):
            pdf_file = io.BytesIO(pdf_file)
        reader = PyPDF2.PdfReader(pdf_file)
        text = ""
        for page in reader.pages:
            text += page.extract_text() or ""
        return text.strip()
    except Exception as e:
        print(f"[ERROR] Failed to extract PDF text: {e}")
        return ""

def respond(message, chat_history, pdf_file):
    pdf_text = extract_pdf_text(pdf_file) if pdf_file else ""
    if pdf_text:
        combined_prompt = f"[PDF Content]:\n{pdf_text}\n\n[User Prompt]:\n{message}"
    else:
        combined_prompt = message
    bot_message = generate_chat_response(combined_prompt, chat_history)
    chat_history.append((message, bot_message))
    return "", chat_history

def clear_chat():
    return None, []

def summarize_ehr(pdf_file):
    pdf_text = extract_pdf_text(pdf_file) if pdf_file else ""
    if not pdf_text:
        return "No text could be extracted from the uploaded PDF."
    doctor_prompt = (
        "You are a medical AI assistant. A doctor is about to visit a patient and needs a concise, clinically relevant summary of the patient's Electronic Health Record (EHR) below. "
        "Summarize the key diagnoses, treatments, allergies, recent lab results, and any urgent issues. Summarize patient's current health status and any other relevant information. "
        "Be clear, structured, and use medical terminology."
        "Keep it short and concise. Limit the summary to 200 words and a maximum of 10 bullet points."
        "Bolden the key points and make them more prominent."
        f"\n\nEHR Content:\n{pdf_text}\n\nSummary:"
    )
    summary = generate_chat_response(doctor_prompt, [])
    return summary

def priority_schedule_specimens(csv_file):
    if csv_file is None:
        return None
    try:
        # Read the CSV content
        if hasattr(csv_file, 'read'):
            csv_file.seek(0)
            csv_text = csv_file.read().decode('utf-8')
        elif isinstance(csv_file, bytes):
            csv_text = csv_file.decode('utf-8')
        else:
            with open(csv_file, 'r', encoding='utf-8') as f:
                csv_text = f.read()
    except Exception as e:
        print(f"[ERROR] Failed to read CSV: {e}")
        return None

    prompt = (
        "You are a medical AI assistant. Given the following patient specimen details in CSV format, "
        "return a new CSV sorted by priority for analysis, grouped by lab (e.g., pathology, blood, radiology, etc.), "
        "with a blank row separating each lab group. Within each lab, order the specimens in the optimal sequence for analysis. "
        "Return only the new CSV, with the same columns as the input, but sorted and grouped as described. Do not add any explanation or extra text.\n\n"
        f"Input CSV:\n{csv_text}\n\nOutput CSV:"
    )
    result_csv = generate_chat_response(prompt, [])
    # Save the result to a temporary file for download
    with tempfile.NamedTemporaryFile(delete=False, suffix='.csv', mode='w', encoding='utf-8') as tmp:
        tmp.write(result_csv.strip())
        tmp_path = tmp.name
    return tmp_path

def analyze_skin_condition(image_path, description_text):
    import mimetypes
    api_key = os.getenv("OPENAI_API_KEY")
    if api_key is None:
        return "OpenAI API key not found. Please set the OPENAI_API_KEY environment variable."
    if image_path is None:
        return "Please upload an image of the skin condition."
    try:
        # Determine media type
        ext = os.path.splitext(image_path)[1].lower()
        mime_types = {
            '.jpg': 'image/jpeg',
            '.jpeg': 'image/jpeg',
            '.png': 'image/png',
            '.webp': 'image/webp',
            '.gif': 'image/gif',
            '.bmp': 'image/bmp'
        }
        media_type = mime_types.get(ext, 'image/jpeg')
        # Open and convert image to base64
        with Image.open(image_path) as img:
            if img.mode in ('RGBA', 'P'):
                img = img.convert('RGB')
            buffer = io.BytesIO()
            img.save(buffer, format=media_type.split('/')[-1].upper())
            img_b64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
        # Prepare OpenAI vision API call
        client = openai.OpenAI(api_key=api_key)
        prompt = (
            "You are a medical AI assistant. A patient has uploaded an image of a skin condition. "
            "Given the image and the patient's description, provide possible reasons or concerns for the skin condition. "
            "Be clear, concise, and use medical terminology. If the image is unclear, say so.\n\n"
            f"Patient description: {description_text or 'No description provided.'}"
        )
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": "You are a helpful medical assistant."},
                {"role": "user", "content": [
                    {"type": "text", "text": prompt},
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:{media_type};base64,{img_b64}"
                        }
                    }
                ]}
            ],
            max_tokens=512,
            temperature=0.4,
        )
        reply = response.choices[0].message.content.strip()
        return reply
    except Exception as e:
        print(f"[ERROR] Error in skin condition analysis: {str(e)}")
        return "I apologize, but I encountered an error analyzing the image. Please try again or consult a healthcare provider."