import streamlit as st
import os
from google import genai
from google.genai import types

# --- 1. CONFIG & UI SETUP ---
st.set_page_config(page_title="Abyssinia Scribe", page_icon="📝", layout="wide")

# Custom CSS for a professional medical dashboard
st.markdown("""
    <style>
    .reportview-container { background: #fdfdfd; }
    .stTextArea textarea { font-family: 'Courier New', monospace; font-size: 14px; }
    </style>
    """, unsafe_allow_html=True)

client = genai.Client(api_key=os.environ.get("GOOGLE_API_KEY"))

# --- 2. CORE SCRIBE LOGIC ---
SCRIBE_PROMPT = """
Act as an expert Clinical Scribe for Abyssinia Intelligence. 
Analyze the provided medical transcript or summary and produce a high-fidelity SOAP note.
Ensure all medical terms are standardized (ICD-11/SNOMED).
"""

st.title("🏥 Abyssinia Intelligence Scribe")
st.info("Clinical Documentation & SOAP Note Automation System")

# Layout: 2 Columns (Input | Output)
col1, col2 = st.columns([1, 1], gap="large")

with col1:
    st.subheader("📋 Clinical Input")
    input_text = st.text_area("Patient Encounter / Transcript:", height=450, 
                             placeholder="Patient presents with 3 days of productive cough and low-grade fever...")
    
    # Advanced 2026 Settings
    with st.expander("Scribe Engine Settings"):
        reasoning_depth = st.select_slider("Reasoning Depth", options=["low", "high"], value="high")
        specialty = st.selectbox("Medical Specialty", ["General Practice", "Cardiology", "Pediatrics"])

with col2:
    st.subheader("📄 Structured SOAP Note")
    if st.button("Generate Final Note", type="primary"):
        if input_text:
            with st.spinner("Abyssinia AI is analyzing clinical data..."):
                try:
                    # Utilizing Gemini 3.0 Flash with Thinking Level
                    response = client.models.generate_content(
                        model="gemini-3-flash-preview",
                        config=types.GenerateContentConfig(
                            system_instruction=f"{SCRIBE_PROMPT} Context: {specialty}",
                            thinking_config=types.ThinkingConfig(
                                thinking_level=reasoning_depth
                            ),
                            temperature=0.2, # Low temperature for medical precision
                        ),
                        contents=input_text
                    )
                    
                    st.success("Analysis Complete")
                    final_note = response.text
                    st.text_area("Review/Edit Note:", value=final_note, height=450)
                    
                    # Download Feature
                    st.download_button("Export as TXT", final_note, file_name="Abyssinia_Note.txt")
                    
                except Exception as e:
                    st.error(f"Scribe Error: {e}")
        else:
            st.warning("Please enter encounter data first.")

# --- 3. AUDIT TRAIL ---
st.markdown("---")
st.caption("Abyssinia Scribe | v3.0-Flash | Powered by AHRI (Ethiopia)")