"""Per-category mission generator.

Generates investigation missions for a single category via one or more LLM calls
with retry handling for transient errors. Supports multi-call generation
for categories that exceed max_tokens limits.
"""

from __future__ import annotations

import functools
import logging
from typing import TYPE_CHECKING

from unfamiliar.constants import DEFAULT_TOKEN_BUDGET, InvestigationCategory
from unfamiliar.llm.exceptions import AuthenticationError, MissionGenerationError
from unfamiliar.llm.missions.context import (
    build_mission_context,
    build_module_descriptions,
)
from unfamiliar.llm.missions.prompts import MISSION_PROMPT, build_theme_prompt_vars
from unfamiliar.llm.missions.schemas import (
    CategoryMissionsResponse,
    ChainResponse,
    MissionItemResponse,
)
from unfamiliar.llm.retry import with_llm_retry
from unfamiliar.llm.structured import create_structured_chain, detect_truncation

if TYPE_CHECKING:
    from langchain_core.language_models.chat_models import BaseChatModel

    from unfamiliar.analysis.models import AnalysisResult
    from unfamiliar.llm.briefing.generator import BriefingResult
    from unfamiliar.tui.themes import ThemeConfig

logger = logging.getLogger(__name__)

# Multi-call generation constants
_MISSIONS_PER_CALL = 3  # Request 3 missions per call to avoid truncation
_MAX_CALLS_PER_CATEGORY = 3  # Maximum 3 calls per category (up to 9 missions)


def generate_category_missions(
    analysis: AnalysisResult,
    briefing_result: BriefingResult,
    category: InvestigationCategory,
    model: BaseChatModel,
    previous_missions: list[str] | None = None,
    token_budget: int = DEFAULT_TOKEN_BUDGET,
    provider_name: str | None = None,
    theme_config: ThemeConfig | None = None,
) -> CategoryMissionsResponse:
    """Generate missions for a single investigation category via one LLM call.

    Args:
        analysis: Complete AnalysisResult from analyze_repository().
        briefing_result: Complete BriefingResult from generate_briefing().
        category: The investigation category to generate missions for.
        model: A configured BaseChatModel instance.
        previous_missions: title_plain values from previously completed categories.
        token_budget: Total token budget (will be divided by 7 per category).
        provider_name: Provider name for retry delays.
        theme_config: Optional theme configuration for narrative style.

    Returns:
        CategoryMissionsResponse containing missions and chains for the category.

    Raises:
        AuthenticationError: If the LLM API returns a 401/403 error.
        MissionGenerationError: If generation fails after max retries.
    """
    if previous_missions is None:
        previous_missions = []

    logger.info("Generating missions for category: %s", category.value)

    # Build context for this category
    context = build_mission_context(
        analysis=analysis,
        briefing_result=briefing_result,
        category=category,
        previous_missions=previous_missions,
        token_budget=token_budget,
    )
    logger.debug(
        "Built mission context with %d characters for %s",
        len(context),
        category.value,
    )

    # Build module descriptions for linking
    module_descriptions = build_module_descriptions(briefing_result)

    # Get briefing summary
    briefing_summary = briefing_result.briefing.summary

    # Create the structured output chain (provider-agnostic)

    chain = create_structured_chain(
        model, MISSION_PROMPT, CategoryMissionsResponse, provider_name=provider_name
    )

    # Prepare input with theme variables
    input_data = {
        "category": category.value.replace("_", " ").title(),
        "briefing_summary": briefing_summary,
        "category_context": context,
        "module_descriptions": module_descriptions,
        "previous_missions": (
            "\n".join(f"- {title}" for title in previous_missions)
            if previous_missions
            else "None (this is the first category)"
        ),
    }
    # Add theme variables (with fallback to generic investigation style)
    input_data.update(build_theme_prompt_vars(theme_config))

    # Call LLM with retry handling
    attempts = 0
    try:

        def invoke_chain() -> CategoryMissionsResponse:
            nonlocal attempts
            attempts += 1
            result = chain.invoke(input_data)
            if not isinstance(result, CategoryMissionsResponse):
                raise TypeError(
                    f"Expected CategoryMissionsResponse, got {type(result).__name__}"
                )
            return result

        response = with_llm_retry(
            invoke_chain, max_attempts=5, provider_name=provider_name
        )
    except AuthenticationError as exc:
        logger.error(
            "Mission generation failed for %s after %d attempts: %s",
            category.value,
            attempts,
            exc,
        )
        raise
    except Exception as exc:
        raise MissionGenerationError(
            f"Failed to generate missions: {exc}",
            category=category.value,
            attempts=attempts,
        ) from exc

    # Validate at least one mission was generated
    if not response.missions:
        logger.warning("LLM returned zero missions for %s", category.value)
        raise MissionGenerationError(
            "LLM returned no missions for this category",
            category=category.value,
            attempts=attempts,
        )

    logger.info(
        "Generated %d missions and %d chains for %s",
        len(response.missions),
        len(response.chains),
        category.value,
    )

    return response


