import streamlit as st
import os
import jwt
from datetime import datetime, timedelta
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
import torch
from PIL import Image
import io
import librosa
import numpy as np
import logging
import tempfile
from streamlit.runtime.uploaded_file_manager import UploadedFile
from diffusers import StableDiffusionPipeline
import sentry_sdk
from supabase import create_client, Client
import types
import pandas as pd
import plotly.express as px
from typing import Dict, Any

# --- Configuração Inicial ---
st.set_page_config(
    page_title="AiiT Services - AI Platform",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded"
)

# --- Estilos Personalizados ---
st.markdown("""
<style>
    :root {
    --primary: #7f5af0;       /* Roxo neon vibrante */
    --secondary: #2cb67d;     /* Verde menta brilhante */
    --accent: #f25f4c;        /* Laranja coral vibrante */
    --light: #fdfdfd;         /* Branco suave */
    --dark: #16161a;          /* Preto carvão elegante */
    --success: #2cb67d;       /* Verde sucesso */
    --info: #3d5afe;          /* Azul elétrico */
}
.main {
    background-color: #f5f7ff;  /* Cor de fundo geral */
}
.stApp {
    background-image: linear-gradient(120deg, #0f0c29, #302b63, #24243e); /* Fundo roxo/azul degradê */
}
.header {
    color: var(--light);
    border-bottom: 2px solid var(--primary);
    padding-bottom: 10px;
}
.card {
    background: var(--light);
    border-radius: 12px;
    box-shadow: 0 4px 20px rgba(127, 90, 240, 0.2);
    padding: 20px;
    margin-bottom: 20px;
    transition: transform 0.3s, box-shadow 0.3s;
}
.card:hover {
    transform: translateY(-5px);
    box-shadow: 0 6px 25px rgba(127, 90, 240, 0.4);
}
.feature-icon {
    font-size: 2.8rem;
    margin-bottom: 15px;
    color: var(--primary);
}
.btn-primary {
    background: linear-gradient(135deg, var(--primary), var(--secondary));
    border: none;
    color: white;
    border-radius: 10px;
    padding: 10px 20px;
    font-weight: 600;
    transition: all 0.3s;
}
.btn-primary:hover {
    transform: scale(1.05);
    box-shadow: 0 6px 15px rgba(44, 182, 125, 0.5);
}
.result-box {
    background: linear-gradient(135deg, #e0f7fa, #b2ebf2);
    border-radius: 12px;
    padding: 20px;
    margin-top: 20px;
}
.demo-box {
    background: linear-gradient(135deg, #f1f8e9, #dcedc8);
    border-radius: 12px;
    padding: 20px;
    margin-top: 20px;
}
.use-case-card {
    background: var(--light);
    border-radius: 12px;
    box-shadow: 0 4px 15px rgba(0, 0, 0, 0.08);
    padding: 20px;
    margin-bottom: 20px;
    border-left: 4px solid var(--accent);
    transition: all 0.3s;
}
.use-case-card:hover {
    transform: translateX(5px);
    box-shadow: 0 6px 20px rgba(242, 95, 76, 0.2);
}
</style>
""", unsafe_allow_html=True)

# --- Monkey-patch para compatibilidade ---
st.context = types.SimpleNamespace(cookies={})

import streamlit_authenticator as stauth

