import streamlit as st
import os
import re
import json
import base64
from datetime import datetime
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Page configuration - MUST be first Streamlit command
st.set_page_config(
    page_title="MediScan AI - Understand Your Medical Reports",
    page_icon="🩺",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Import libraries after page config
try:
    import pdfplumber
except ImportError:
    st.error("pdfplumber not installed. Please check requirements.")
    st.stop()

try:
    from groq import Groq
except ImportError:
    st.error("groq not installed. Please check requirements.")
    st.stop()

try:
    from PIL import Image
except ImportError:
    st.error("Pillow not installed. Please check requirements.")
    st.stop()

try:
    import pytesseract
except ImportError:
    st.warning("pytesseract not fully configured. OCR may not work properly.")
    pytesseract = None

# Custom CSS
st.markdown("""
<style>
    .main-header {
        font-size: 2.5rem;
        color: #2c3e50;
        text-align: center;
        margin-bottom: 1rem;
    }
    .sub-header {
        font-size: 1.2rem;
        color: #7f8c8d;
        text-align: center;
        margin-bottom: 2rem;
    }
    .alert-box {
        background-color: #fee2e2;
        padding: 1rem;
        border-radius: 10px;
        border-left: 5px solid #ef4444;
        margin: 1rem 0;
    }
    .info-box {
        background-color: #e0f2fe;
        padding: 1rem;
        border-radius: 10px;
        border-left: 5px solid #3b82f6;
        margin: 1rem 0;
    }
    .success-box {
        background-color: #dcfce7;
        padding: 1rem;
        border-radius: 10px;
        border-left: 5px solid #22c55e;
        margin: 1rem 0;
    }
    .stButton > button {
        background-color: #2c3e50;
        color: white;
        font-weight: bold;
        width: 100%;
    }
</style>
""", unsafe_allow_html=True)

# Initialize Groq client
GROQ_API_KEY = os.getenv("GROQ_API_KEY")

def extract_text_from_pdf(uploaded_file):
    """Extract text from PDF using pdfplumber"""
    try:
        text = ""
        with pdfplumber.open(uploaded_file) as pdf:
            for page in pdf.pages:
                page_text = page.extract_text()
                if page_text:
                    text += page_text + "\n"
        return text if text.strip() else None
    except Exception as e:
        st.error(f"Error reading PDF: {str(e)}")
        return None

def extract_text_from_image(uploaded_file):
    """Extract text from image using Tesseract OCR"""
    try:
        if pytesseract is None:
            st.error("Tesseract OCR is not configured")
            return None
        
        image = Image.open(uploaded_file)
        # Convert to grayscale for better OCR
        image = image.convert('L')
        text = pytesseract.image_to_string(image, config='--psm 3')
        return text if text.strip() else None
    except Exception as e:
        st.error(f"Error reading image: {str(e)}")
        return None

def analyze_medical_report(report_text, language="english"):
    """Analyze medical report using Groq API"""
    
    if not GROQ_API_KEY:
        return {"error": "GROQ_API_KEY not set. Please add your API key to Secrets."}
    
    try:
        client = Groq(api_key=GROQ_API_KEY)
    except Exception as e:
        return {"error": f"Failed to initialize Groq client: {str(e)}"}
    
    system_prompt = """You are MediScan AI, a compassionate medical assistant. Analyze the lab report and provide response in this exact JSON format:

{
    "extracted_values": [
        {"test_name": "Test Name", "value": "value", "unit": "unit", "reference_range": "range", "status": "normal/high/low"}
    ],
    "summary": "Brief 2-sentence summary in simple language",
    "critical_alerts": ["Alert 1", "Alert 2"],
    "simple_explanations": {"Test Name": "Simple explanation in everyday language"},
    "questions_for_doctor": ["Question 1", "Question 2"],
    "recommendations": ["Recommendation 1", "Recommendation 2"],
    "disclaimer": "This is AI assistance, not medical advice"
}

Rules:
- Mark status as 'low', 'normal', or 'high' based on reference range
- If value is critically abnormal, add to critical_alerts
- Keep explanations at 6th grade reading level
- Be supportive and not alarming"""

    try:
        response = client.chat.completions.create(
            model="llama-3.1-8b-instant",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"Analyze this medical report in {language}: {report_text[:8000]}"}
            ],
            temperature=0.2,
            response_format={"type": "json_object"}
        )
        
        result = json.loads(response.choices[0].message.content)
        return result
    
    except json.JSONDecodeError:
        # Fallback - try to extract JSON
        content = response.choices[0].message.content
        json_match = re.search(r'\{.*\}', content, re.DOTALL)
        if json_match:
            return json.loads(json_match.group())
        return {"error": "Failed to parse AI response"}
    except Exception as e:
        return {"error": f"Analysis error: {str(e)}"}

