app.py

"""Streamlit app for ProfillyBot."""

import logging
import os
import sys
from datetime import datetime
from pathlib import Path

import streamlit as st

# Add src to path
sys.path.insert(0, str(Path(__file__).parent))

from src.config_loader import get_config
from src.llm_handler import (
    detect_default_provider,
    get_available_groq_models,
    get_available_ollama_models,
    get_default_groq_model,
    get_default_ollama_model,
    get_llm_handler,
)
from src.rag_pipeline import RAGPipeline

# Configure logging
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

# Log LangSmith tracing status
if os.getenv("LANGCHAIN_TRACING_V2", "").lower() == "true":
    logger.info(
        f"LangSmith tracing enabled for project: {os.getenv('LANGCHAIN_PROJECT', 'profillybot')}"
    )


def init_session_state():
    """Initialize Streamlit session state."""
    if "messages" not in st.session_state:
        st.session_state.messages = []

    if "rag_pipeline" not in st.session_state:
        st.session_state.rag_pipeline = None

    if "config" not in st.session_state:
        st.session_state.config = None

    if "pending_question" not in st.session_state:
        st.session_state.pending_question = None

    # LLM Provider settings
    if "llm_provider" not in st.session_state:
        st.session_state.llm_provider = detect_default_provider()

    if "groq_api_key" not in st.session_state:
        st.session_state.groq_api_key = ""

    if "groq_model" not in st.session_state:
        st.session_state.groq_model = get_default_groq_model()

    if "ollama_model" not in st.session_state:
        st.session_state.ollama_model = get_default_ollama_model()


def load_config():
    """Load configuration."""
    if st.session_state.config is None:
        try:
            st.session_state.config = get_config()
            logger.info("Configuration loaded successfully")
        except Exception as e:
            st.error(f"Error loading configuration: {e}")
            st.stop()


def load_rag_pipeline(force_reload: bool = False):
    """Load RAG pipeline with current provider settings.

    Args:
        force_reload: Force reload even if pipeline exists
    """
    if st.session_state.rag_pipeline is None or force_reload:
        with st.spinner("Loading RAG pipeline... This may take a moment."):
            try:
                # Get provider settings from session state
                provider = st.session_state.llm_provider
                api_key = st.session_state.groq_api_key or None
                model = None

                if provider == "groq":
                    model = st.session_state.groq_model
                elif provider == "ollama":
                    model = st.session_state.ollama_model

                # Create LLM handler with overrides
                llm_handler = get_llm_handler(
                    provider_override=provider,
                    api_key_override=api_key,
                    model_override=model,
                )

                st.session_state.rag_pipeline = RAGPipeline(llm_handler=llm_handler)
                logger.info(f"RAG pipeline loaded with provider: {provider}")

            except FileNotFoundError:
                st.error(
                    "❌ Vector store not found! Please build it first:\n\n"
                    "```bash\n"
                    "python -m src.build_vectorstore\n"
                    "```"
                )
                st.info(
                    "Make sure you have:\n"
                    "1. Added your documents to `data/documents/`\n"
                    "2. Run the build command above\n"
                    "3. Started Ollama and pulled the model"
                )
                st.stop()

            except Exception as e:
                error_str = str(e).lower()

                # Check if it's a Groq API error - try fallback to Ollama
                if st.session_state.llm_provider == "groq" and (
                    "401" in error_str
                    or "invalid" in error_str
                    or "api_key" in error_str
                    or "429" in error_str
                    or "rate" in error_str
                ):
                    logger.warning(f"Groq API error: {e}. Falling back to Ollama.")

                    # Show warning to user
                    st.warning(
                        "⚠️ **Groq API Error** - Falling back to Ollama\n\n"
                        "Possible causes:\n"
                        "- Invalid API key\n"
                        "- Rate limit exceeded (free tier cap)\n"
                        "- API key expired\n\n"
                        "**Get a new key at:** https://console.groq.com\n\n"
                        "Using local Ollama instead..."
                    )

                    # Fallback to Ollama
                    try:
                        st.session_state.llm_provider = "ollama"
                        ollama_model = st.session_state.ollama_model or get_default_ollama_model()

                        llm_handler = get_llm_handler(
                            provider_override="ollama",
                            model_override=ollama_model,
                        )

                        st.session_state.rag_pipeline = RAGPipeline(llm_handler=llm_handler)
                        logger.info("Fallback to Ollama successful")
                        st.success(f"✅ Now using Ollama with model: {ollama_model}")
                        return

                    except Exception as fallback_error:
                        st.error(f"Fallback to Ollama also failed: {fallback_error}")
                        st.info(
                            "Make sure Ollama is running:\n"
                            "- Start Ollama: `ollama serve`\n"
                            "- Pull a model: `ollama pull llama3.2:3b`"
                        )
                        logger.error(f"Ollama fallback error: {fallback_error}")
                        st.stop()

                # Not a Groq error or other error
                st.error(f"Error loading RAG pipeline: {e}")
                if st.session_state.llm_provider == "groq":
                    st.info(
                        "Groq issues:\n"
                        "- Check your API key is valid\n"
                        "- Get a free key at https://console.groq.com"
                    )
                else:
                    st.info(
                        "Ollama issues:\n"
                        "- Ollama not running: `ollama serve`\n"
                        "- Model not pulled: `ollama pull llama3.2:3b`"
                    )
                logger.error(f"RAG pipeline error: {e}", exc_info=True)
                st.stop()


