import os

import streamlit as st
from esperanto import AIFactory

from open_notebook.domain.models import DefaultModels, Model, model_manager
from pages.components.model_selector import model_selector
from pages.stream_app.utils import setup_page

setup_page("🤖 Models", only_check_mandatory_models=False, stop_on_model_error=False)


st.title("🤖 Models")

provider_status = {}

model_types = [
    # "vision",
    "language",
    "embedding",
    "text_to_speech",
    "speech_to_text",
]


def check_available_providers():
    provider_status["ollama"] = os.environ.get("OLLAMA_API_BASE") is not None
    provider_status["openai"] = os.environ.get("OPENAI_API_KEY") is not None
    provider_status["groq"] = os.environ.get("GROQ_API_KEY") is not None
    provider_status["xai"] = os.environ.get("XAI_API_KEY") is not None
    provider_status["vertexai"] = (
        os.environ.get("VERTEX_PROJECT") is not None
        and os.environ.get("VERTEX_LOCATION") is not None
        and os.environ.get("GOOGLE_APPLICATION_CREDENTIALS") is not None
    )
    provider_status["gemini"] = os.environ.get("GOOGLE_API_KEY") is not None
    provider_status["openrouter"] = (
        os.environ.get("OPENROUTER_API_KEY") is not None
        and os.environ.get("OPENAI_API_KEY") is not None
        and os.environ.get("OPENROUTER_BASE_URL") is not None
    )
    provider_status["anthropic"] = os.environ.get("ANTHROPIC_API_KEY") is not None
    provider_status["elevenlabs"] = os.environ.get("ELEVENLABS_API_KEY") is not None
    provider_status["voyage"] = os.environ.get("VORAGE_API_KEY") is not None
    provider_status["azure"] = (
        os.environ.get("AZURE_OPENAI_API_KEY") is not None
        and os.environ.get("AZURE_OPENAI_ENDPOINT") is not None
        and os.environ.get("AZURE_OPENAI_DEPLOYMENT_NAME") is not None
        and os.environ.get("AZURE_OPENAI_API_VERSION") is not None
    )
    provider_status["mistral"] = os.environ.get("MISTRAL_API_KEY") is not None
    provider_status["deepseek"] = os.environ.get("DEEPSEEK_API_KEY") is not None

    available_providers = [k for k, v in provider_status.items() if v]
    unavailable_providers = [k for k, v in provider_status.items() if not v]

    return available_providers, unavailable_providers


default_models = DefaultModels()
all_models = Model.get_all()
esperanto_available_providers = AIFactory.get_available_providers()


st.subheader("Provider Availability")
st.markdown(
    "Below, you'll find all AI providers supported and their current availability status. To enable more providers, you need to setup some of their ENV Variables. Please check [the documentation](https://github.com/lfnovo/open-notebook/blob/main/docs/models.md) for instructions on how to do so."
)
available_providers, unavailable_providers = check_available_providers()
with st.expander("Available Providers"):
    st.write(available_providers)
with st.expander("Unavailable Providers"):
    st.write(unavailable_providers)

st.divider()


# Helper function to add model with auto-save
def add_model_form(model_type, container_key):
    available_providers = esperanto_available_providers.get(model_type, [])
    # Sort providers alphabetically for easier navigation
    available_providers.sort()

    # Remove perplexity from available_providers if it exists
    if "perplexity" in available_providers:
        available_providers.remove("perplexity")

    if not available_providers:
        st.info(f"No providers available for {model_type}")
        return

    st.markdown("**Add New Model**")

    with st.form(key=f"add_{model_type}_{container_key}"):
        provider = st.selectbox(
            "Provider",
            available_providers,
            key=f"provider_{model_type}_{container_key}",
        )

        model_name = st.text_input(
            "Model Name",
            key=f"name_{model_type}_{container_key}",
            help="gpt-4o-mini, claude, gemini, llama3, etc. For azure, use the deployment_name as the model_name",
        )

        if st.form_submit_button("Add Model"):
            if model_name:
                model = Model(name=model_name, provider=provider, type=model_type)
                model.save()
                st.success("Model added!")
                st.rerun()


# Helper function to handle default model selection with auto-save
def handle_default_selection(
    label, key, current_value, help_text, model_type, caption=None
):
    selected_model = model_selector(
        label,
        key,
        selected_id=current_value,
        help=help_text,
        model_type=model_type,
    )
    # Auto-save when selection changes
    if selected_model and (not current_value or selected_model.id != current_value):
        setattr(default_models, key, selected_model.id)
        default_models.update()
        model_manager.refresh_defaults()
        st.toast(f"Default {model_type} model set to {selected_model.name}")
    elif not selected_model and current_value:
        setattr(default_models, key, None)
        default_models.update()
        model_manager.refresh_defaults()
        st.toast(f"Default {model_type} model removed")

    if caption:
        st.caption(caption)
    return selected_model


# Group models by type
models_by_type = {
    "language": [],
    "embedding": [],
    "text_to_speech": [],
    "speech_to_text": [],
}