def display_results(analysis_data):
    """Display analysis results in a beautiful format"""
    
    if not analysis_data or "error" in analysis_data:
        st.error(analysis_data.get("error", "Unable to analyze report"))
        return
    
    # Critical Alerts
    if analysis_data.get('critical_alerts') and analysis_data['critical_alerts']:
        st.markdown('<div class="alert-box">', unsafe_allow_html=True)
        st.markdown("## 🚨 CRITICAL ALERTS")
        for alert in analysis_data['critical_alerts']:
            st.markdown(f"⚠️ {alert}")
        st.markdown('</div>', unsafe_allow_html=True)
    
    # Summary
    if analysis_data.get('summary'):
        st.markdown('<div class="info-box">', unsafe_allow_html=True)
        st.markdown("## 📋 Summary")
        st.write(analysis_data['summary'])
        st.markdown('</div>', unsafe_allow_html=True)
    
    # Test Results Table
    if analysis_data.get('extracted_values'):
        st.markdown("## 📊 Test Results")
        for test in analysis_data['extracted_values']:
            status_emoji = {
                "normal": "✅", 
                "high": "🔴", 
                "low": "🔵"
            }.get(test.get('status', 'normal'), "⚪")
            
            st.markdown(f"""
            **{status_emoji} {test.get('test_name', 'Unknown')}**  
            - Your value: `{test.get('value', '?')} {test.get('unit', '')}`  
            - Normal range: `{test.get('reference_range', 'N/A')}`  
            - Status: **{test.get('status', 'unknown').upper()}**
            """)
            st.markdown("---")
    
    # Simple Explanations
    if analysis_data.get('simple_explanations'):
        st.markdown("## 💡 What This Means")
        for test_name, explanation in analysis_data['simple_explanations'].items():
            with st.expander(f"📖 {test_name}"):
                st.write(explanation)
    
    # Recommendations
    if analysis_data.get('recommendations'):
        st.markdown("## ✅ Recommendations")
        for rec in analysis_data['recommendations']:
            st.markdown(f"• {rec}")
    
    # Questions for Doctor
    if analysis_data.get('questions_for_doctor'):
        st.markdown("## 🗣️ Questions to Ask Your Doctor")
        for q in analysis_data['questions_for_doctor']:
            st.markdown(f"• {q}")
    
    # Disclaimer
    st.markdown("---")
    st.caption(analysis_data.get('disclaimer', "⚠️ This is AI assistance. Always consult a healthcare provider."))

# Main App
def main():
    st.markdown('<p class="main-header">🩺 MediScan AI</p>', unsafe_allow_html=True)
    st.markdown('<p class="sub-header">Upload your medical report — Get simple, understandable insights in plain language</p>', unsafe_allow_html=True)
    
    # Sidebar
    with st.sidebar:
        st.image("https://img.icons8.com/color/96/medical-report.png", width=80)
        st.markdown("## About")
        st.info("""
        **MediScan AI** helps you understand medical reports by:
        - Extracting values from PDFs or images
        - Flagging abnormal results
        - Explaining in simple language
        - Suggesting questions for your doctor
        """)
        
        st.markdown("## Language")
        language = st.radio("Select Language", ["English", "Roman Urdu"], index=0)
        
        if not GROQ_API_KEY:
            st.error("⚠️ GROQ_API_KEY not found!")
            st.markdown("Please add your API key in Hugging Face Spaces Secrets:")
            st.code("Settings → Repository Secrets → New Secret\nName: GROQ_API_KEY\nValue: your_groq_api_key")
        
        st.markdown("## How to Use")
        st.markdown("""
        1. Upload your lab report (PDF or photo)
        2. Click 'Analyze Report'
        3. Review the simplified results
        """)
    
    # File Upload
    uploaded_file = st.file_uploader(
        "📄 Upload your medical report",
        type=["pdf", "png", "jpg", "jpeg"],
        help="Supports PDF files and images of lab reports"
    )
    
    if uploaded_file:
        # Preview
        if uploaded_file.type == "application/pdf":
            st.success(f"✅ PDF loaded: {uploaded_file.name}")
            # Show PDF preview
            try:
                base64_pdf = base64.b64encode(uploaded_file.getvalue()).decode('utf-8')
                pdf_display = f'<iframe src="data:application/pdf;base64,{base64_pdf}" width="100%" height="400" type="application/pdf"></iframe>'
                st.markdown(pdf_display, unsafe_allow_html=True)
            except Exception as e:
                st.warning(f"Preview not available: {str(e)}")
        else:
            st.image(uploaded_file, caption="Uploaded Report", use_container_width=True)
        
        # Analyze button
        if st.button("🔍 Analyze Report", type="primary"):
            with st.spinner("📖 Extracting text from report..."):
                if uploaded_file.type == "application/pdf":
                    report_text = extract_text_from_pdf(uploaded_file)
                else:
                    report_text = extract_text_from_image(uploaded_file)
            
            if report_text:
                with st.spinner("🧠 AI is analyzing your medical data..."):
                    analysis = analyze_medical_report(report_text, "urdu" if language == "Roman Urdu" else "english")
                
                if analysis and "error" not in analysis:
                    st.session_state['analysis'] = analysis
                    st.success("✅ Analysis complete!")
                else:
                    st.error(analysis.get("error", "Failed to analyze report"))
            else:
                st.error("Could not extract text. Please ensure the report is clear and readable.")
    
    # Display results
    if 'analysis' in st.session_state:
        st.markdown("---")
        display_results(st.session_state['analysis'])

if __name__ == "__main__":
    main()