import os
import io
import tempfile
import logging
from datetime import datetime, timedelta
from typing import Dict, Any

import torch
import librosa
import numpy as np
import pandas as pd
import plotly.express as px
from PIL import Image

import streamlit as st
import streamlit_authenticator as stauth
from streamlit.runtime.uploaded_file_manager import UploadedFile

from transformers import pipeline
from diffusers import StableDiffusionPipeline
import sentry_sdk
from supabase import create_client, Client
import types

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

# Monkey-patch compat (keeps old code working if needed)
st.context = types.SimpleNamespace(cookies={})

# -----------------------
# Estilos
# -----------------------
st.markdown(
    """
<style>
    :root {
    --primary: #7f5af0;
    --secondary: #2cb67d;
    --accent: #f25f4c;
    --light: #fdfdfd;
    --dark: #16161a;
    --success: #2cb67d;
    --info: #3d5afe;
}
.stApp {
    background-image: linear-gradient(120deg, #0f0c29, #302b63, #24243e);
}
.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.2rem;
    margin-bottom: 10px;
    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;
}
.result-box {
    background: linear-gradient(135deg, #e0f7fa, #b2ebf2);
    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;
}
</style>
""",
    unsafe_allow_html=True,
)

# -----------------------
# Logging & Sentry
# -----------------------
sentry_sdk.init(os.getenv("SENTRY_DSN", ""), traces_sample_rate=1.0)
os.makedirs("private", exist_ok=True)
logging.basicConfig(
    filename="private/app_errors.log",
    level=logging.ERROR,
    format="%(asctime)s - %(levelname)s - %(message)s",
)

# -----------------------
# Supabase (optional)
# -----------------------
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


# -----------------------
# DB helper functions
# -----------------------
def connect_db():
    """Connects to Supabase (if configured)."""
    if not supabase:
        return None
    return supabase