def configure_page():
    """Configure Streamlit page settings (must run before other st.* calls)."""
    # Defaults first so set_page_config can run even if config is not ready
    page_title = "💬 ProfillyBot"
    page_icon = "🤖"
    layout = "centered"
    if st.session_state.config is not None:
        page_title = st.session_state.config.get("ui.page_title", page_title)
        page_icon = st.session_state.config.get("ui.page_icon", page_icon)
        layout = st.session_state.config.get("ui.layout", layout)

    st.set_page_config(page_title=page_title, page_icon=page_icon, layout=layout)


def display_header():
    """Display app header."""
    config = st.session_state.config

    profile_name = config.get("profile.name", "Candidate")
    profile_title = config.get("profile.title", "Professional")
    greeting = config.format_template(config.get("profile.greeting", ""))

    st.title(f"💬 ProfillyBot — Chat with {profile_name}'s Profile")
    st.markdown(f"**{profile_title}**")

    if greeting:
        st.info(greeting)

    st.divider()


def display_sidebar():
    """Display sidebar with information and controls."""
    config = st.session_state.config

    with st.sidebar:
        st.header("ℹ️ About")

        profile_name = config.get("profile.name", "the candidate")
        st.markdown(
            f"This chatbot can answer questions about **{profile_name}'s** "
            "professional background, skills, experience, and projects."
        )

        st.divider()

        # Example questions
        st.subheader("💡 Example Questions")
        example_questions = config.get("ui.example_questions", [])

        for question in example_questions:
            formatted_q = config.format_template(question)
            if st.button(formatted_q, key=f"example_{hash(formatted_q)}"):
                st.session_state.pending_question = formatted_q
                st.rerun()

        st.divider()

        # LLM Provider Selection
        st.subheader("🤖 LLM Provider")

        # Provider selector
        providers = ["ollama", "groq"]
        current_provider_idx = (
            providers.index(st.session_state.llm_provider)
            if st.session_state.llm_provider in providers
            else 0
        )

        new_provider = st.radio(
            "Select Provider",
            providers,
            index=current_provider_idx,
            format_func=lambda x: "🦙 Ollama (Local)" if x == "ollama" else "⚡ Groq (Cloud)",
            key="provider_radio",
            help="Ollama runs locally. Groq requires an API key but is faster.",
        )

        # Groq-specific settings
        if new_provider == "groq":
            st.markdown("---")
            groq_key = st.text_input(
                "Groq API Key",
                value=st.session_state.groq_api_key,
                type="password",
                placeholder="Enter your Groq API key",
                help="Get a free key at https://console.groq.com",
                key="groq_key_input",
            )

            groq_models = get_available_groq_models()
            current_model_idx = (
                groq_models.index(st.session_state.groq_model)
                if st.session_state.groq_model in groq_models
                else 0
            )

            new_model = st.selectbox(
                "Groq Model",
                groq_models,
                index=current_model_idx,
                key="groq_model_select",
            )

            # Check if settings changed
            settings_changed = (
                new_provider != st.session_state.llm_provider
                or groq_key != st.session_state.groq_api_key
                or new_model != st.session_state.groq_model
            )

            if settings_changed and st.button("🔄 Apply Changes", type="primary"):
                st.session_state.llm_provider = new_provider
                st.session_state.groq_api_key = groq_key
                st.session_state.groq_model = new_model
                st.session_state.rag_pipeline = None  # Force reload
                st.rerun()
        else:
            # Ollama selected - show model selector
            st.markdown("---")
            ollama_models = get_available_ollama_models()
            current_ollama_idx = (
                ollama_models.index(st.session_state.ollama_model)
                if st.session_state.ollama_model in ollama_models
                else 0
            )

            new_ollama_model = st.selectbox(
                "Ollama Model",
                ollama_models,
                index=current_ollama_idx,
                key="ollama_model_select",
                help="Make sure the model is pulled: ollama pull <model>",
            )

            settings_changed = (
                new_provider != st.session_state.llm_provider
                or new_ollama_model != st.session_state.ollama_model
            )

            if settings_changed and st.button("🔄 Apply Changes", type="primary"):
                st.session_state.llm_provider = new_provider
                st.session_state.ollama_model = new_ollama_model
                st.session_state.rag_pipeline = None
                st.rerun()

        # Show current active settings
        st.markdown("---")
        st.caption(f"**Active:** {st.session_state.llm_provider.upper()}")
        if st.session_state.llm_provider == "groq":
            st.caption(f"**Model:** {st.session_state.groq_model}")
        else:
            st.caption(f"**Model:** {st.session_state.ollama_model}")

        st.divider()

        # Clear chat button
        if st.button("🗑️ Clear Chat History"):
            st.session_state.messages = []
            st.rerun()

        # Settings
        with st.expander("⚙️ Advanced Settings"):
            st.markdown("**Vector Store**")
            st.markdown(f"- Collection: {config.get('vectorstore.collection_name', 'N/A')}")
            st.markdown(f"- Top K: {config.get('vectorstore.search_kwargs.k', 'N/A')}")

            st.markdown("**Document Processing**")
            st.markdown(f"- Chunk Size: {config.get('document_processing.chunk_size', 'N/A')}")
            st.markdown(f"- Overlap: {config.get('document_processing.chunk_overlap', 'N/A')}")


