import streamlit as st
import os
import base64
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.messages import HumanMessage
from dotenv import load_dotenv

# Load API key from .env file
load_dotenv()
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")

# Initialize the Gemini model
llm = ChatGoogleGenerativeAI(
    model="gemini-2.0-flash",
    max_tokens=4000,
    api_key=GEMINI_API_KEY
)

def encode_file(file):
    return base64.b64encode(file.read()).decode()

def summarize_report(file, file_type):
    encoded_string = encode_file(file)
    
    prompt = """You are a highly advanced medical document analysis AI designed to assist live agents and doctors by extracting and summarizing fluctuating or improper test results from uploaded medical reports. Your task is to carefully parse the document, identify irregularities in test values, highlight trends that indicate instability, and generate a clear, structured summary.
    Ensure the summary includes:
    1) The test names with significant fluctuations or abnormal values.
    2) A brief comparison of fluctuating values over time (if historical data is available).
    3) Possible concerns based on deviations from normal ranges.
    4) A concise, professional summary that can be quickly reviewed by a medical expert or triage bot for primary screening.
    Format the output to be clear, concise, and medically relevant. Do not make any medical diagnoses—only summarize fluctuations and abnormalities in an easily interpretable manner."""
    
    content = [{"type": "text", "text": prompt}]
    
    if file_type in ["image/jpeg", "image/png", "image/webp"]:
        content.append({"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encoded_string}"}})
    else:
        content.append({"type": "text", "text": encoded_string})
    
    message = HumanMessage(content=content)
    ai_msg = llm.invoke([message])
    return ai_msg.content

# Streamlit UI
st.set_page_config(page_title="Medical Report Summarizer", layout="wide", page_icon="📄")

# Custom CSS for styling
st.markdown("""
    <style>
        .main {
            background-color: #f5f5f5;
        }
        .stButton>button {
            background-color: #4CAF50;
            color: white;
            border-radius: 8px;
            padding: 10px 20px;
        }
        .stSpinner {
            color: #4CAF50;
        }
    </style>
""", unsafe_allow_html=True)

# Header Section
st.title("📄 Call on Doc Medical Report Summarizer")
st.markdown("""
    <div style="background-color: #4CAF50; padding: 10px; border-radius: 10px;">
        <h2 style="color: white; text-align: center;">Upload your medical report and get a summarized analysis!</h2>
    </div>
""", unsafe_allow_html=True)

# Add demo images to the app
st.write("### Test with Demo Images:")
demo_images = {
    "Test Report 1": "report1.webp",
    "Test Report 2": "report2.webp",
    "Test Report 3": "report3.webp"
}

# Dropdown to select a demo image
selected_demo = st.selectbox("Choose a demo image to test:", options=["None"] + list(demo_images.keys()))

if selected_demo != "None":
    demo_image_path = demo_images[selected_demo]
    with open(demo_image_path, "rb") as file:
        demo_file = file.read()
    # st.image(demo_image_path, caption=f"Selected Demo: {selected_demo}", use_column_width=True)
    
    # Process the selected demo image
    with st.spinner("🔄 Processing the demo image..."):
        summary = summarize_report(file=open(demo_image_path, "rb"), file_type="image/jpeg")
    
    # Display Summary
    st.success("✅ Demo image processed successfully!")
    st.subheader("Summarized Report:")
    st.write(summary)

    # Download Button for Summary
    st.download_button(
        label="📥 Download Summary",
        data=summary,
        file_name=f"{selected_demo}_summary.txt",
        mime="text/plain"
    )

# File Upload Section
st.write("### Upload your medical report (image, PDF, or Word document):")
uploaded_file = st.file_uploader("", type=["pdf", "doc", "docx", "jpg", "jpeg", "png", "webp"])

# Process Uploaded File
if uploaded_file is not None:
    file_type = uploaded_file.type
    with st.spinner("🔄 Processing the document..."):
        summary = summarize_report(uploaded_file, file_type)
    
    # Display Summary
    st.success("✅ Document processed successfully!")
    st.subheader("Summarized Report:")
    st.write(summary)

    # Download Button for Summary
    st.download_button(
        label="📥 Download Summary",
        data=summary,
        file_name="summary.txt",
        mime="text/plain"
    )
else:
    st.info("ℹ️ Please upload a file to get started.")