import streamlit as st
import keras
from keras.preprocessing.image import img_to_array
import numpy as np
from PIL import Image
import base64
import torch
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
from reportlab.lib.pagesizes import letter
from reportlab.pdfgen import canvas
import io
import logging

# -------------- Logging Configuration --------------
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger()

# -------------- Page Config --------------
st.set_page_config(page_title="VisiHealth AI", layout="centered")

# -------------- Custom CSS --------------
def set_custom_styles(bg_img_path):
    try:
        with open(bg_img_path, "rb") as img_file:
            encoded = base64.b64encode(img_file.read()).decode()
        css = f"""
        <style>
        .stApp {{
            background-image: url("data:image/png;base64,{encoded}");
            background-size: cover;
            background-repeat: no-repeat;
            background-position: center;
            color: #FFFFFF;
        }}
        .st-emotion-cache-1v0mbdj {{
            background-color: rgba(0, 0, 0, 0.75);
            border-radius: 15px;
            padding: 20px;
        }}
        .stButton > button {{
            background-color: #0ff;
            color: black;
            font-weight: bold;
            border: none;
            padding: 10px 20px;
            border-radius: 10px;
            box-shadow: 0 0 15px #0ff;
            transition: all 0.3s ease-in-out;
        }}
        .stButton > button:hover {{
            background-color: #00e5ff;
            box-shadow: 0 0 25px #00e5ff;
            transform: scale(1.05);
        }}
        </style>
        """
        st.markdown(css, unsafe_allow_html=True)
    except Exception as e:
        logger.error(f"Error in set_custom_styles: {e}")

set_custom_styles("img1.png")

# -------------- Cache and Load Models --------------
@st.cache_resource
def load_hair_model():
    try:
        model = keras.models.load_model('VGG19-Final.h5')
        return model
    except Exception as e:
        logger.error(f"Error loading hair model: {e}")

@st.cache_resource
def load_skin_model():
    try:
        model = keras.models.load_model('skin_cancer_detection_model.h5')
        return model
    except Exception as e:
        logger.error(f"Error loading skin model: {e}")

@st.cache_resource
def load_dental_model():
    try:
        model = keras.models.load_model('oral_disease.h5')
        return model
    except Exception as e:
        logger.error(f"Error loading dental model: {e}")

@st.cache_resource
def load_cure_generator():
    try:
        tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
        model = AutoModelForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
        return pipeline("text-generation", model=model, tokenizer=tokenizer, device=0)
    except Exception as e:
        logger.error(f"Error loading cure generator: {e}")

hair_model = load_hair_model()
skin_model = load_skin_model()
dental_model = load_dental_model()
cure_generator = load_cure_generator()

# -------------- Disease Labels --------------
hair_diseases = ['Alopecia Areata', 'Contact Dermatitis', 'Folliculitis', 'Head Lice', 'Lichen Planus',
                 'Male Pattern Baldness', 'Psoriasis', 'Seborrheic Dermatitis', 'Telogen Effluvium', 'Tinea Capitis']

skin_diseases = ['Actinic Keratoses', 'Basal Cell Carcinoma', 'Benign Keratosis', 'Dermatofibroma',
                 'Melanoma', 'Nevus', 'Vascular Lesion']

dental_diseases = ['Calculus', 'Caries_Gingivitus_ToothDiscoloration', 'Data caries', 'Gingivitis',
                   'Mouth Ulcer', 'Tooth Discoloration', 'hypodontia']

# -------------- Cure Generator --------------
def generate_cure_details(disease_name):
    try:
        prompt = f"""<|user|> Provide detailed medical information for the disease: {disease_name}:
1. Description
2. Symptoms
3. Possible Treatments
4. Cost Estimate (in Rs. or USD)
5. Estimated Recovery Time
<|assistant|>"""
        result = cure_generator(prompt, max_length=2048, do_sample=True, temperature=0.7)[0]['generated_text']
        return result
    except Exception as e:
        logger.error(f"Error generating cure details for {disease_name}: {e}")