def extract_chat_history(messages: list, exclude_last: bool = True) -> list[dict]:
    """Extract and format chat history for LLM.

    Args:
        messages: List of Streamlit message dicts
        exclude_last: Whether to exclude the last message (usually the current one being processed)

    Returns:
        List of formatted message dicts with 'role' and 'content' keys
    """
    if not messages:
        return []

    # Exclude the last message if requested (current user message being processed)
    history_messages = messages[:-1] if exclude_last else messages

    # Filter to only user and assistant messages, format for LLM
    formatted_history = []
    for msg in history_messages:
        role = msg.get("role")
        content = msg.get("content", "")

        if role in ("user", "assistant") and content:
            formatted_history.append({"role": role, "content": content})

    return formatted_history


def truncate_chat_history(
    history: list[dict], max_turns: int | None = None, max_tokens: int | None = None
) -> list[dict]:
    """Truncate chat history based on turn count or token limit.

    Args:
        history: List of message dicts
        max_turns: Maximum number of Q&A pairs to keep (None = no limit)
        max_tokens: Maximum tokens allowed (None = no limit)

    Returns:
        Truncated history list
    """
    from src.main_document_loader import get_main_document_loader

    if not history:
        return []

    # Apply turn-based truncation first (keep most recent)
    if max_turns is not None and max_turns > 0:
        # Each turn = 1 user + 1 assistant message
        max_messages = max_turns * 2
        history = history[-max_messages:] if len(history) > max_messages else history

    # Apply token-based truncation if needed
    if max_tokens is not None and max_tokens > 0:
        loader = get_main_document_loader()
        current_tokens = 0
        truncated_history = []

        # Go through history in reverse, adding messages until we hit token limit
        for msg in reversed(history):
            content = msg.get("content", "")
            msg_tokens = loader.count_tokens(content)

            if current_tokens + msg_tokens <= max_tokens:
                truncated_history.insert(0, msg)
                current_tokens += msg_tokens
            else:
                break

        return truncated_history

    return history


def display_chat_messages():
    """Render conversation history with Streamlit chat elements.

    Uses ``st.chat_message`` containers as documented at:
    https://docs.streamlit.io/develop/api-reference/chat
    """
    config = st.session_state.config
    show_timestamps = config.get("ui.show_timestamps", True)

    for message in st.session_state.messages:
        role = message["role"]
        avatar = "🧑‍💻" if role == "user" else "🤖"
        with st.chat_message(role, avatar=avatar):
            st.markdown(message["content"])

            if show_timestamps and message.get("timestamp"):
                st.caption(f"_{message['timestamp']}_")

            if message.get("sources"):
                with st.expander("📚 Sources"):
                    st.markdown(message["sources"])

            if message.get("retrieval_panel"):
                with st.expander("🔍 Retrieval (query + chunks)", expanded=False):
                    st.markdown(message["retrieval_panel"])


def _prepare_chat_history(config) -> list[dict]:
    """Build truncated prior-turn history for the RAG pipeline."""
    if not config.get("chat.enable_history", True):
        return []

    raw_history = extract_chat_history(st.session_state.messages, exclude_last=True)
    max_turns = config.get("chat.max_history_turns", 10)
    max_tokens = config.get("chat.max_history_tokens", 2000)

    chat_history = truncate_chat_history(
        raw_history,
        max_turns=max_turns if max_turns > 0 else None,
        max_tokens=max_tokens if max_tokens > 0 else None,
    )
    if chat_history:
        logger.debug("Including %s messages from chat history", len(chat_history))
    return chat_history


