"""Mission generation pipeline with checkpointing.

Processes all 7 investigation categories sequentially with SQLite
checkpointing for fault-tolerant, resumable execution.
"""

from __future__ import annotations

import hashlib
import logging
import sqlite3
import time
from contextlib import contextmanager
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING

from langgraph.checkpoint.sqlite import SqliteSaver

from unfamiliar.constants import DEFAULT_TOKEN_BUDGET, InvestigationCategory
from unfamiliar.llm.exceptions import AuthenticationError, MissionGenerationError
from unfamiliar.llm.missions.generator import generate_category_missions_multi
from unfamiliar.llm.missions.validation import (
    ValidationResult,
    validate_category_missions,
)

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.llm.missions.schemas import CategoryMissionsResponse
    from unfamiliar.tui.themes import ThemeConfig

logger = logging.getLogger(__name__)

# Category processing order (matters for cross-category prerequisites)
CATEGORY_ORDER: list[InvestigationCategory] = [
    InvestigationCategory.ARCHITECTURE,
    InvestigationCategory.DATA_FLOW,
    InvestigationCategory.DOMAIN_BUSINESS_LOGIC,
    InvestigationCategory.DEPENDENCIES,
    InvestigationCategory.CONVENTIONS,
    InvestigationCategory.TESTING,
    InvestigationCategory.TECHNICAL_DEBT,
]


@dataclass
class PipelineResult:
    """Complete output of the mission generation pipeline.

    Attributes:
        category_missions: Generated missions keyed by category name.
        category_validations: Validation results keyed by category name.
        completed_categories: List of successfully completed category names.
        errors: Error messages for failed categories.
    """

    category_missions: dict[str, CategoryMissionsResponse] = field(default_factory=dict)
    category_validations: dict[str, ValidationResult] = field(default_factory=dict)
    completed_categories: list[str] = field(default_factory=list)
    errors: list[str] = field(default_factory=list)


@contextmanager
def _checkpoint_context(campaign_dir: Path):
    """Context manager for checkpointing pipeline state."""
    campaign_dir.mkdir(parents=True, exist_ok=True)
    db_path = campaign_dir / "mission_checkpoint.db"
    conn = sqlite3.connect(str(db_path), check_same_thread=False)
    try:
        yield SqliteSaver(conn)
    finally:
        conn.close()