# --- Configuração de Logging e Sentry ---
sentry_sdk.init(os.getenv('SENTRY_DSN', ''), traces_sample_rate=1.0)
logging.basicConfig(
    filename='private/app_errors.log',
    level=logging.ERROR,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

# --- Configuração do Supabase ---
SUPABASE_URL = os.getenv("SUPABASE_URL")
SUPABASE_KEY = os.getenv("SUPABASE_KEY")
supabase: Client = create_client(SUPABASE_URL, SUPABASE_KEY) if SUPABASE_URL and SUPABASE_KEY else None

# --- Funções de Banco de Dados ---
def connect_db():
    """Connects to Supabase"""
    if not supabase:
        st.error("❌ Supabase configuration missing")
        return None
    return supabase

def init_db():
    """Initializes the users table in Supabase"""
    try:
        sb = connect_db()
        if not sb:
            return
        _ = sb.table("users").select("username").limit(1).execute()
    except Exception as e:
        logging.error(f"Error initializing database: {e}")

def load_users_from_db():
    """Loads users from Supabase"""
    try:
        sb = connect_db()
        if not sb:
            return None

        resp = sb.table("users").select("username, name, password, role").execute()
        rows = resp.data if hasattr(resp, 'data') else []

        users = {
            row['username']: {
                'name': row['name'],
                'password': row['password'],
                'role': row['role']
            }
            for row in rows
        }

        return {
            'credentials': {'usernames': users},
            'cookie': {'name': 'ai_app_cookie', 'key': 'random_key_123', 'expiry_days': 30},
            'preauthorized': {'emails': []}
        }

    except Exception as e:
        logging.error(f"Error loading users: {e}")
        return None

def admin_panel(authenticator):
    sb = connect_db()
    if not sb or st.session_state.get('role') != 'admin':
        st.warning("⚠️ Admins only")
        return

    st.sidebar.title("⚙️ Admin Panel")

    with st.sidebar.expander("➕ Add User"):
        with st.form("add_user_form"):
            u = st.text_input("Username")
            n = st.text_input("Name")
            p = st.text_input("Password", type="password")
            r = st.selectbox("Role", ["user", "admin"])
            if st.form_submit_button("Add", use_container_width=True):
                try:
                    hashed = stauth.Hasher([p]).generate()[0]
                    sb.table("users").insert({"username": u, "name": n, "password": hashed, "role": r}).execute()
                    st.success(f"✅ User '{u}' created!")
                except Exception as e:
                    st.error(f"❌ Error adding user: {str(e)}")

    with st.sidebar.expander("🗑️ Remove User"):
        with st.form("remove_user_form"):
            ur = st.text_input("Username to remove")
            if st.form_submit_button("Remove", use_container_width=True):
                try:
                    sb.table("users").delete().eq("username", ur).execute()
                    st.success(f"✅ User '{ur}' removed!")
                except Exception as e:
                    st.error(f"❌ Error removing user: {str(e)}")

    with st.sidebar.expander("🔄 Update Role"):
        with st.form("update_role_form"):
            uu = st.text_input("Username to update")
            nr = st.selectbox("New Role", ["user", "admin"])
            if st.form_submit_button("Update", use_container_width=True):
                try:
                    sb.table("users").update({"role": nr}).eq("username", uu).execute()
                    st.success(f"✅ Role of '{uu}' changed to '{nr}'")
                except Exception as e:
                    st.error(f"❌ Error updating role: {str(e)}")

    st.sidebar.subheader("👥 Registered Users")
    try:
        resp = sb.table("users").select("username,name,role").execute()
        for row in resp.data or []:
            st.sidebar.info(f"**{row['username']}** ({row['name']}) - `{row['role']}`")
    except Exception as e:
        st.error(f"❌ Error listing users: {str(e)}")

# --- Cache for Models ---
@st.cache_resource(show_spinner="Loading model...")
def load_model(model_key):
    """Loads specific model with persistent cache"""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    cache_dir = "model_cache"
    os.makedirs(cache_dir, exist_ok=True)
    
    model_loaders = {
        'sentiment_analysis': lambda: pipeline("sentiment-analysis", model="cardiffnlp/twitter-roberta-base-sentiment-latest", device=device),
        'text_classification': lambda: pipeline("text-classification", model="distilbert-base-uncased-finetuned-sst-2-english", device=device),
        'summarization': lambda: pipeline("summarization", model="facebook/bart-large-cnn", device=device, max_length=150, min_length=30),
        'question_answering': lambda: pipeline("question-answering", model="deepset/roberta-base-squad2", device=device),
        'translation': lambda: pipeline("translation", model="Helsinki-NLP/opus-mt-tc-big-en-pt", device=device),
        'text_generation': lambda: pipeline("text-generation", model="gpt2", device=device, pad_token_id=50256),
        'ner': lambda: pipeline("ner", model="dbmdz/bert-large-cased-finetuned-conll03-english", device=device, aggregation_strategy="simple"),
        'image_classification': lambda: pipeline("image-classification", model="google/vit-base-patch16-224", device=device),
        'object_detection': lambda: pipeline("object-detection", model="facebook/detr-resnet-50", device=device),
        'image_segmentation': lambda: pipeline("image-segmentation", model="facebook/detr-resnet-50-panoptic", device=device),
        'facial_recognition': lambda: pipeline("image-classification", model="mo-thecreator/vit-Facial-Expression-Recognition", device=device),
        'speech_to_text': lambda: pipeline("automatic-speech-recognition", model="openai/whisper-base", device=device),
        'audio_classification': lambda: pipeline("audio-classification", model="superb/hubert-base-superb-er", device=device),
        'text_to_image': lambda: StableDiffusionPipeline.from_pretrained(
            "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float32, use_safetensors=True, safety_checker=None, cache_dir=cache_dir
        ).to(device)
    }
    
    if model_key in model_loaders:
        try:
            return model_loaders[model_key]()
        except Exception as e:
            st.error(f"❌ Error loading model {model_key}: {str(e)}")
            logging.error(f"Error loading model {model_key}: {e}")
            sentry_sdk.capture_exception(e)
    return None

# --- Funções de Processamento ---
def validate_audio_file(file: UploadedFile) -> bool:
    valid_extensions = ['.wav', '.mp3', '.flac', '.m4a']
    return any(file.name.lower().endswith(ext) for ext in valid_extensions)

def validate_image_file(file: UploadedFile) -> bool:
    valid_extensions = ['.jpg', '.jpeg', '.png', '.bmp']
    if not any(file.name.lower().endswith(ext) for ext in valid_extensions):
        return False
    try:
        Image.open(file).verify()
        return True
    except Exception:
        return False

def process_audio_file(audio_file):
    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(audio_file.name)[1]) as tmp_file:
            tmp_file.write(audio_file.read())
            tmp_file_path = tmp_file.name
        
        audio_array, sample_rate = librosa.load(tmp_file_path, sr=16000)
        os.unlink(tmp_file_path)
        return audio_array
    except Exception as e:
        st.error(f"❌ Error processing audio: {str(e)}")
        logging.error(f"Audio processing error: {e}")
        return None