def init_db():
    """Quick check to ensure table exists (non-blocking)."""
    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 and returns credentials structure compatible with streamlit_authenticator.
    If Supabase is not configured or fails, returns a default admin user for local dev.
    """
    try:
        sb = connect_db()
        if not sb:
            # fallback default (useful for dev)
            default_password = stauth.Hasher(["admin"]).generate()[0]
            return {
                "credentials": {
                    "usernames": {
                        "admin": {
                            "name": "Administrator",
                            "password": default_password,
                            "role": "admin",
                        }
                    }
                },
                "cookie": {"name": "ai_app_cookie", "key": "random_key_123", "expiry_days": 30},
                "preauthorized": {"emails": []},
            }

        resp = sb.table("users").select("username, name, password, role").execute()
        rows = getattr(resp, "data", []) or []

        users = {}
        for row in rows:
            # Expecting password already hashed (using stauth.Hasher)
            users[row["username"]] = {
                "name": row.get("name", row["username"]),
                "password": row.get("password", ""),
                "role": row.get("role", "user"),
            }

        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}")
        sentry_sdk.capture_exception(e)
        st.error("❌ Error loading authentication configuration")
        return None


# -----------------------
# Admin panel
# -----------------------
def admin_panel(authenticator):
    sb = connect_db()
    role = st.session_state.get("role")
    if not sb or 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)}")
                    logging.error(f"Add user error: {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)}")
                    logging.error(f"Remove user error: {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)}")
                    logging.error(f"Update role error: {e}")

    st.sidebar.subheader("👥 Registered Users")
    try:
        resp = sb.table("users").select("username,name,role").execute()
        for row in getattr(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)}")
        logging.error(f"List users error: {e}")


# -----------------------
# Model caching/loader
# -----------------------
@st.cache_resource(show_spinner="Loading model...")
def load_model(model_key: str):
    """
    Loads and returns a Hugging Face pipeline or Diffusers pipeline.
    Uses device index (0 for first GPU or -1 for CPU).
    """
    device_index = 0 if torch.cuda.is_available() else -1
    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_index),
        "text_classification": lambda: pipeline("text-classification", model="distilbert-base-uncased-finetuned-sst-2-english", device=device_index),
        "summarization": lambda: pipeline("summarization", model="facebook/bart-large-cnn", device=device_index),
        "question_answering": lambda: pipeline("question-answering", model="deepset/roberta-base-squad2", device=device_index),
        "translation": lambda: pipeline("translation", model="Helsinki-NLP/opus-mt-tc-big-en-pt", device=device_index),
        "text_generation": lambda: pipeline("text-generation", model="gpt2", device=device_index, pad_token_id=50256),
        "ner": lambda: pipeline("ner", model="dbmdz/bert-large-cased-finetuned-conll03-english", device=device_index, aggregation_strategy="simple"),
        "image_classification": lambda: pipeline("image-classification", model="google/vit-base-patch16-224", device=device_index),
        "object_detection": lambda: pipeline("object-detection", model="facebook/detr-resnet-50", device=device_index),
        "image_segmentation": lambda: pipeline("image-segmentation", model="facebook/detr-resnet-50-panoptic", device=device_index),
        "facial_recognition": lambda: pipeline("image-classification", model="mo-thecreator/vit-Facial-Expression-Recognition", device=device_index),
        "speech_to_text": lambda: pipeline("automatic-speech-recognition", model="openai/whisper-base", device=device_index),
        "audio_classification": lambda: pipeline("audio-classification", model="superb/hubert-base-superb-er", device=device_index),
        "text_to_image": lambda: StableDiffusionPipeline.from_pretrained(
            "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float32, cache_dir=cache_dir
        ).to("cuda" if torch.cuda.is_available() else "cpu"),
    }

    loader = model_loaders.get(model_key)
    if not loader:
        return None

    try:
        return loader()
    except Exception as e:
        logging.error(f"Error loading model {model_key}: {e}")
        sentry_sdk.capture_exception(e)
        st.error(f"❌ Error loading model '{model_key}': {str(e)}")
        return None


# -----------------------
# File validation & processing
# -----------------------
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: UploadedFile):
    """
    Saves uploaded audio to a temporary file and returns the path and sample array.
    (ASR pipelines usually accept file path or raw array depending on implementation.)
    """
    try:
        suffix = os.path.splitext(audio_file.name)[1]
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) 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)
        return tmp_file_path, audio_array
    except Exception as e:
        st.error(f"❌ Error processing audio: {str(e)}")
        logging.error(f"Audio processing error: {e}")
        sentry_sdk.capture_exception(e)
        return None, None


def process_image_file(image_file: UploadedFile):
    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}")
        sentry_sdk.capture_exception(e)
        return None


# -----------------------
# Result display (single canonical function)
# -----------------------
def display_results(result: Any, model_key: str, input_text: str = 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)
        if isinstance(result, list) and result:
            st.info(result[0].get("summary_text", ""))
        else:
            st.info(result)

    elif model_key == "translation":
        st.subheader("🌍 Translation")
        if isinstance(result, list) and result:
            st.success(result[0].get("translation_text", result[0].get("generated_text", "")))
        else:
            st.success(result)

    elif model_key in ["sentiment_analysis", "text_classification"]:
        st.subheader("📊 Results")
        df = pd.DataFrame(result)
        if "score" in df.columns:
            df["score_pct"] = df["score"].apply(lambda x: float(x) * 100)
            fig = px.bar(df, x="label", y="score_pct", color="label", labels={"score_pct": "Confidence (%)", "label": "Category"}, title="Probability Distribution")
            st.plotly_chart(fig, use_container_width=True)
            for _, row in df.iterrows():
                st.write(f"- **{row['label']}** — {row['score']:.2%}")

    elif model_key == "ner":
        st.subheader("🔍 Recognized Entities")
        entities = []
        for entity in result:
            entities.append({"Entity": entity.get("word", ""), "Type": entity.get("entity_group", entity.get("entity")), "Confidence": f"{entity.get('score', 0):.2%}"})
        st.table(pd.DataFrame(entities))

    elif model_key == "text_generation":
        st.subheader("🧠 Generated Text")
        if isinstance(result, list) and result:
            st.write(result[0].get("generated_text", ""))
        else:
            st.write(result)

    elif model_key == "image_classification":
        st.subheader("🏷️ Classification")
        top5 = result[:5] if isinstance(result, list) else []
        labels = [res["label"] for res in top5]
        scores = [res["score"] * 100 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 = [{"Object": obj.get("label", ""), "Confidence": f"{obj.get('score', 0):.2%}"} for obj in result]
        st.table(pd.DataFrame(objects))

    elif model_key == "image_segmentation":
        st.subheader("🧩 Segmentation")
        # result format varies by pipeline - attempt to show masks/images
        if isinstance(result, dict) and "segmentation" in result:
            st.image(result["segmentation"], caption="Segmentation Mask", use_column_width=True)
        elif isinstance(result, list) and result and isinstance(result[0], dict) and "mask" in result[0]:
            st.image(result[0]["mask"], caption="Segmentation Mask", use_column_width=True)
        else:
            st.write(result)

    elif model_key == "facial_recognition":
        st.subheader("😊 Facial Recognition")
        df = pd.DataFrame(result[:3])
        if "score" in df.columns:
            df["score_pct"] = df["score"].apply(lambda x: f"{x:.2%}")
            st.table(df[["label", "score_pct"]])

    elif model_key == "speech_to_text":
        st.subheader("🔈 Transcription")
        # pipeline returns a dict with 'text' sometimes or list - handle both
        if isinstance(result, dict) and "text" in result:
            st.success(result["text"])
        elif isinstance(result, str):
            st.success(result)
        else:
            st.success(str(result))

    elif model_key == "audio_classification":
        st.subheader("🎧 Audio Classification")
        df = pd.DataFrame(result[:3])
        if "score" in df.columns:
            df["score_pct"] = df["score"].apply(lambda x: f"{x:.2%}")
            st.table(df[["label", "score_pct"]])

    elif model_key == "text_to_image":
        st.subheader("🎨 Generated Image")
        # Expecting a list of PIL images or np arrays
        if isinstance(result, (list, tuple)):
            for idx, img in enumerate(result):
                st.image(img, caption=f"Image {idx + 1}", use_column_width=True)
        else:
            st.image(result, use_column_width=True)

    else:
        st.write(result)

    st.markdown("</div>", unsafe_allow_html=True)


# -----------------------
# UI Helpers
# -----------------------
def load_branding(username: str):
    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: str, description: str, icon: str) -> str:
    return f"""
    <div class="card">
        <div class="feature-icon">{icon}</div>
        <h3>{title}</h3>
        <p>{description}</p>
    </div>
    """


def get_use_cases():
    return {
        "chatbot_support": {
            "title": "Support Chatbot",
            "description": "Automated technical support assistant for first-level tickets.",
            "benefit": "Faster response and ticket triage.",
        },
        "visual_audit": {
            "title": "Visual Audit",
            "description": "Image-based quality inspection and anomaly detection.",
            "benefit": "Scale visual QA with automation.",
        },
    }


# -----------------------
# Model Handlers (UI)
# -----------------------
def handle_text_models(models: Dict[str, Any], model_key: str, model_name: str):
    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, mainly due to increased adoption of product A and operational efficiencies.",
        "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 and input_text.strip():
            with st.spinner("Processing..."):
                try:
                    result = models[model_key](input_text, **advanced_params) if advanced_params else models[model_key](input_text)
                    display_results(result, model_key, input_text=input_text)
                except Exception as e:
                    st.error(f"❌ Error processing text: {str(e)}")
                    logging.error(f"Text model error ({model_key}): {e}")
                    sentry_sdk.capture_exception(e)
        else:
            st.warning("⚠️ Please enter valid text")


def handle_qa_model(models: Dict[str, Any], model_key: str):
    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)
                    score = result.get("score", 0) if isinstance(result, dict) else 0
                    if score < confidence_threshold:
                        st.warning(f"⚠️ Low confidence in answer ({score:.2%})")
                    st.success("🔍 Answer found:")
                    st.markdown(f"**Question:** {question}")
                    st.markdown(f"**Answer:** {result.get('answer', str(result))}")
                    st.markdown(f"**Confidence:** {score:.2%}")
                except Exception as e:
                    st.error(f"❌ Error processing Q&A: {str(e)}")
                    logging.error(f"QA model error: {e}")
                    sentry_sdk.capture_exception(e)
        else:
            st.warning("⚠️ Please fill in context and question")


def handle_image_models(models: Dict[str, Any], model_key: str, model_name: str):
    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:
                            # Some pipelines accept PIL image directly
                            result = models[model_key](image)
                            display_results(result, model_key)
                        except Exception as e:
                            st.error(f"❌ Error processing image: {str(e)}")
                            logging.error(f"Image model error ({model_key}): {e}")
                            sentry_sdk.capture_exception(e)


def handle_audio_models(models: Dict[str, Any], model_key: str):
    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)
        if st.button(f"🚀 Run {model_name}", type="primary", use_container_width=True):
            with st.spinner("Processing audio..."):
                try:
                    tmp_path, audio_array = process_audio_file(uploaded_file)
                    if tmp_path:
                        # Many ASR pipelines accept a path
                        result = models[model_key](tmp_path)
                        display_results(result, model_key)
                        try:
                            os.unlink(tmp_path)
                        except Exception:
                            pass
                    else:
                        st.error("❌ Error preparing audio for model")
                except Exception as e:
                    st.error(f"❌ Error processing audio: {str(e)}")
                    logging.error(f"Audio model error ({model_key}): {e}")
                    sentry_sdk.capture_exception(e)


def handle_generative_models(models: Dict[str, Any], model_key: str):
    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 and prompt.strip():
            with st.spinner("Creating image..."):
                try:
                    pipe: StableDiffusionPipeline = models[model_key]
                    images = []
                    for _ in range(num_images):
                        out = pipe(prompt, height=height, width=width, guidance_scale=guidance_scale)
                        img = out.images[0]
                        images.append(img)
                    display_results(images, model_key)
                except Exception as e:
                    st.error(f"❌ Error generating image: {str(e)}")
                    logging.error(f"Text-to-image error: {e}")
                    sentry_sdk.capture_exception(e)
        else:
            st.warning("⚠️ Please enter a description for the image")


def handle_use_case_demo(models: Dict[str, Any], key: str, case: Dict[str, str]):
    st.info(f"Demo for: {case['title']}")
    st.write(case["description"])
    st.write("Benefit:", case["benefit"])
    # Minimal demo: show which models would be used
    if key == "chatbot_support":
        st.write("- Uses: question_answering, text_generation, ner, translation")
    elif key == "visual_audit":
        st.write("- Uses: image_classification, object_detection, image_segmentation")


# -----------------------
# Main UI rendering
# -----------------------
def render_login(authenticator):
    st.title("AiiT Services - AI Platform")
    st.markdown(
        """
    <div style="text-align:center; margin-bottom:30px;">
        <h2 style="color:#fdfdfd;">Multimodal Artificial Intelligence</h2>
        <p style="font-size:1.1rem; color:#e6e6e6;">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)

    # capture authenticator output and set session state
    name, authentication_status, username = authenticator.login("Login", location="main")

    st.session_state["name"] = name
    st.session_state["authentication_status"] = authentication_status
    st.session_state["username"] = username