# -------------- Preprocessing --------------
def preprocess_image(img):
    try:
        img = img.resize((224, 224))
        return np.expand_dims(np.array(img) / 255.0, axis=0)
    except Exception as e:
        logger.error(f"Error preprocessing image: {e}")

# -------------- Predict Functions --------------
def predict_hair_disease(img_array):
    try:
        return hair_diseases[np.argmax(hair_model.predict(img_array), axis=1)[0]]
    except Exception as e:
        logger.error(f"Error predicting hair disease: {e}")

def predict_skin_disease(img_array):
    try:
        return skin_diseases[np.argmax(skin_model.predict(img_array))]
    except Exception as e:
        logger.error(f"Error predicting skin disease: {e}")

def predict_dental_disease(img_array):
    try:
        return dental_diseases[np.argmax(dental_model.predict(img_array))]
    except Exception as e:
        logger.error(f"Error predicting dental disease: {e}")

# -------------- PDF Generation --------------
def generate_pdf(cure_info, disease_name):
    try:
        buffer = io.BytesIO()
        c = canvas.Canvas(buffer, pagesize=letter)
        c.setFont("Helvetica", 12)
        c.drawString(100, 750, f"VisiHealth AI - Disease: {disease_name}")
        c.drawString(100, 730, "--------------------------------------------------------")
        y_position = 710
        for line in cure_info.split("\n"):
            c.drawString(100, y_position, line)
            y_position -= 15
        c.save()
        buffer.seek(0)
        return buffer
    except Exception as e:
        logger.error(f"Error generating PDF: {e}")

# -------------- Main App --------------
st.markdown("<h1 style='text-align:center; color:#00e5ff;'>VisiHealth AI 🧬</h1>", unsafe_allow_html=True)
st.markdown("<h4 style='text-align:center;'>AI-powered Disease Detection and Cure Report Generator</h4>", unsafe_allow_html=True)

if "last_disease" not in st.session_state:
    st.session_state.last_disease = None
    st.session_state.last_cure_info = None

disease_type = st.radio("🩺 Select Disease Type:", ['Hair Disease', 'Skin Cancer', 'Dental Disease'])
uploaded_file = st.file_uploader("📸 Upload an Image", type=["jpg", "jpeg", "png"])

chat_history = []

if uploaded_file:
    img = Image.open(uploaded_file)
    st.image(img, caption="Uploaded Image", use_column_width=True)
    img_array = preprocess_image(img)

    if disease_type == "Hair Disease":
        disease_name = predict_hair_disease(img_array)
    elif disease_type == "Skin Cancer":
        disease_name = predict_skin_disease(img_array)
    else:
        disease_name = predict_dental_disease(img_array)

    st.success(f"✅ Detected: {disease_name}")

    if disease_name != st.session_state.last_disease:
        with st.spinner("Generating Cure Info... Please wait ⏳"):
            raw_info = generate_cure_details(disease_name)

        st.session_state.last_disease = disease_name
        st.session_state.last_cure_info = raw_info
        chat_history.append({'Disease': disease_name, 'Cure Info': raw_info})
    else:
        raw_info = st.session_state.last_cure_info

    st.markdown("### 💊 Cure Info:")
    st.text(raw_info)

    # PDF Download Button
    pdf_buffer = generate_pdf(raw_info, disease_name)
    st.download_button(
        label="Download Cure Info as PDF",
        data=pdf_buffer,
        file_name=f"{disease_name}_cure_info.pdf",
        mime="application/pdf"
    )

if st.button("💾 Save Session"):
    try:
        with open("chat_history.txt", "w") as f:
            for entry in chat_history:
                f.write(f"Disease: {entry['Disease']}\nCure Info:\n{entry['Cure Info']}\n\n")
        st.success("💬 Chat history saved!")
    except Exception as e:
        logger.error(f"Error saving chat history: {e}")
        st.error("Error saving chat history.")