def process_image_file(image_file):
    try:
        image = Image.open(image_file)
        if image.mode != 'RGB':
            image = image.convert('RGB')
        return image
    except Exception as e:
        st.error(f"❌ Error processing image: {str(e)}")
        logging.error(f"Image processing error: {e}")
        return None

# --- Result Display Functions ---
def display_results(result, model_key, input_text=None):
    st.markdown("""<div class='result-box'>""", unsafe_allow_html=True)
    
    if model_key == 'summarization':
        st.subheader("📝 Generated Summary")
        if input_text:
            with st.expander("🔍 Original Text"):
                st.write(input_text)
        st.info(result[0]['summary_text'])
    
    elif model_key == 'translation':
        st.subheader("🌍 Translation")
        st.success(result[0]['translation_text'])
    
    elif model_key in ['sentiment_analysis', 'text_classification']:
        st.subheader("📊 Results")
        df = pd.DataFrame(result)
        df['score'] = df['score'].apply(lambda x: f"{x:.2%}")
        
        fig = px.bar(df, x='label', y='score', color='label', 
                     labels={'score': 'Confidence', 'label': 'Category'},
                     title="Probability Distribution")
        st.plotly_chart(fig, use_container_width=True)
        
        for _, row in df.iterrows():
            st.progress(float(row['score'].strip('%'))/100, text=f"{row['label']} ({row['score']})")
    
    elif model_key == 'ner':
        st.subheader("🔍 Recognized Entities")
        entities = []
        for entity in result:
            entities.append({
                "Entity": entity['word'],
                "Type": entity['entity_group'],
                "Confidence": f"{entity['score']:.2%}"
            })
        st.table(pd.DataFrame(entities))
    
    elif model_key == 'text_generation':
        st.subheader("🧠 Generated Text")
        st.write(result[0]['generated_text'])
    
    elif model_key == 'image_classification':
        st.subheader("🏷️ Classification")
        top5 = result[:5]
        labels = [res['label'] for res in top5]
        scores = [res['score'] for res in top5]
        
        fig = px.bar(x=scores, y=labels, orientation='h', 
                     labels={'x': 'Confidence', 'y': 'Category'},
                     title="Top 5 Classifications")
        st.plotly_chart(fig, use_container_width=True)
    
    elif model_key == 'object_detection':
        st.subheader("📦 Detected Objects")
        objects = []
        for obj in result:
            objects.append({
                "Object": obj['label'],
                "Confidence": f"{obj['score']:.2%}"
            })
        st.table(pd.DataFrame(objects))
    
    elif model_key == 'image_segmentation':
        st.subheader("🧩 Segmentation")
        st.image(result[0]['mask'], caption="Segmentation Mask", use_column_width=True)
    
    elif model_key == 'facial_recognition':
        st.subheader("😊 Facial Recognition")
        df = pd.DataFrame(result[:3])
        df['score'] = df['score'].apply(lambda x: f"{x:.2%}")
        st.table(df[['label', 'score']])
    
    elif model_key == 'speech_to_text':
        st.subheader("🔈 Transcription")
        st.success(result['text'])
    
    elif model_key == 'audio_classification':
        st.subheader("🎧 Audio Classification")
        df = pd.DataFrame(result[:3])
        df['score'] = df['score'].apply(lambda x: f"{x:.2%}")
        st.table(df[['label', 'score']])
    
    elif model_key == 'text_to_image':
        st.subheader("🎨 Generated Image")
        st.image(result[0], caption="Image generated from text", use_column_width=True)
    
    st.markdown("</div>", unsafe_allow_html=True)

