"""Suggested questions generation for the Databao chat interface."""

import logging
from concurrent.futures import Future, ThreadPoolExecutor
from concurrent.futures import TimeoutError as FuturesTimeoutError
from typing import TYPE_CHECKING

import streamlit as st
from langchain_core.messages import HumanMessage, SystemMessage
from pydantic import BaseModel, Field

if TYPE_CHECKING:
    from databao.core.agent import Agent

logger = logging.getLogger(__name__)

SuggestionsResult = tuple[list[str], bool]


@st.cache_resource
def _get_suggestions_executor() -> ThreadPoolExecutor:
    """Get a shared thread pool executor for suggestions generation."""
    return ThreadPoolExecutor(max_workers=1, thread_name_prefix="suggestions")


FALLBACK_QUESTIONS = [
    "What tables are available in my database?",
    "Show me a summary of the main data",
    "What are the key metrics I can analyze?",
]


class SuggestedQuestions(BaseModel):
    """Structured output model for suggested questions."""

    questions: list[str] = Field(
        description="List of exactly 3 short questions about the data",
        min_length=3,
        max_length=3,
    )


def _build_context_from_sources(agent: "Agent") -> str:
    """Extract schema/context information from agent sources."""
    context_parts: list[str] = []

    for name, db_source in agent.dbs.items():
        if db_source.description:
            context_parts.append(f"Database '{name}':\n{db_source.description}")

    for name, df_source in agent.dfs.items():
        if df_source.description:
            context_parts.append(f"DataFrame '{name}':\n{df_source.description}")

    for ctx in agent.additional_description:
        context_parts.append(ctx)

    if not context_parts:
        return "No detailed schema information available."

    return "\n\n".join(context_parts)


SUGGESTION_PROMPT = """You are helping suggest questions for a data analysis assistant.
The user has connected the following data sources:

{context}

Generate exactly 3 short questions (maximum 60 characters each) that:
1. Would produce interesting data tables or visualizations
2. Are answerable using the available data
3. Cover different aspects of the data (e.g., summaries, trends, comparisons)

The questions should be simple and direct, suitable for display as clickable buttons.
Do not include numbering or bullet points in the questions themselves."""


def generate_suggested_questions(agent: "Agent") -> tuple[list[str], bool]:
    """Generate suggested questions based on the agent's data sources.

    Uses the agent's LLM with structured output to generate contextual questions.
    Falls back to generic questions if LLM generation fails.

    Args:
        agent: The Databao agent with configured data sources.

    Returns:
        A tuple of (questions, is_llm_generated) where:
        - questions: List of 3 question strings
        - is_llm_generated: True if questions were generated by LLM, False if fallback
    """
    try:
        context = _build_context_from_sources(agent)

        prompt = SUGGESTION_PROMPT.format(context=context)

        llm_with_structure = agent.llm.with_structured_output(SuggestedQuestions)

        result = llm_with_structure.invoke(
            [
                SystemMessage(content=prompt),
                HumanMessage(content="Generate 3 suggested questions for exploring this data."),
            ]
        )

        if isinstance(result, SuggestedQuestions) and result.questions and len(result.questions) == 3:
            logger.info(f"Successfully generated {len(result.questions)} suggested questions via LLM")
            return list(result.questions), True

        logger.warning("LLM returned invalid result, using fallback questions")
        return FALLBACK_QUESTIONS.copy(), False

    except Exception as e:
        logger.warning(f"Failed to generate suggested questions: {e}. Using fallback.")
        return FALLBACK_QUESTIONS.copy(), False


def _generate_suggestions_task(agent: "Agent") -> SuggestionsResult:
    """Background task that generates suggestions.

    Args:
        agent: The Databao agent.

    Returns:
        Tuple of (questions, is_llm_generated).
    """
    try:
        return generate_suggested_questions(agent)
    except Exception as e:
        logger.warning(f"Background suggestions generation failed: {e}")
        return FALLBACK_QUESTIONS.copy(), False


def start_suggestions_generation(agent: "Agent") -> bool:
    """Start background suggestions generation.

    Returns True if generation was started, False if already running or ready.
    """
    status = st.session_state.get("suggestions_status", "not_started")

    if status != "not_started":
        return False

    executor = _get_suggestions_executor()
    future: Future[SuggestionsResult] = executor.submit(_generate_suggestions_task, agent)

    st.session_state.suggestions_future = future
    st.session_state.suggestions_status = "loading"

    logger.info("Started background suggestions generation")
    return True


def reset_suggestions_state() -> None:
    """Reset suggestions state to allow regeneration.

    Call this when Reload button is pressed to regenerate suggestions.
    """
    st.session_state.suggestions_future = None
    st.session_state.suggestions_status = "not_started"
    st.session_state.suggested_questions = []
    st.session_state.suggestions_are_llm_generated = False

    logger.info("Reset suggestions state")


def check_suggestions_completion() -> bool:
    """Check if background suggestions generation has completed.

    If completed, updates session state with results.

    Returns True if suggestions became ready (caller should rerun to show them).
    """
    status = st.session_state.get("suggestions_status")
    if status != "loading":
        return False

    future: Future[SuggestionsResult] | None = st.session_state.get("suggestions_future")

    if future is None:
        return False

    if not future.done():
        return False

    try:
        result = future.result(timeout=1.0)
    except FuturesTimeoutError:
        logger.warning("Timeout getting completed future result")
        result = None
    except Exception as e:
        logger.warning(f"Failed to get suggestions result: {e}")
        result = None

    st.session_state.suggestions_future = None

    if result is None:
        st.session_state.suggested_questions = FALLBACK_QUESTIONS.copy()
        st.session_state.suggestions_are_llm_generated = False
        st.session_state.suggestions_status = "ready"
        logger.info("Suggestions generation failed, using fallback")
        return True

    questions, is_llm_generated = result
    st.session_state.suggested_questions = questions
    st.session_state.suggestions_are_llm_generated = is_llm_generated
    st.session_state.suggestions_status = "ready"

    logger.info(f"Suggestions generation completed with {len(questions)} questions")
    return True


def is_suggestions_loading() -> bool:
    """Check if suggestions are currently being generated.

    Useful for the polling fragment to know when to stop.
    """
    return st.session_state.get("suggestions_status") == "loading"