def process_user_input(prompt: str):
    """Handle a user prompt with ``st.chat_message`` + ``st.write_stream``.

    Args:
        prompt: User's question
    """
    config = st.session_state.config
    rag_pipeline = st.session_state.rag_pipeline
    timestamp = datetime.now().strftime("%H:%M:%S")

    # 1) Persist + show user turn
    st.session_state.messages.append(
        {"role": "user", "content": prompt, "timestamp": timestamp}
    )
    with st.chat_message("user", avatar="🧑‍💻"):
        st.markdown(prompt)
        if config.get("ui.show_timestamps", True):
            st.caption(f"_{timestamp}_")

    chat_history = _prepare_chat_history(config)

    # 2) Stream assistant reply into a chat_message container
    with st.chat_message("assistant", avatar="🤖"):
        try:
            # Retrieve once, then stream the answer
            sources = []
            retrieval_panel = ""
            if config.get("rag.include_sources", True) or config.get(
                "rag.show_retrieval_panel", True
            ):
                sources = rag_pipeline.get_source_documents(prompt, chat_history=chat_history)
                if config.get("rag.show_retrieval_panel", True):
                    retrieval_panel = rag_pipeline.format_retrieval_panel(
                        prompt, sources=sources
                    )
                    with st.expander("🔍 Retrieval (query + chunks)", expanded=True):
                        st.markdown(retrieval_panel)

            streamed_content = st.write_stream(
                rag_pipeline.stream_answer(prompt, chat_history=chat_history)
            )
            if streamed_content is None:
                streamed_content = ""
            elif not isinstance(streamed_content, str):
                streamed_content = str(streamed_content)

            # Keep the streamed UI as-is; persist an enhanced copy in history if enabled
            answer = streamed_content
            if config.get("rag.enhance_responses", True) and streamed_content:
                from src.response_enhancer import get_response_enhancer

                enhancer = get_response_enhancer()
                enhanced = enhancer.enhance_with_context(streamed_content, prompt)
                if enhanced != streamed_content:
                    answer = enhanced
                    logger.debug("Response enhanced for better tone (stored in history)")

            sources_text = ""
            if config.get("rag.include_sources", True) and sources:
                sources_text = rag_pipeline.format_sources(sources)
                with st.expander("📚 Sources"):
                    st.markdown(sources_text)

            st.session_state.messages.append(
                {
                    "role": "assistant",
                    "content": answer,
                    "timestamp": datetime.now().strftime("%H:%M:%S"),
                    "sources": sources_text,
                    "retrieval_panel": retrieval_panel,
                }
            )

        except Exception as e:
            error_msg = f"I apologize, but I encountered an error: {e}"
            st.error(error_msg)
            logger.error("Error generating response: %s", e, exc_info=True)
            st.session_state.messages.append(
                {
                    "role": "assistant",
                    "content": error_msg,
                    "timestamp": timestamp,
                }
            )


def main():
    """Main application entry point."""
    init_session_state()
    # set_page_config must be the first Streamlit command
    configure_page()
    load_config()
    # Re-apply title from config if needed is already handled via defaults above
    load_rag_pipeline()

    display_header()
    display_sidebar()

    # Conversation history (st.chat_message)
    display_chat_messages()

    # Example-question buttons inject into the same chat flow
    if st.session_state.pending_question:
        pending = st.session_state.pending_question
        st.session_state.pending_question = None
        process_user_input(pending)
        st.rerun()

    # Pinned chat composer (st.chat_input)
    # https://docs.streamlit.io/develop/api-reference/chat/st.chat_input
    chat_input_kwargs = {
        "key": "profilly_chat_input",
        "max_chars": 4000,
    }
    # submit_mode is available on newer Streamlit versions
    try:
        prompt = st.chat_input(
            "Ask me anything about the profile...",
            submit_mode="disable",
            **chat_input_kwargs,
        )
    except TypeError:
        prompt = st.chat_input(
            "Ask me anything about the profile...",
            **chat_input_kwargs,
        )
    if prompt:
        # Newer Streamlit may return a dict-like object when files/audio are enabled
        if not isinstance(prompt, str):
            prompt = getattr(prompt, "text", None) or str(prompt)
        if prompt:
            process_user_input(prompt)

    max_history = st.session_state.config.get("ui.max_chat_history", 50)
    if len(st.session_state.messages) > max_history:
        st.session_state.messages = st.session_state.messages[-max_history:]


if __name__ == "__main__":
    main()