def render_main_interface(authenticator, config, models: Dict[str, Any]):
    name = st.session_state.get("name", "User")
    username = st.session_state.get("username", "unknown")
    # Determine role from session or config
    role = st.session_state.get("role") or config["credentials"]["usernames"].get(username, {}).get("role", "user")
    st.session_state["role"] = 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("<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")
        else:
            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)}")
                logging.error(f"Handler execution error ({model_key}): {e}")
                sentry_sdk.capture_exception(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>
                </div>
                """,
                    unsafe_allow_html=True,
                )
                if st.button("View Demo", key=f"useCase_{key}"):
                    handle_use_case_demo(models, key, case)


# -----------------------
# Main
# -----------------------
# -----------------------
# Main App
# -----------------------
def main():
    init_db()

    # Carregar config de utilizadores do Supabase (ou fallback local)
    config = load_users_from_db()
    if not config:
        st.stop()

    # Instanciar autenticação (sem o preauthorized!)
    authenticator = stauth.Authenticate(
        config["credentials"],
        config["cookie"]["name"],
        config["cookie"]["key"],
        config["cookie"]["expiry_days"],
    )

    # Login
    name, auth_status, username = authenticator.login("Login", "main")

    if auth_status:
        st.sidebar.success(f"✅ Logged in as {name}")
        st.session_state["role"] = config["credentials"]["usernames"][username].get("role", "user")

        # Branding
        branding = load_branding(username)
        st.title(branding["title"])

        # Painel admin (apenas para admins)
        if st.session_state["role"] == "admin":
            st.sidebar.subheader("⚙️ Admin Panel")
            with st.sidebar.expander("Manage Users"):
                # Criar novo utilizador
                if st.checkbox("Add new user"):
                    try:
                        email_of_registered_user, username_of_registered_user, name_of_registered_user = authenticator.register_user(
                            preauthorization=config.get("preauthorized", [])  # <<-- só aqui o preauth
                        )
                        st.success(f"✅ User {username_of_registered_user} registered successfully")
                    except Exception as e:
                        st.error(f"❌ Error registering user: {e}")

                # Reset password
                if st.checkbox("Reset password"):
                    try:
                        if authenticator.reset_password(username):
                            st.success("✅ Password reset successful")
                    except Exception as e:
                        st.error(f"❌ Error resetting password: {e}")

                # Update user details
                if st.checkbox("Update user details"):
                    try:
                        if authenticator.update_user_details(username):
                            st.success("✅ User details updated")
                    except Exception as e:
                        st.error(f"❌ Error updating user: {e}")

        # ----------------------------
        # Modelos carregados on-demand
        # ----------------------------
        models = {
            key: load_model(key)
            for key in [
                "sentiment_analysis", "text_classification", "summarization",
                "translation", "ner", "text_generation", "question_answering",
                "image_classification", "object_detection", "image_segmentation",
                "facial_recognition", "speech_to_text", "audio_classification",
                "text_to_image"
            ]
        }

        # Sidebar de navegação
        st.sidebar.title("📂 Features")
        choice = st.sidebar.radio(
            "Select a module",
            [
                "📖 Text Analysis",
                "❓ Question Answering",
                "🖼️ Image Analysis",
                "🎧 Audio Processing",
                "🎨 Generative AI"
            ]
        )

        if choice == "📖 Text Analysis":
            task = st.selectbox("Select task", [
                "Sentiment Analysis", "Text Classification", "Summarization",
                "Translation", "Named Entity Recognition", "Text Generation"
            ])
            mapping = {
                "Sentiment Analysis": "sentiment_analysis",
                "Text Classification": "text_classification",
                "Summarization": "summarization",
                "Translation": "translation",
                "Named Entity Recognition": "ner",
                "Text Generation": "text_generation"
            }
            handle_text_models(models, mapping[task], task)

        elif choice == "❓ Question Answering":
            handle_qa_model(models, "question_answering")

        elif choice == "🖼️ Image Analysis":
            task = st.selectbox("Select task", [
                "Image Classification", "Object Detection", "Image Segmentation", "Facial Recognition"
            ])
            mapping = {
                "Image Classification": "image_classification",
                "Object Detection": "object_detection",
                "Image Segmentation": "image_segmentation",
                "Facial Recognition": "facial_recognition"
            }
            handle_image_models(models, mapping[task], task)

        elif choice == "🎧 Audio Processing":
            task = st.selectbox("Select task", ["Speech to Text", "Audio Classification"])
            mapping = {
                "Speech to Text": "speech_to_text",
                "Audio Classification": "audio_classification"
            }
            handle_audio_models(models, mapping[task])

        elif choice == "🎨 Generative AI":
            handle_generative_models(models, "text_to_image")

        # Logout
        authenticator.logout("Logout", "sidebar")

    elif auth_status is False:
        st.error("❌ Invalid username or password")
    else:
        st.warning("⚠️ Please login")

if __name__ == "__main__":
    main()