# --- Use Case Functions ---
def get_use_cases():
    """Returns detailed use cases for each model"""
    return {
        'sentiment_analysis': {
            'title': "📊 Social Media Sentiment Analysis",
            'description': "Automatically identify sentiment (positive, negative, or neutral) in comments about your brand.",
            'example': "Monitor mentions of your product on Twitter to detect customer dissatisfaction and respond quickly.",
            'benefit': "Reduce crisis response time by 40% and increase customer satisfaction by 25%.",
            'demo_input': "I love the new product design! But the battery could last longer.",
            'demo_type': 'text',
            'model_key': 'sentiment_analysis'
        },
        'text_classification': {
            'title': "🗂️ Automated Ticket Classification",
            'description': "Automatically classify support tickets into categories for efficient routing.",
            'example': "Direct technical tickets to engineering and refund requests to finance.",
            'benefit': "Reduce resolution time by 60% and increase routing accuracy by 30%.",
            'demo_input': "My order #12345 hasn't arrived. I need help tracking it.",
            'demo_type': 'text',
            'model_key': 'text_classification'
        },
        'summarization': {
            'title': "📑 Automatic Report Summarization",
            'description': "Condense long reports into easy-to-understand executive summaries.",
            'example': "Turn 50-page quarterly financial reports into 1-page summaries for management.",
            'benefit': "Save 15 hours/month in document analysis and speed up decision making.",
            'demo_input': """The company reported 15% revenue growth last quarter, reaching $120 million. 
            The increase was driven by the launch of the new product X2, contributing $18 million in sales. 
            Operating costs rose 5% due to logistics investments.""",
            'demo_type': 'text',
            'model_key': 'summarization'
        },
        'question_answering': {
            'title': "❓ Intelligent FAQ for Support",
            'description': "Automated Q&A system based on your manuals and documents.",
            'example': "Chatbot that answers technical product questions directly from the user manual.",
            'benefit': "Reduce support call volume by 50%.",
            'demo_input': {
                'context': "Product X has a 24-month warranty. To activate, register on the website with the invoice. Technical support within 48 business hours.",
                'question': "What is the warranty period?"
            },
            'demo_type': 'qa',
            'model_key': 'question_answering'
        },
        'translation': {
            'title': "🌍 Content Translation for Global Expansion",
            'description': "Automatically translate marketing content for new markets.",
            'example': "Translate product descriptions from English to Portuguese for launch in Brazil.",
            'benefit': "Reduce human translation costs by 70% and speed up 10x.",
            'demo_input': "Our premium subscription offers ad-free experience and exclusive content.",
            'demo_type': 'text',
            'model_key': 'translation'
        },
        'ner': {
            'title': "🔍 Information Extraction from Contracts",
            'description': "Automatically identify parties, amounts, and deadlines in legal documents.",
            'example': "Extract party names, amounts, and due dates from supplier contracts.",
            'benefit': "Reduce contract review time by 80% and minimize human errors.",
            'demo_input': "Agreement between ABC Ltd. (Tax ID 12.345/0001-01) and XYZ Inc. (Tax ID 98.765/0001-02) signed on 01/15/2023.",
            'demo_type': 'text',
            'model_key': 'ner'
        },
        'text_generation': {
            'title': "✍️ Marketing Content Generation",
            'description': "Create creative texts for campaigns, ads, and social media.",
            'example': "Generate copywriting variations for email marketing campaigns.",
            'benefit': "Triple content production and reduce copywriting costs.",
            'demo_input': "Write a creative ad for a new smartphone with a 108MP camera:",
            'demo_type': 'text',
            'model_key': 'text_generation'
        },
        'image_classification': {
            'title': "🖼️ Automated Quality Control",
            'description': "Automatically classify products on production lines.",
            'example': "Identify defects in manufactured parts using production line images.",
            'benefit': "Increase defect detection by 25% and reduce rework by 40%.",
            'demo_input': None,
            'demo_type': 'image',
            'model_key': 'image_classification'
        },
        'object_detection': {
            'title': "📦 Smart Inventory Monitoring",
            'description': "Count and identify items on shelves using images.",
            'example': "Monitor stock levels in real-time through shelf images.",
            'benefit': "Reduce stockouts by 30% and increase inventory accuracy by 50%.",
            'demo_input': None,
            'demo_type': 'image',
            'model_key': 'object_detection'
        },
        'image_segmentation': {
            'title': "🏥 Medical Image Analysis",
            'description': "Segment areas of interest in medical imaging exams.",
            'example': "Identify tumors in MRI scans to assist diagnosis.",
            'benefit': "Increase diagnostic accuracy by 20% and reduce analysis time by 30%.",
            'demo_input': None,
            'demo_type': 'image',
            'model_key': 'image_segmentation'
        },
        'facial_recognition': {
            'title': "😊 Customer Satisfaction Analysis",
            'description': "Measure customer emotions at points of sale.",
            'example': "Analyze facial expressions of customers interacting with products in stores.",
            'benefit': "Gain valuable insights into customer preferences and pain points.",
            'demo_input': None,
            'demo_type': 'image',
            'model_key': 'facial_recognition'
        },
        'speech_to_text': {
            'title': "🎙️ Automatic Meeting Transcription",
            'description': "Convert meeting recordings into text minutes automatically.",
            'example': "Transcribe board meetings to create actionable records.",
            'benefit': "Save 5 hours/week on documentation and improve record accuracy.",
            'demo_input': None,
            'demo_type': 'audio',
            'model_key': 'speech_to_text'
        },
        'audio_classification': {
            'title': "📞 Call Center Emotion Analysis",
            'description': "Classify emotional tone in support calls.",
            'example': "Identify frustrated customers for priority escalation.",
            'benefit': "Increase customer satisfaction by 35% and reduce churn by 25%.",
            'demo_input': None,
            'demo_type': 'audio',
            'model_key': 'audio_classification'
        },
        'text_to_image': {
            'title': "🎨 Visual Content Generation for Marketing",
            'description': "Create custom images for campaigns from text descriptions.",
            'example': "Generate promotional banners for social media campaigns.",
            'benefit': "Reduce visual production time by 90% and design costs.",
            'demo_input': "A futuristic landscape with flying cars over a high-tech city at night",
            'demo_type': 'text',
            'model_key': 'text_to_image'
        }
    }

