import logging
import tomllib
from contextlib import contextmanager
from pathlib import Path
from typing import TYPE_CHECKING, Any

import tomli_w

from ._context import (  # noqa: F401 - get_input_base_dir re-exported
    get_input_base_dir,
    reset_input_base_dir,
    set_input_base_dir,
)
from ._path import AttributePart, CalcPath, ItemPart, ModelPath, PartBase, ProjectPath, VerificationPath
from ._table import Table

if TYPE_CHECKING:
    from collections.abc import Iterator, Mapping

    from pydantic import BaseModel

    from ._models import Project

logger = logging.getLogger(__name__)


# =============================================================================
# Base Directory Context for Path Resolution
# =============================================================================


@contextmanager
def input_base_dir(base_dir: Path) -> Iterator[None]:
    """Context manager to set the base directory for resolving relative paths.

    When loading input files (e.g., TOML), use this context to resolve relative
    paths (e.g., FileRef.path) relative to a base directory (typically the
    input file's parent directory).

    Example:
        with input_base_dir(toml_path.parent):
            model_data = load_model_data(...)

    """
    token = set_input_base_dir(base_dir)
    try:
        yield
    finally:
        reset_input_base_dir(token)


def _serialize_value(value: Any) -> Any:
    """Recursively serialize a value for TOML export, handling special types.

    Handles:
    - Pydantic BaseModel: Converts to dict via model_dump()
    - Table: Converts enum keys to strings and recursively serializes values
    - dict: Recursively serializes values
    - list/tuple: Recursively serializes items
    - Primitives and TOML-native types: Returns as-is
    """
    # Import BaseModel locally to avoid potential circular imports
    from pydantic import BaseModel  # noqa: PLC0415

    # Handle Pydantic BaseModel - convert to dict and recursively serialize
    if isinstance(value, BaseModel):
        # Use mode='python' to preserve Python types (datetime, Decimal, etc.)
        # which are TOML-compatible, rather than converting to JSON types
        return _serialize_value(value.model_dump(mode="python"))

    # Handle Table (dict with enum keys) - convert keys to strings and serialize values
    if isinstance(value, Table):
        result = {}
        for k, v in value.items():
            if isinstance(k, tuple):
                # Tuple of enums - join their values
                key_str = ",".join(item.value if hasattr(item, "value") else str(item) for item in k)
            elif hasattr(k, "value"):
                # Single enum
                key_str = k.value
            else:
                # Other types
                key_str = str(k)
            # Recursively serialize the value
            result[key_str] = _serialize_value(v)
        return result

    # Handle dict - recursively serialize values, excluding None (TOML doesn't support None)
    if isinstance(value, dict):
        return {k: _serialize_value(v) for k, v in value.items() if v is not None}

    # Handle list/tuple - recursively serialize items
    if isinstance(value, (list, tuple)):
        return [_serialize_value(item) for item in value]

    # Handle Path objects - convert to string
    if isinstance(value, Path):
        return str(value)

    # Primitives and TOML-native types (str, int, float, bool, datetime, etc.)
    return value


def _set_nested_value(data: dict[str, Any], keys: list[str], value: Any) -> None:
    """Set a value in a nested dictionary using a list of keys."""
    current = data
    for key in keys[:-1]:
        if key not in current:
            current[key] = {}
        current = current[key]
    current[keys[-1]] = _serialize_value(value)


def _parts_to_keys(parts: tuple[PartBase, ...]) -> list[str]:
    """Convert path parts to a list of keys for nested dictionary access."""
    keys: list[str] = []
    for part in parts:
        match part:
            case AttributePart(name):
                keys.append(name)
            case ItemPart(key):
                # For Table items, use the key directly as a dictionary key
                if isinstance(key, tuple):
                    # Multi-key access (though not common in this context)
                    keys.append(",".join(key))
                else:
                    keys.append(key)
            case _:
                msg = f"Unknown part type: {type(part)}"
                raise TypeError(msg)
    return keys