def generate_category_missions_multi(
    analysis: AnalysisResult,
    briefing_result: BriefingResult,
    category: InvestigationCategory,
    model: BaseChatModel,
    previous_missions: list[str] | None = None,
    token_budget: int = DEFAULT_TOKEN_BUDGET,
    provider_name: str | None = None,
    missions_per_call: int = _MISSIONS_PER_CALL,
    max_calls: int = _MAX_CALLS_PER_CATEGORY,
    theme_config: ThemeConfig | None = None,
) -> CategoryMissionsResponse:
    """Generate missions for a category using multiple LLM calls with truncation resume.

    This function makes multiple calls (up to max_calls) requesting
    missions_per_call missions each. If a response is truncated, it
    continues from the last complete mission.

    Args:
        analysis: Complete AnalysisResult from analyze_repository().
        briefing_result: Complete BriefingResult from generate_briefing().
        category: The investigation category to generate missions for.
        model: A configured BaseChatModel instance.
        previous_missions: title_plain values from previously completed categories.
        token_budget: Total token budget (will be divided by 7).
        provider_name: Provider name for retry delay (e.g., "glm" or "anthropic").
        missions_per_call: Number of missions to request per call (default 3).
        max_calls: Maximum number of calls to make (default 3).
        theme_config: Optional theme configuration for narrative style.

    Returns:
        CategoryMissionsResponse containing all merged missions and chains.

    Raises:
        AuthenticationError: If the LLM API returns a 401/403 error.
        MissionGenerationError: If all calls fail or no missions generated.
    """
    if previous_missions is None:
        previous_missions = []

    logger.info(
        "Starting multi-call mission generation for %s "
        "(up to %d calls, %d missions each)",
        category.value,
        max_calls,
        missions_per_call,
    )

    # Build context for this category
    context = build_mission_context(
        analysis=analysis,
        briefing_result=briefing_result,
        category=category,
        previous_missions=previous_missions,
        token_budget=token_budget,
    )
    logger.debug(
        "Built mission context with %d characters for %s",
        len(context),
        category.value,
    )

    # Build module descriptions for linking
    module_descriptions = build_module_descriptions(briefing_result)

    # Get briefing summary
    briefing_summary = briefing_result.briefing.summary

    # Create the structured output chain (provider-agnostic)

    chain = create_structured_chain(
        model, MISSION_PROMPT, CategoryMissionsResponse, provider_name=provider_name
    )

    all_missions: list[MissionItemResponse] = []
    all_chains: list[ChainResponse] = []
    all_warnings: list[str] = []
    call_count = 0
    seen_titles: set[str] = set()  # Track titles to avoid duplicates

    try:
        while call_count < max_calls:
            call_count += 1

            # Prepare input if resuming
            input_data = {
                "category": category.value.replace("_", " ").title(),
                "briefing_summary": briefing_summary,
                "category_context": context,
                "module_descriptions": module_descriptions,
                "previous_missions": (
                    "\n".join(f"- {title}" for title in previous_missions)
                    if previous_missions
                    else "None (this is the first category)"
                ),
            }
            # Add theme variables (with fallback to generic investigation style)
            input_data.update(build_theme_prompt_vars(theme_config))

            invoke_chain = functools.partial(chain.invoke, input_data)

            # Call LLM with retry handling
            try:
                raw_result = with_llm_retry(
                    invoke_chain,
                    max_attempts=5,
                    provider_name=provider_name,
                )
                if not isinstance(raw_result, CategoryMissionsResponse):
                    raise TypeError(
                        f"Expected CategoryMissionsResponse, "
                        f"got {type(raw_result).__name__}"
                    )
                response = raw_result
            except AuthenticationError as exc:
                logger.error(
                    "Mission generation failed for %s after %d attempts: %s",
                    category.value,
                    call_count,
                    exc,
                )
                raise
            except Exception as exc:
                # If we have some missions already, return partial results
                if all_missions:
                    logger.warning(
                        "Call %d/%d failed for %s, returning %d partial missions: %s",
                        call_count,
                        max_calls,
                        category.value,
                        len(all_missions),
                        exc,
                    )
                    all_warnings.append(str(exc))
                    break
                raise MissionGenerationError(
                    f"Failed to generate missions: {exc}",
                    category=category.value,
                    attempts=call_count,
                ) from exc

            # Check for truncation
            truncation = detect_truncation(
                response.model_dump_json(),
                CategoryMissionsResponse,
            )

            # Add non-duplicate missions to our collection
            new_missions = 0
            for mission in response.missions:
                if mission.title_plain not in seen_titles:
                    seen_titles.add(mission.title_plain)
                    all_missions.append(mission)
                    new_missions += 1

            logger.info(
                "Call %d/%d: %d missions (+%d new), truncated=%s",
                call_count,
                max_calls,
                len(response.missions),
                new_missions,
                truncation.is_truncated,
            )

            # Collect chains (avoid duplicates by name)
            for chain_item in response.chains:
                if not any(c.name_plain == chain_item.name_plain for c in all_chains):
                    all_chains.append(chain_item)

            # If not truncated, we're done
            if not truncation.is_truncated:
                logger.info(
                    "Response complete after %d calls for %s",
                    call_count,
                    category.value,
                )
                break

            logger.info(
                "Response truncated, preparing continuation call %d with same prompt.",
                call_count + 1,
            )

    except AuthenticationError:
        raise
    except (MissionGenerationError, TimeoutError, OSError) as exc:
        # External errors - return partial results if we have any
        # Programming errors (NameError, TypeError, etc.) propagate
        if all_missions:
            logger.error(
                "Unexpected error in multi-call generation for %s, "
                "returning %d partial missions: %s",
                category.value,
                len(all_missions),
                exc,
            )
            return CategoryMissionsResponse(
                missions=all_missions,
                chains=all_chains,
                warnings=[str(exc)],
            )
        else:
            raise MissionGenerationError(
                f"Failed to generate missions: {exc}",
                category=category.value,
                attempts=call_count,
            ) from exc

    # Validate at least one mission was generated
    if not all_missions:
        logger.warning(
            "Zero missions generated for %s after %d calls",
            category.value,
            call_count,
        )
        raise MissionGenerationError(
            "No missions generated after all calls",
            category=category.value,
            attempts=call_count,
        )

    logger.info(
        "Multi-call generation complete for %s: "
        "%d total missions, %d chains from %d calls",
        category.value,
        len(all_missions),
        len(all_chains),
        call_count,
    )

    return CategoryMissionsResponse(
        missions=all_missions, chains=all_chains, warnings=all_warnings
    )