def handle_use_case_demo(models, use_case_key, use_case):
    """
    Executes and displays a demo of a specific use case
    
    Args:
        models: Dictionary of loaded models
        use_case_key: Identifier key of the use case
        use_case: Dictionary with use case settings
    """
    if use_case['demo_input'] is None:
        st.warning("⚠️ Demo requires file upload")
        return
    
    st.markdown("""<div class='demo-box'>""", unsafe_allow_html=True)
    st.subheader("🚀 Live Demo")
    
    try:
        # Load model if needed
        model = models.get(use_case_key)
        if model is None:
            with st.spinner(f"Loading {use_case_key.replace('_', ' ').title()} model..."):
                model = load_model(use_case_key)
                models[use_case_key] = model
        
        if model is None:
            st.error("Failed to load model for demo")
            return
            
        # Process based on demo type
        if use_case['demo_type'] == 'text':
            with st.spinner("Analyzing text..."):
                result = model(use_case['demo_input'])
                display_demo_results(result, use_case_key, use_case['demo_input'])
                
        elif use_case['demo_type'] == 'qa':
            with st.spinner("Processing question..."):
                result = model(
                    question=use_case['demo_input']['question'],
                    context=use_case['demo_input']['context']
                )
                st.success("💡 Answer Found")
                st.markdown(f"**Context:** {use_case['demo_input']['context']}")
                st.markdown(f"**Question:** {use_case['demo_input']['question']}")
                st.markdown(f"**Answer:** `{result['answer']}`")
                st.markdown(f"**Confidence:** {result['score']:.2%}")
                
    except Exception as e:
        st.error(f"❌ Demo failed: {str(e)}")
        logging.error(f"Use case error {use_case_key}: {str(e)}")
    finally:
        st.markdown("</div>", unsafe_allow_html=True)

def display_demo_results(result, model_key, input_text=None):
    """
    Displays the formatted demo results
    
    Args:
        result: Result returned by the model
        model_key: Model identifier
        input_text: Original text (optional)
    """
    with st.expander("🔍 Analysis Details", expanded=True):
        if input_text:
            st.markdown("**Input:**")
            st.write(input_text)
        
        st.markdown("**Results:**")
        
        if model_key in ['sentiment_analysis', 'text_classification']:
            for item in result:
                label = item['label']
                score = item['score']
                st.progress(float(score), 
                           text=f"{label} ({score:.2%})")
        
        elif model_key == 'summarization':
            st.markdown("**Generated Summary:**")
            st.info(result[0]['summary_text'])
        
        elif model_key == 'translation':
            st.markdown("**Translation:**")
            st.success(result[0]['translation_text'])
        
        elif model_key == 'ner':
            st.markdown("**Identified Entities:**")
            for entity in result:
                st.write(f"- {entity['word']} ({entity['entity_group']})")
        
        elif model_key == 'text_generation':
            st.markdown("**Generated Text:**")
            st.write(result[0]['generated_text'])
        
        elif model_key == 'text_to_image':
            st.image(result[0], caption="Image generated by the model")