def results_to_dict(results: dict[ProjectPath, Any]) -> dict[str, Any]:
    """Convert evaluation results to a nested dictionary structure.

    This is a pure function that converts the flat results dictionary
    (with ProjectPath keys) into a nested dictionary suitable for TOML export.

    Args:
        results: The evaluation results from evaluate_project

    Returns:
        A nested dictionary with the structure:
        {
            "ScopeName": {
                "model": {...},
                "calc": {...},
                "verification": {...}
            }
        }

    """
    toml_data: dict[str, Any] = {}

    # Process all results (includes both model data and calculated/verified values)
    for ppath, value in results.items():
        scope_name = ppath.scope
        path = ppath.path

        # Build the section name based on the path type
        if isinstance(path, ModelPath):
            # Model paths: {scope}.model.{field_path}
            section_keys = [scope_name, "model"]
            field_keys = _parts_to_keys(path.parts)
        elif isinstance(path, CalcPath):
            # Calculation paths: {scope}.calc.{calc_name}.{field_path}
            section_keys = [scope_name, "calc", path.calc_name]
            field_keys = _parts_to_keys(path.parts)
        elif isinstance(path, VerificationPath):
            # Verification paths: {scope}.verification.{verification_name} or
            # {scope}.verification.{verification_name}.{field_path} for Table[K, bool]
            if path.parts:
                # Table[K, bool] verification - has parts for table indexing
                section_keys = [scope_name, "verification", path.verification_name]
                field_keys = _parts_to_keys(path.parts)
            else:
                # Plain bool verification - no parts
                section_keys = [scope_name, "verification"]
                field_keys = [path.verification_name]
        else:
            msg = f"Unknown path type: {type(path)}"
            raise TypeError(msg)

        # Combine section and field keys
        all_keys = section_keys + field_keys

        # Set the value in the nested dictionary
        _set_nested_value(toml_data, all_keys, value)

    return toml_data


def export_to_toml(
    _project: Project,
    _model_data: Mapping[str, BaseModel],
    results: dict[ProjectPath, Any],
    output_path: Path | str,
) -> None:
    """Export model data and evaluation results to a TOML file.

    Args:
        project: The project containing scope definitions
        model_data: The input model data for each scope
        results: The evaluation results from evaluate_project
        output_path: Path to the output TOML file

    """
    toml_data = results_to_dict(results)

    # Write to TOML file
    output_path = Path(output_path)
    with output_path.open("wb") as f:
        tomli_w.dump(toml_data, f)

    logger.debug(f"Exported results to {output_path}")


def toml_to_model_data(
    project: Project,
    toml_contents: dict[str, Any],
) -> dict[str, BaseModel]:
    """Convert TOML dictionary contents to validated model data.

    This is a pure function that validates and converts TOML contents
    into Pydantic model instances for each scope in the project.

    Args:
        project: The project containing scope definitions
        toml_contents: The parsed TOML dictionary

    Returns:
        A dictionary mapping scope names to their validated root model instances

    """
    from pydantic import BaseModel as PydanticBaseModel  # noqa: PLC0415, TC002

    model_data: dict[str, PydanticBaseModel] = {}
    for scope_name, scope in project.scopes.items():
        if scope_name in toml_contents and "model" in toml_contents[scope_name]:
            root_model = scope.get_root_model()
            model_data[scope_name] = root_model.model_validate(toml_contents[scope_name]["model"])
            logger.debug(f"Loaded model data for scope '{scope_name}'")
        else:
            logger.debug(f"No model data found for scope '{scope_name}'")

    return model_data


def load_model_data_from_toml(
    project: Project,
    input_path: Path | str,
) -> dict[str, BaseModel]:
    """Load model data from a TOML file for each scope in the project.

    FileRef paths in the TOML are resolved relative to the TOML file's directory,
    making projects portable and self-contained.

    Args:
        project: The project containing scope definitions
        input_path: Path to the input TOML file containing model data

    Returns:
        A dictionary mapping scope names to their validated root model instances

    """
    input_path = Path(input_path).resolve()
    base_dir = input_path.parent

    with input_path.open("rb") as f:
        toml_contents = tomllib.load(f)

    # Use context manager so FileRef resolves paths relative to TOML directory
    with input_base_dir(base_dir):
        model_data = toml_to_model_data(project, toml_contents)

    logger.debug(f"Loaded model data from {input_path}")
    return model_data