for model in all_models:
    if model.type in models_by_type:
        models_by_type[model.type].append(model)


st.markdown("""
**Model Management Guide:** For optimal performance, refer to [Which model to choose?](https://github.com/lfnovo/open-notebook/blob/main/docs/models.md) 
You can test models in the [Transformations](Transformations) page.
""")

# Language Models Section
st.subheader("🗣️ Language Models")
with st.container(border=True):
    col1, col2 = st.columns([2, 1])

    with col1:
        st.markdown("**Configured Models**")
        language_models = models_by_type["language"]
        if language_models:
            for model in language_models:
                subcol1, subcol2 = st.columns([4, 1])
                with subcol1:
                    st.markdown(f"• {model.provider}/{model.name}")
                with subcol2:
                    if st.button(
                        "🗑️", key=f"delete_lang_{model.id}", help="Delete model"
                    ):
                        model.delete()
                        st.rerun()
        else:
            st.info("No language models configured")

    with col2:
        add_model_form("language", "main")

    st.markdown("**Default Model Assignments**")
    col1, col2 = st.columns(2)

    with col1:
        handle_default_selection(
            "Chat Model",
            "default_chat_model",
            default_models.default_chat_model,
            "Used for chat conversations",
            "language",
            "Pick the one that vibes with you.",
        )

        handle_default_selection(
            "Tools Model",
            "default_tools_model",
            default_models.default_tools_model,
            "Used for calling tools - use OpenAI or Anthropic",
            "language",
            "Recommended: gpt-4o, claude, qwen3, etc.",
        )

    with col2:
        handle_default_selection(
            "Transformation Model",
            "default_transformation_model",
            default_models.default_transformation_model,
            "Used for summaries, insights, etc.",
            "language",
            "Can use cheaper models: gpt-4o-mini, llama3, gemma3, etc.",
        )

        handle_default_selection(
            "Large Context Model",
            "large_context_model",
            default_models.large_context_model,
            "Used for large context processing",
            "language",
            "Recommended: Gemini models",
        )

# Embedding Models Section
st.subheader("🔍 Embedding Models")
with st.container(border=True):
    col1, col2 = st.columns([2, 1])

    with col1:
        st.markdown("**Configured Models**")
        embedding_models = models_by_type["embedding"]
        if embedding_models:
            for model in embedding_models:
                subcol1, subcol2 = st.columns([4, 1])
                with subcol1:
                    st.markdown(f"• {model.provider}/{model.name}")
                with subcol2:
                    if st.button(
                        "🗑️", key=f"delete_emb_{model.id}", help="Delete model"
                    ):
                        model.delete()
                        st.rerun()
        else:
            st.info("No embedding models configured")

        handle_default_selection(
            "Default Embedding Model",
            "default_embedding_model",
            default_models.default_embedding_model,
            "Used for semantic search and embeddings",
            "embedding",
        )
        st.warning("⚠️ Changing embedding models requires regenerating all embeddings")

    with col2:
        add_model_form("embedding", "main")

# Text-to-Speech Models Section
st.subheader("🎙️ Text-to-Speech Models")
with st.container(border=True):
    col1, col2 = st.columns([2, 1])

    with col1:
        st.markdown("**Configured Models**")
        tts_models = models_by_type["text_to_speech"]
        if tts_models:
            for model in tts_models:
                subcol1, subcol2 = st.columns([4, 1])
                with subcol1:
                    st.markdown(f"• {model.provider}/{model.name}")
                with subcol2:
                    if st.button(
                        "🗑️", key=f"delete_tts_{model.id}", help="Delete model"
                    ):
                        model.delete()
                        st.rerun()
        else:
            st.info("No text-to-speech models configured")

        handle_default_selection(
            "Default TTS Model",
            "default_text_to_speech_model",
            default_models.default_text_to_speech_model,
            "Used for podcasts and audio generation",
            "text_to_speech",
            "Can be overridden per podcast",
        )

    with col2:
        add_model_form("text_to_speech", "main")

# Speech-to-Text Models Section
st.subheader("🎤 Speech-to-Text Models")
with st.container(border=True):
    col1, col2 = st.columns([2, 1])

    with col1:
        st.markdown("**Configured Models**")
        stt_models = models_by_type["speech_to_text"]
        if stt_models:
            for model in stt_models:
                subcol1, subcol2 = st.columns([4, 1])
                with subcol1:
                    st.markdown(f"• {model.provider}/{model.name}")
                with subcol2:
                    if st.button(
                        "🗑️", key=f"delete_stt_{model.id}", help="Delete model"
                    ):
                        model.delete()
                        st.rerun()
        else:
            st.info("No speech-to-text models configured")

        handle_default_selection(
            "Default STT Model",
            "default_speech_to_text_model",
            default_models.default_speech_to_text_model,
            "Used for audio transcriptions",
            "speech_to_text",
        )

    with col2:
        add_model_form("speech_to_text", "main")