def render_use_cases(models):
    """
    Renders all use cases in the interface
    
    Args:
        models: Dictionary of loaded models
    """
    st.header("📂 Practical Use Cases")
    st.markdown("Explore real examples of application of our AI models")
    
    use_cases = get_use_cases()
    
    for key, case in use_cases.items():
        with st.container():
            st.subheader(case['title'])
            st.markdown(f"**{case['description']}**")
            
            with st.expander("ℹ️ Use Case Details", expanded=False):
                st.markdown(f"📌 **Practical Example:** {case['example']}")
                st.markdown(f"✅ **Benefit:** {case['benefit']}")
            
            # Interactive demo section
            if case.get('demo_input') is not None:
                if st.button(
                    "▶ Run Demo", 
                    key=f"demo_{key}",
                    type="primary",
                    help=f"Test {case['title'].split(' ')[-1]} with example data"
                ):
                    handle_use_case_demo(models, key, case)
            else:
                # Automatic uploader depending on case type
                if "image" in key or "visual" in key:
                    file = st.file_uploader("📸 Upload an image", type=["jpg", "jpeg", "png"], key=f"upload_{key}")
                    if file:
                        case["demo_input"] = file
                        handle_use_case_demo(models, key, case)

                elif "speech" in key or "audio" in key or "call" in key:
                    file = st.file_uploader("🎙️ Upload an audio file", type=["wav", "mp3", "m4a"], key=f"upload_{key}")
                    if file:
                        case["demo_input"] = file
                        handle_use_case_demo(models, key, case)

                elif "medical" in key:
                    file = st.file_uploader("🏥 Upload a medical image (PNG/JPG/DICOM)", type=["png", "jpg", "jpeg", "dcm"], key=f"upload_{key}")
                    if file:
                        case["demo_input"] = file
                        handle_use_case_demo(models, key, case)

                elif "document" in key or "summarization" in key or "extraction" in key:
                    file = st.file_uploader("📄 Upload a document", type=["pdf", "docx"], key=f"upload_{key}")
                    if file:
                        case["demo_input"] = file
                        handle_use_case_demo(models, key, case)

                else:
                    st.warning("⚠️ This use case requires a specific file input")
            
            st.divider()

# --- Layout and Components ---
def load_branding(username):
    branding = {
        'admin': {'logo': None, 'title': 'Welcome, Administrator!', 'color': '#2c3e50'},
        'cliente': {'logo': None, 'title': 'Welcome, Client!', 'color': '#3498db'},
        'empresa1': {'logo': None, 'title': 'Welcome, Company One!', 'color': '#e74c3c'}
    }
    return branding.get(username, {'logo': None, 'title': 'Welcome!', 'color': '#3498db'})

def feature_card(title, description, icon):
    return f"""
    <div class="card">
        <div class="feature-icon">{icon}</div>
        <h3>{title}</h3>
        <p>{description}</p>
    </div>
    """

def render_login(authenticator):
    st.title("AiiT Services - AI Platform")
    st.markdown("""
    <div style="text-align:center; margin-bottom:30px;">
        <h2 style="color:#2c3e50;">Multimodal Artificial Intelligence</h2>
        <p style="font-size:1.1rem;">Advanced AI solutions for text, image, and audio analysis</p>
    </div>
    """, unsafe_allow_html=True)
    
    cols = st.columns(3)
    with cols[0]:
        st.markdown(feature_card("Smart Text", "Text analysis, classification, and generation", "📝"), unsafe_allow_html=True)
    with cols[1]:
        st.markdown(feature_card("Computer Vision", "Image and object recognition", "🖼️"), unsafe_allow_html=True)
    with cols[2]:
        st.markdown(feature_card("Audio Processing", "Transcription and sentiment analysis in audio", "🎵"), unsafe_allow_html=True)
    
    authenticator.login(location='main')