def run_mission_pipeline(
    analysis: AnalysisResult,
    briefing_result: BriefingResult,
    model: BaseChatModel,
    campaign_dir: Path,
    token_budget: int = DEFAULT_TOKEN_BUDGET,
    provider_name: str | None = None,
    category_cooldown: float = 0.0,
    categories: list[InvestigationCategory] | None = None,
    theme_config: ThemeConfig | None = None,
) -> PipelineResult:
    """Execute the mission generation pipeline for selected categories.

    Processes categories in defined order (Architecture first, Technical Debt last).
    Each completed category is checkpointed - resume skips completed categories.

    Args:
        analysis: Complete AnalysisResult from analyze_repository().
        briefing_result: Complete BriefingResult from generate_briefing().
        model: Any BaseChatModel instance.
        campaign_dir: Path to campaign storage directory (for checkpoint file).
        token_budget: Positive integer for context token budget.
        provider_name: Provider name for provider-specific retry delays
            (e.g., "glm" for 30s base).
        category_cooldown: Seconds to sleep between category processing.
        categories: Subset of categories to process, in CATEGORY_ORDER sequence.
            None means all 7 categories.
        theme_config: Optional theme configuration for narrative style.

    Returns:
        PipelineResult with missions for all completed categories.

    Raises:
        AuthenticationError: Immediately on 401/403 errors (partial results
            returned in PipelineResult).
    """
    # Filter to requested categories, preserving canonical order
    if categories is not None:
        requested = set(categories)
        run_categories = [c for c in CATEGORY_ORDER if c in requested]
    else:
        run_categories = list(CATEGORY_ORDER)

    logger.info(
        "Starting mission pipeline for %s (%d categories)",
        analysis.repo_path,
        len(run_categories),
    )

    # Initialize result
    result = PipelineResult()

    with _checkpoint_context(campaign_dir) as checkpointer:
        thread_id = _derive_thread_id(campaign_dir)

        # Track state across categories
        all_mission_titles: list[str] = []

        # Process categories in order
        for cat_idx, category in enumerate(run_categories):
            category_name = category.value
            logger.info("Processing category: %s", category_name)

            # Apply cooldown between categories (skip before first category)
            if cat_idx > 0 and category_cooldown > 0:
                logger.debug(
                    "Cooldown %.1fs before %s", category_cooldown, category_name
                )
                time.sleep(category_cooldown)

            try:
                # Generate missions (multi-call with truncation resume)
                response = generate_category_missions_multi(
                    analysis=analysis,
                    briefing_result=briefing_result,
                    category=category,
                    model=model,
                    previous_missions=all_mission_titles.copy(),
                    token_budget=token_budget,
                    provider_name=provider_name,
                    theme_config=theme_config,
                )

                # Validate the generated missions
                validation = validate_category_missions(
                    missions_response=response.missions,
                    repo_path=Path(analysis.repo_path),
                    category=category,
                    chains=response.chains,
                )

                # Replace response data with validated versions before storing
                # This ensures persist_pipeline_data receives cleaned data
                response.missions = validation.missions
                response.chains = validation.chains

                # Store results
                result.category_missions[category_name] = response
                result.category_validations[category_name] = validation
                result.completed_categories.append(category_name)

                # Collect mission titles for subsequent categories
                for mission in validation.missions:
                    all_mission_titles.append(mission.title_plain)

                # Log validation warnings (separate auto-corrected from real issues)
                # "Clamped" and "Fixed line_" are trivial LLM number format errors
                autocorrected_warnings = [
                    w
                    for w in validation.warnings
                    if w.startswith("Clamped") or w.startswith("Fixed line_")
                ]
                real_warnings = [
                    w for w in validation.warnings if w not in autocorrected_warnings
                ]

                # Auto-corrections are logged at debug level (not user-facing)
                if autocorrected_warnings:
                    logger.debug(
                        "Auto-fixed %d line range issues in %s",
                        len(autocorrected_warnings),
                        category_name,
                    )

                # Real issues (missing files, broken chains) need attention
                if real_warnings:
                    logger.warning(
                        "Validation issues for %s: %s",
                        category_name,
                        real_warnings[:3],  # First 3 only
                    )

                # Checkpoint after each successful category
                _save_checkpoint(checkpointer, thread_id, result)

            except AuthenticationError:
                error_msg = f"{category_name}: authentication failed"
                result.errors.append(error_msg)
                logger.error("Authentication error - stopping pipeline")
                break

            except MissionGenerationError as exc:
                error_msg = f"{category_name}: {exc}"
                result.errors.append(error_msg)
                logger.error("Category %s failed: %s", category_name, exc)

            except TimeoutError as exc:
                error_msg = f"{category_name}: {exc}"
                result.errors.append(error_msg)
                logger.error("Category %s timed out: %s", category_name, exc)

        logger.info(
            "Pipeline complete: %d categories, %d errors",
            len(result.completed_categories),
            len(result.errors),
        )

    return result


def _derive_thread_id(campaign_dir: Path) -> str:
    """Derive a deterministic thread ID from campaign directory.

    Same directory = same thread = resume capability.

    Args:
        campaign_dir: Path to campaign directory.

    Returns:
        Thread ID string for checkpointing.
    """
    dir_str = str(campaign_dir.resolve())
    return hashlib.sha256(dir_str.encode()).hexdigest()[:16]


def _save_checkpoint(
    checkpointer: object,
    thread_id: str,
    result: PipelineResult,
) -> None:
    """Save current pipeline state to checkpointer.

    Args:
        checkpointer: The checkpointer instance.
        thread_id: Thread ID for this pipeline run.
        result: Current pipeline result to save.
    """
    # For InMemorySaver, state is automatically retained in the session
    # For SqliteSaver, we need to ensure the connection is committed
    try:
        if hasattr(checkpointer, "conn"):
            checkpointer.conn.commit()  # type: ignore[union-attr]  # hasattr guard verifies conn exists; mypy can't narrow
    except (OSError, sqlite3.Error) as exc:
        logger.warning("Checkpoint save failed (resume may not work): %s", exc)
