import streamlit as st
import os
from google import genai
from google.genai import types
from pydantic import BaseModel, Field
from PIL import Image
import requests
from datetime import datetime

# 1. Enhanced Schema (Added NPI for the API call)
class ProviderLicense(BaseModel):
    provider_name: str = Field(description="Full name of the healthcare provider")
    license_number: str = Field(description="The professional license number")
    npi_number: str = Field(description="The 10-digit National Provider Identifier (NPI)")
    state: str = Field(description="The state where the license was issued")
    expiration_date: str = Field(description="Format: YYYY-MM-DD")

# 2. Page Configuration
st.set_page_config(page_title="AI Credentialing Assistant", layout="wide")
st.title("🩺 Provider Credentialing AI")

# 3. Helper Function for NPPES
def get_nppes_data(npi_number):
    url = "https://npiregistry.cms.hhs.gov/api/?version=2.1"
    params = {"number": npi_number}
    try:
        response = requests.get(url, params=params)
        if response.status_code == 200:
            data = response.json()
            return data["results"][0] if data.get("result_count", 0) > 0 else None
        return None
    except Exception as e:
        st.error(f"NPPES API Error: {e}")
        return None

# 4. Initialize Client
api_key = os.environ.get("GEMINI_API_KEY")
client = genai.Client(api_key=api_key)

# 5. Sidebar & Main Logic
uploaded_file = st.sidebar.file_uploader("Upload Medical License", type=["jpg", "jpeg", "png", "pdf"])

if uploaded_file:
    col1, col2 = st.columns(2)
    image = Image.open(uploaded_file)
    
    with col1:
        st.subheader("Document Preview")
        st.image(image, use_container_width=True)

    with col2:
        st.subheader("AI Extraction & Registry Match")
        with st.spinner("Analyzing and Verifying..."):
            try:
                # AI Extraction
                response = client.models.generate_content(
                    model="gemini-2.5-flash", # Verify model ID in your environment
                    contents=["Extract details from this license.", image],
                    config=types.GenerateContentConfig(
                        response_mime_type="application/json",
                        response_schema=ProviderLicense,
                    ),
                )
                data = response.parsed

                # Logic: Is Valid? (Calculated, not guessed)
                expiry = datetime.strptime(data.expiration_date, "%Y-%m-%d").date()
                is_active = expiry >= datetime.today().date()

                # Display Results
                st.metric("Provider Name", data.provider_name)
                st.write(f"**License:** {data.license_number} ({data.state})")
                
                status_color = "green" if is_active else "red"
                st.markdown(f"**Status:** :{status_color}[{ 'Valid' if is_active else 'Expired'}] (Expires: {data.expiration_date})")

                # --- NEW: Automated NPPES Verification ---
                st.divider()
                st.subheader("Federal Registry Verification")
                registry_data = get_nppes_data(data.npi_number)

                if registry_data:
                    # Based on your st.json, registry_data is the single record dictionary
                    basic_info = registry_data.get('basic', {})
                    
                    # Check if it's an Organization (NPI-2) or Individual (NPI-1)
                    if registry_data.get('enumeration_type') == 'NPI-2':
                        display_name = basic_info.get('organization_name', 'Unnamed Organization')
                        st.info(f"Verified as Organizational NPI: {display_name}")
                    else:
                        f_name = basic_info.get('first_name', '')
                        l_name = basic_info.get('last_name', '')
                        display_name = f"{f_name} {l_name}".strip()
                        st.success(f"Verified as Individual Provider: {display_name}")
                
                    st.write(f"**Registry Name:** {display_name}")
                    
                    # Taxonomy parsing (based on your JSON structure)
                    taxonomies = registry_data.get('taxonomies', [])
                    if taxonomies:
                        st.write(f"**Primary Taxonomy:** {taxonomies[0].get('desc', 'N/A')}")
                        st.write(f"**Associated License:** {taxonomies[0].get('license', 'N/A')}")
            
            except Exception as e:
                st.error(f"Processing Error: {e}")