def render_main_interface(authenticator, config, models):
    name = st.session_state["name"]
    username = st.session_state["username"]
    role = st.session_state.get("role") or config['credentials']['usernames'][username]['role']
    
    branding = load_branding(username)
    st.sidebar.title(f"👤 {name}")
    authenticator.logout("Logout", "sidebar")
    
    if role == "admin":
        admin_panel(authenticator)
    
    st.title(f"{branding['title']}")
    st.markdown(f"<style>.sidebar .sidebar-content {{background: {branding['color']};}}</style>", unsafe_allow_html=True)
    
    model_categories = {
        "📝 Text Processing": [
            ("Sentiment Analysis", "sentiment_analysis"),
            ("Text Classification", "text_classification"),
            ("Text Summarization", "summarization"),
            ("Question Answering", "question_answering"),
            ("Translation (EN→PT)", "translation"),
            ("Named Entity Recognition", "ner"),
            ("Text Generation", "text_generation")
        ],
        "🖼️ Image Processing": [
            ("Image Classification", "image_classification"),
            ("Object Detection", "object_detection"),
            ("Image Segmentation", "image_segmentation"),
            ("Facial Recognition", "facial_recognition")
        ],
        "🎵 Audio Processing": [
            ("Audio Transcription", "speech_to_text"),
            ("Emotion Classification", "audio_classification")
        ],
        "✨ Generative Models": [
            ("Text to Image", "text_to_image")
        ]
    }
    
    tab1, tab2 = st.tabs(["🧠 Explore Models", "💼 Use Cases"])
    
    with tab1:
        selected_category = st.sidebar.selectbox(
            "Category",
            list(model_categories.keys()),
            index=0
        )
        
        selected_model = st.sidebar.selectbox(
            "Model", 
            [name for name, key in model_categories[selected_category]],
            index=0
        )
        
        model_key = next(key for name, key in model_categories[selected_category] if name == selected_model)
        
        st.header(f"{selected_model}")
        st.markdown(f"<div class='card'>", unsafe_allow_html=True)
        
        if model_key not in models:
            models[model_key] = load_model(model_key)
            
        if not models[model_key]:
            st.error("❌ Model not available")
            return
            
        try:
            if model_key in ['sentiment_analysis', 'text_classification', 'summarization', 
                            'translation', 'text_generation', 'ner']:
                handle_text_models(models, model_key, selected_model)
            
            elif model_key == 'question_answering':
                handle_qa_model(models, model_key)
            
            elif model_key in ['image_classification', 'object_detection', 
                              'image_segmentation', 'facial_recognition']:
                handle_image_models(models, model_key, selected_model)
            
            elif model_key in ['speech_to_text', 'audio_classification']:
                handle_audio_models(models, model_key)
            
            elif model_key == 'text_to_image':
                handle_generative_models(models, model_key)
            
        except Exception as e:
            st.error(f"❌ Execution error: {str(e)}")
        
        st.markdown("</div>", unsafe_allow_html=True)
        
    with tab2:
        st.header("💼 Practical Use Cases")
        st.markdown("Explore real applications of our AI models")
        
        use_cases = get_use_cases()
        cols = st.columns(2)
        
        for idx, (key, case) in enumerate(use_cases.items()):
            with cols[idx % 2]:
                st.markdown(f"""
                <div class="use-case-card">
                    <h4>{case['title']}</h4>
                    <p><strong>Description:</strong> {case['description']}</p>
                    <p><strong>Benefit:</strong> {case['benefit']}</p>
                    <button class="btn-primary" onclick="window['useCase_{key}'].click()">View Demo</button>
                </div>
                """, unsafe_allow_html=True)
                
                if st.button("View Demo", key=f"useCase_{key}"):
                    handle_use_case_demo(models, key, case)

# --- Handlers for Model Types ---
def handle_text_models(models, model_key, model_name):
    examples = {
        'sentiment_analysis': "The delivery was super fast, I loved it!",
        'text_classification': "I am dissatisfied with the product",
        'summarization': "Company XYZ reported a 15% growth last quarter...",
        'translation': "Our product ensures high performance",
        'ner': "Microsoft signed a contract with company XYZ in New York",
        'text_generation': "A future where technology connects everyone"
    }
    
    col1, col2 = st.columns([0.85, 0.15])
    with col1:
        input_text = st.text_area(
            f"Enter text for {model_name.lower()}:", 
            height=200,
            placeholder="Paste or type your text here..."
        )
    with col2:
        st.write("")
        st.write("")
        if st.button("📋 Example", use_container_width=True):
            input_text = examples.get(model_key, "")
    
    advanced_params = {}
    if model_key == 'summarization':
        with st.expander("⚙️ Advanced Parameters"):
            advanced_params['max_length'] = st.slider("Max Length", 50, 300, 150)
            advanced_params['min_length'] = st.slider("Min Length", 10, 100, 30)
    
    if model_key == 'text_generation':
        with st.expander("⚙️ Advanced Parameters"):
            advanced_params['max_length'] = st.slider("Text Length", 50, 500, 100)
            advanced_params['temperature'] = st.slider("Creativity", 0.1, 1.0, 0.7)
            advanced_params['num_return_sequences'] = st.slider("Number of outputs", 1, 5, 1)
    
    if st.button(f"🚀 Run {model_name}", type="primary", use_container_width=True):
        if input_text.strip():
            with st.spinner("Processing..."):
                try:
                    result = models[model_key](input_text, **advanced_params)
                    display_demo_results(result, model_key, input_text=input_text)
                except Exception as e:
                    st.error(f"❌ Error processing text: {str(e)}")
        else:
            st.warning("⚠️ Please enter valid text")

def handle_qa_model(models, model_key):
    col1, col2 = st.columns(2)
    with col1:
        context = st.text_area(
            "Context:", 
            height=200,
            placeholder="Text containing the information...",
            value="Product X has a 2-year warranty and can be configured via app in 5 minutes."
        )
    with col2:
        question = st.text_area(
            "Question:", 
            height=150,
            placeholder="Ask your question about the context...",
            value="What is the warranty period of product X?"
        )
    
    with st.expander("⚙️ Advanced Parameters"):
        confidence_threshold = st.slider("Confidence threshold", 0.0, 1.0, 0.5, 0.01)
    
    if st.button("🚀 Run Question and Answer", type="primary", use_container_width=True):
        if context.strip() and question.strip():
            with st.spinner("Searching for answer..."):
                try:
                    result = models[model_key](question=question, context=context)
                    
                    if result['score'] < confidence_threshold:
                        st.warning(f"⚠️ Low confidence in answer ({result['score']:.2%})")
                    
                    st.success("🔍 Answer found:")
                    st.markdown(f"**Question:** {question}")
                    st.markdown(f"**Answer:** {result['answer']}")
                    st.markdown(f"**Confidence:** {result['score']:.2%}")
                except Exception as e:
                    st.error(f"❌ Error processing Q&A: {str(e)}")
        else:
            st.warning("⚠️ Please fill in context and question")

def handle_image_models(models, model_key, model_name):
    uploaded_file = st.file_uploader(
        "Upload an image", 
        type=["jpg", "png", "jpeg", "bmp"],
        help="Supported formats: JPG, PNG, JPEG, BMP"
    )
    
    if uploaded_file:
        if not validate_image_file(uploaded_file):
            st.error("⚠️ Invalid format or corrupted file")
            return
            
        col1, col2 = st.columns(2)
        
        with col1:
            st.subheader("🖼️ Original Image")
            image = process_image_file(uploaded_file)
            if image:
                st.image(image, use_column_width=True)
        
        with col2:
            st.subheader("📊 Results")
            if st.button(f"🚀 Run {model_name}", type="primary", use_container_width=True):
                if image:
                    with st.spinner("Analyzing image..."):
                        try:
                            result = models[model_key](image)
                            display_demo_results(result, model_key)
                        except Exception as e:
                            st.error(f"❌ Error processing image: {str(e)}")

def handle_audio_models(models, model_key):
    model_name = "Audio Transcription" if model_key == 'speech_to_text' else "Audio Classification"
    
    uploaded_file = st.file_uploader(
        f"Upload an audio file for {model_name}", 
        type=["wav", "mp3", "flac", "m4a"],
        help="Supported formats: WAV, MP3, FLAC, M4A"
    )
    
    if uploaded_file:
        if not validate_audio_file(uploaded_file):
            st.error("⚠️ Unsupported file format")
            return
            
        st.audio(uploaded_file, format="audio/wav")
        
        if st.button(f"🚀 Run {model_name}", type="primary", use_container_width=True):
            with st.spinner("Processing audio..."):
                try:
                    audio_array = process_audio_file(uploaded_file)
                    if audio_array is not None:
                        result = models[model_key](audio_array)
                        display_demo_results(result, model_key)
                except Exception as e:
                    st.error(f"❌ Error processing audio: {str(e)}")

def handle_generative_models(models, model_key):
    prompt = st.text_area(
        "Image description:",
        height=150,
        placeholder="Describe the image you want to generate...",
        value="A tropical landscape at sunset"
    )
    
    with st.expander("⚙️ Advanced Parameters"):
        cols = st.columns(2)
        with cols[0]:
            width = st.slider("Width", 256, 1024, 512, 64)
        with cols[1]:
            height = st.slider("Height", 256, 1024, 512, 64)
        num_images = st.slider("Number of images", 1, 4, 1)
        guidance_scale = st.slider("Guidance scale", 1.0, 20.0, 7.5)
    
    if st.button("🚀 Generate Image", type="primary", use_container_width=True):
        if prompt.strip():
            with st.spinner("Creating image..."):
                try:
                    result = models[model_key](
                        prompt,
                        height=height,
                        width=width,
                        num_images_per_prompt=num_images,
                        guidance_scale=guidance_scale
                    ).images
                    display_demo_results(result, model_key)
                except Exception as e:
                    st.error(f"❌ Error generating image: {str(e)}")
        else:
            st.warning("⚠️ Please enter a description for the image")

# --- Main function ---
def main():
    init_db()
    config = load_users_from_db()
    
    if not config:
        st.error("⚠️ Authentication system not configured")
        return
        
    authenticator = stauth.Authenticate(
        config['credentials'],
        config['cookie']['name'],
        config['cookie']['key'],
        config['cookie']['expiry_days']
    )
    
    if 'authentication_status' not in st.session_state:
        st.session_state.authentication_status = None
        
    if st.session_state.get("authentication_status"):
        models = st.session_state.setdefault("models", {})
        render_main_interface(authenticator, config, models)
        
        # Moved inside the authenticated block
        tab1, tab2 = st.tabs(["Models", "Use Cases"])
        with tab2:
            render_use_cases(models)  # Render all use cases
    else:
        render_login(authenticator)
        
        if st.session_state.get("authentication_status") is False:
            st.error("❌ Invalid credentials")
        elif st.session_state.get("authentication_status") is None:
            st.info("🔑 Please log in to access the platform")

if __name__ == "__main__":
    main()