"""Hierarchical token resolver with formula evaluation, audit trail, and provenance tracking."""

from __future__ import annotations

import copy
import json
import logging
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from collections.abc import Mapping
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Union

from tokenmoulds.dtcg.generator import TokenGenerator
from tokenmoulds.dtcg.errors import (
    FormulaError,
    FormulaReferenceError,
    FormulaResolutionError,
)
from tokenmoulds.dtcg.formula_evaluator import (
    FormulaEvaluator,
    is_formula,
    strip_formula,
)
from tokenmoulds.dtcg.loader import load_token_file

logger = logging.getLogger(__name__)

# Layer precedence order (lowest to highest)
LAYER_ORDER = ("core", "fork", "org", "group", "personal")


@dataclass(frozen=True)
class TokenLayer:
    """A named token layer with its payload."""

    name: str
    payload: Mapping[str, Any]
    source_path: Optional[str] = None


@dataclass
class Override:
    """Records a single value override during layer merging."""

    path: str
    from_layer: str
    to_layer: str
    old_value: Any
    new_value: Any


@dataclass
class FormulaEvaluation:
    """Records a formula evaluation."""

    path: str
    expression: str
    result: Any
    source_layer: Optional[str] = None


@dataclass
class LayerConflict:
    """Records when multiple layers define the same key."""

    path: str
    layers: List[str]
    final_layer: str
    final_value: Any


@dataclass
class ResolutionAudit:
    """Complete audit trail for a token resolution operation."""

    timestamp: datetime
    layers_used: List[str]
    layer_sources: Dict[str, Optional[str]]
    overrides: List[Override]
    formulas_evaluated: List[FormulaEvaluation]
    conflicts: List[LayerConflict]
    license_id: Optional[str] = None
    license_valid: Optional[bool] = None
    errors: List[str] = field(default_factory=list)
    warnings: List[str] = field(default_factory=list)

    def to_dict(self) -> Dict[str, Any]:
        """Convert audit to dictionary for serialization."""
        return {
            "timestamp": self.timestamp.isoformat(),
            "layers_used": self.layers_used,
            "layer_sources": self.layer_sources,
            "overrides": [
                {
                    "path": o.path,
                    "from_layer": o.from_layer,
                    "to_layer": o.to_layer,
                    "old_value": _safe_serialize(o.old_value),
                    "new_value": _safe_serialize(o.new_value),
                }
                for o in self.overrides
            ],
            "formulas_evaluated": [
                {
                    "path": f.path,
                    "expression": f.expression,
                    "result": _safe_serialize(f.result),
                    "source_layer": f.source_layer,
                }
                for f in self.formulas_evaluated
            ],
            "conflicts": [
                {
                    "path": c.path,
                    "layers": c.layers,
                    "final_layer": c.final_layer,
                    "final_value": _safe_serialize(c.final_value),
                }
                for c in self.conflicts
            ],
            "license_id": self.license_id,
            "license_valid": self.license_valid,
            "errors": self.errors,
            "warnings": self.warnings,
            "summary": {
                "total_overrides": len(self.overrides),
                "total_formulas": len(self.formulas_evaluated),
                "total_conflicts": len(self.conflicts),
                "total_errors": len(self.errors),
                "total_warnings": len(self.warnings),
            },
        }


@dataclass
class Provenance:
    """Tracks the source layer for each token value."""

    layer: str
    source_path: Optional[str] = None
    overridden_by: Optional[List[str]] = None


@dataclass
class ResolutionResult:
    """Result of token resolution including tokens, provenance, and audit."""

    tokens: Dict[str, Any]
    provenance: Dict[str, Provenance]
    audit: ResolutionAudit

    def to_dict(
        self, include_provenance: bool = True, include_audit: bool = True
    ) -> Dict[str, Any]:
        """Convert to dictionary for serialization."""
        result: Dict[str, Any] = {"tokens": self.tokens}
        if include_provenance:
            result["_provenance"] = {
                path: {
                    "layer": prov.layer,
                    "source_path": prov.source_path,
                    "overridden_by": prov.overridden_by,
                }
                for path, prov in self.provenance.items()
            }
        if include_audit:
            result["_audit"] = self.audit.to_dict()
        return result


def _safe_serialize(value: Any) -> Any:
    """Safely serialize a value for JSON output."""
    if isinstance(value, (str, int, float, bool, type(None))):
        return value
    if isinstance(value, (list, tuple)):
        return [_safe_serialize(v) for v in value]
    if isinstance(value, Mapping):
        return {k: _safe_serialize(v) for k, v in value.items()}
    return str(value)


def _deep_merge_with_tracking(
    base: Dict[str, Any],
    override: Mapping[str, Any],
    base_layer: str,
    override_layer: str,
    path_prefix: str = "",
    overrides: Optional[List[Override]] = None,
    provenance: Optional[Dict[str, Provenance]] = None,
    source_path: Optional[str] = None,
    preserve_provenance: bool = False,
) -> Dict[str, Any]:
    """Deep merge with override tracking and provenance.

    Args:
        base: Base dictionary to merge into
        override: Override dictionary to merge from
        base_layer: Name of the base layer
        override_layer: Name of the override layer
        path_prefix: Current path prefix for nested tracking
        overrides: List to append override records to
        provenance: Dict to update with provenance info
        source_path: Source file path for the override layer
        preserve_provenance: If True, don't overwrite existing provenance entries
    """
    if overrides is None:
        overrides = []
    if provenance is None:
        provenance = {}

    merged = copy.deepcopy(base)

    for key, value in override.items():
        current_path = f"{path_prefix}.{key}" if path_prefix else key

        if key in merged:
            old_value = merged[key]
            if isinstance(merged[key], dict) and isinstance(value, Mapping):
                # Recursive merge for nested dicts
                merged[key] = _deep_merge_with_tracking(
                    merged[key],
                    value,
                    base_layer,
                    override_layer,
                    current_path,
                    overrides,
                    provenance,
                    source_path,
                    preserve_provenance,
                )
            else:
                # Value override
                merged[key] = copy.deepcopy(value)

                # Skip override tracking if preserving provenance and not a real override
                if preserve_provenance and current_path in provenance:
                    continue

                overrides.append(
                    Override(
                        path=current_path,
                        from_layer=base_layer,
                        to_layer=override_layer,
                        old_value=old_value,
                        new_value=value,
                    )
                )
                # Update provenance
                existing_prov = provenance.get(current_path)
                if existing_prov:
                    if existing_prov.overridden_by is None:
                        existing_prov = Provenance(
                            layer=override_layer,
                            source_path=source_path,
                            overridden_by=[existing_prov.layer],
                        )
                    else:
                        existing_prov = Provenance(
                            layer=override_layer,
                            source_path=source_path,
                            overridden_by=existing_prov.overridden_by
                            + [existing_prov.layer],
                        )
                    provenance[current_path] = existing_prov
                else:
                    provenance[current_path] = Provenance(
                        layer=override_layer,
                        source_path=source_path,
                        overridden_by=[base_layer],
                    )
        else:
            # New key
            merged[key] = copy.deepcopy(value)

            # Skip provenance tracking if preserving and already tracked
            if preserve_provenance and current_path in provenance:
                continue

            if isinstance(value, Mapping):
                # Track provenance for nested values
                _track_nested_provenance(
                    value,
                    current_path,
                    override_layer,
                    source_path,
                    provenance,
                    preserve_existing=preserve_provenance,
                )
            else:
                if current_path not in provenance:
                    provenance[current_path] = Provenance(
                        layer=override_layer, source_path=source_path
                    )

    return merged


def _track_nested_provenance(
    data: Mapping[str, Any],
    path_prefix: str,
    layer: str,
    source_path: Optional[str],
    provenance: Dict[str, Provenance],
    preserve_existing: bool = False,
) -> None:
    """Recursively track provenance for nested structures."""
    for key, value in data.items():
        current_path = f"{path_prefix}.{key}"
        if isinstance(value, Mapping):
            _track_nested_provenance(
                value, current_path, layer, source_path, provenance, preserve_existing
            )
        else:
            # Don't overwrite existing provenance if preserving
            if not preserve_existing or current_path not in provenance:
                provenance[current_path] = Provenance(
                    layer=layer, source_path=source_path
                )


def _deep_merge(base: Dict[str, Any], override: Mapping[str, Any]) -> Dict[str, Any]:
    """Simple deep merge without tracking (for backward compatibility)."""
    merged = copy.deepcopy(base)
    for key, value in override.items():
        if (
            key in merged
            and isinstance(merged[key], dict)
            and isinstance(value, Mapping)
        ):
            merged[key] = _deep_merge(merged[key], value)
        else:
            merged[key] = copy.deepcopy(value)
    return merged


def _iter_formula_paths(
    node: Any, path: Tuple[Any, ...] = ()
) -> Iterable[Tuple[Any, ...]]:
    if is_formula(node):
        yield path
        return
    if isinstance(node, Mapping):
        for key, value in node.items():
            yield from _iter_formula_paths(value, path + (key,))
    elif isinstance(node, list):
        for idx, value in enumerate(node):
            yield from _iter_formula_paths(value, path + (idx,))


def _get_path(data: Any, path: Tuple[Any, ...]) -> Any:
    value = data
    for key in path:
        value = value[key]
    return value


def _set_path(data: Any, path: Tuple[Any, ...], value: Any) -> None:
    current = data
    for key in path[:-1]:
        current = current[key]
    current[path[-1]] = value


def _path_to_str(path: Tuple[Any, ...]) -> str:
    return ".".join(str(p) for p in path)


def _detect_conflicts(layers: Sequence[TokenLayer]) -> List[LayerConflict]:
    """Detect keys that are defined in multiple layers."""
    key_layers: Dict[str, List[Tuple[str, Any]]] = {}

    def collect_keys(
        data: Mapping[str, Any], layer_name: str, path_prefix: str = ""
    ) -> None:
        for key, value in data.items():
            current_path = f"{path_prefix}.{key}" if path_prefix else key
            if isinstance(value, Mapping):
                collect_keys(value, layer_name, current_path)
            else:
                if current_path not in key_layers:
                    key_layers[current_path] = []
                key_layers[current_path].append((layer_name, value))

    for layer in layers:
        collect_keys(layer.payload, layer.name)

    conflicts = []
    for path, layer_values in key_layers.items():
        if len(layer_values) > 1:
            layers_list = [lv[0] for lv in layer_values]
            final_layer, final_value = layer_values[-1]
            conflicts.append(
                LayerConflict(
                    path=path,
                    layers=layers_list,
                    final_layer=final_layer,
                    final_value=final_value,
                )
            )

    return conflicts


class TokenResolver:
    """Resolve layered token sets with formula evaluation, audit trail, and provenance tracking."""

    def __init__(
        self,
        evaluator: FormulaEvaluator | None = None,
        generator: TokenGenerator | None = None,
        *,
        use_engine: bool = True,
        license_validator: Any = None,
        layer_validator: Any = None,
    ) -> None:
        self._evaluator = evaluator or FormulaEvaluator()
        self._generator = generator or TokenGenerator()
        self._use_engine = use_engine
        self._license_validator = license_validator
        self._layer_validator = layer_validator

    def resolve_layers(
        self,
        layers: Sequence[TokenLayer],
        *,
        track_provenance: bool = True,
        track_audit: bool = True,
        license_context: Any = None,
        validate_layers: bool = False,
        strict_validation: bool = False,
    ) -> Union[Dict[str, Any], ResolutionResult]:
        """Resolve layered tokens with optional provenance and audit tracking.

        Args:
            layers: Sequence of TokenLayer objects in precedence order (lowest first)
            track_provenance: If True, return ResolutionResult with provenance
            track_audit: If True, include audit trail in result
            license_context: Optional LicenseContext for validation
            validate_layers: If True, validate each layer against schema before merging
            strict_validation: If True, treat validation warnings as errors

        Returns:
            If track_provenance is False: Dict of resolved tokens (backward compatible)
            If track_provenance is True: ResolutionResult with tokens, provenance, and audit
        """
        audit = ResolutionAudit(
            timestamp=datetime.now(timezone.utc),
            layers_used=[layer.name for layer in layers],
            layer_sources={layer.name: layer.source_path for layer in layers},
            overrides=[],
            formulas_evaluated=[],
            conflicts=[],
        )

        # Layer validation
        if validate_layers:
            self._validate_layers(layers, strict_validation, audit)

        # License validation
        if license_context and self._license_validator:
            try:
                result = self._license_validator.validate(license_context)
                audit.license_id = result.license_id
                audit.license_valid = result.allowed

                # Filter layers based on license permissions
                if result.layers:
                    allowed_layers = []
                    for layer in layers:
                        if result.layers.get(layer.name, False):
                            allowed_layers.append(layer)
                        else:
                            audit.warnings.append(
                                f"Layer '{layer.name}' not permitted by license '{result.license_id}'"
                            )
                    layers = allowed_layers

                if not result.allowed:
                    audit.errors.append(
                        f"License validation failed: {result.license_id}"
                    )
            except Exception as e:
                audit.errors.append(f"License validation error: {e}")
                logger.warning(f"License validation failed: {e}")

        # Detect conflicts before merging
        if track_audit:
            audit.conflicts = _detect_conflicts(layers)
            for conflict in audit.conflicts:
                audit.warnings.append(
                    f"Key '{conflict.path}' defined in multiple layers: {conflict.layers} "
                    f"(using '{conflict.final_layer}')"
                )

        # Merge layers with tracking
        provenance: Dict[str, Provenance] = {}
        merged: Dict[str, Any] = {}
        prev_layer = "base"

        for layer in layers:
            if track_provenance:
                merged = _deep_merge_with_tracking(
                    merged,
                    layer.payload,
                    prev_layer,
                    layer.name,
                    overrides=audit.overrides,
                    provenance=provenance,
                    source_path=layer.source_path,
                )
            else:
                merged = _deep_merge(merged, layer.payload)
            prev_layer = layer.name

        # Resolve inputs formulas
        inputs = merged.get("inputs", {})
        if isinstance(inputs, Mapping) and inputs:
            resolved_inputs, input_formulas = self._resolve_formulas_with_tracking(
                {"inputs": inputs}, "inputs"
            )
            resolved_inputs = resolved_inputs.get("inputs", {})
            audit.formulas_evaluated.extend(input_formulas)
        else:
            resolved_inputs = {}

        # Generate derived tokens
        generated: Dict[str, Any] = {}
        if (
            self._use_engine
            and isinstance(resolved_inputs, Mapping)
            and resolved_inputs
        ):
            generated = self._generator.generate(resolved_inputs)
            # Track generated tokens provenance
            if track_provenance:
                _track_nested_provenance(generated, "", "generated", None, provenance)

        # Merge generated with overrides
        overrides = dict(merged)
        if resolved_inputs:
            overrides["inputs"] = resolved_inputs

        if track_provenance:
            # When merging overrides, preserve existing provenance from layers
            # Only update provenance for values that come from generated tokens
            resolved = _deep_merge_with_tracking(
                generated,
                overrides,
                "generated",
                "overrides",
                overrides=audit.overrides,
                provenance=provenance,
                preserve_provenance=True,  # Don't overwrite layer provenance with "overrides"
            )
        else:
            resolved = _deep_merge(generated, overrides)

        # Resolve remaining formulas
        resolved, final_formulas = self._resolve_formulas_with_tracking(
            resolved, "final"
        )
        audit.formulas_evaluated.extend(final_formulas)

        # Log audit
        self._log_audit(audit)

        if track_provenance:
            return ResolutionResult(
                tokens=resolved,
                provenance=provenance,
                audit=audit,
            )
        return resolved

    def resolve_files(
        self,
        files: Sequence[Tuple[str, Path]],
        *,
        track_provenance: bool = False,
        track_audit: bool = True,
        license_context: Any = None,
    ) -> Union[Dict[str, Any], ResolutionResult]:
        """Resolve tokens from file paths."""
        layers: List[TokenLayer] = []
        for name, path in files:
            payload = load_token_file(path)
            layers.append(TokenLayer(name=name, payload=payload, source_path=str(path)))
        return self.resolve_layers(
            layers,
            track_provenance=track_provenance,
            track_audit=track_audit,
            license_context=license_context,
        )

    def _resolve_formulas_with_tracking(
        self,
        data: Mapping[str, Any],
        phase: str = "",
    ) -> Tuple[Dict[str, Any], List[FormulaEvaluation]]:
        """Resolve formulas and track evaluations."""
        resolved = copy.deepcopy(data)
        formula_evals: List[FormulaEvaluation] = []
        pending = list(_iter_formula_paths(resolved))

        if not pending:
            return dict(resolved), formula_evals

        context = self._build_context(resolved)
        max_passes = max(5, len(pending) * 2)

        for _ in range(max_passes):
            progress = False
            next_pending: List[Tuple[Any, ...]] = []

            for path in pending:
                value = _get_path(resolved, path)
                if not is_formula(value):
                    progress = True
                    continue

                expression = strip_formula(value)
                try:
                    evaluated = self._evaluator.evaluate(expression, context)
                except FormulaReferenceError:
                    next_pending.append(path)
                    continue
                except FormulaError as exc:
                    raise FormulaResolutionError(
                        f"Failed to evaluate formula at {_path_to_str(path)}: {exc}"
                    ) from exc

                _set_path(resolved, path, evaluated)
                context = self._build_context(resolved)
                progress = True

                formula_evals.append(
                    FormulaEvaluation(
                        path=_path_to_str(path),
                        expression=expression,
                        result=evaluated,
                        source_layer=phase,
                    )
                )

            pending = next_pending
            if not pending:
                return dict(resolved), formula_evals
            if not progress:
                break

        if pending:
            unresolved = ", ".join(_path_to_str(path) for path in pending)
            raise FormulaResolutionError(
                f"Unresolved formulas after evaluation: {unresolved}"
            )

        return dict(resolved), formula_evals

    def _resolve_formulas(self, data: Mapping[str, Any]) -> Dict[str, Any]:
        """Resolve formulas without tracking (backward compatibility)."""
        resolved, _ = self._resolve_formulas_with_tracking(data)
        return resolved

    @staticmethod
    def _build_context(data: Mapping[str, Any]) -> Dict[str, Any]:
        context = dict(data)
        inputs = data.get("inputs")
        if isinstance(inputs, Mapping):
            for key, value in inputs.items():
                context.setdefault(key, value)
        return context

    def _validate_layers(
        self,
        layers: Sequence[TokenLayer],
        strict: bool,
        audit: ResolutionAudit,
    ) -> Dict[str, Any]:
        """Validate layers against schema before merging."""
        results = {}

        # Use provided validator or create one
        if self._layer_validator:
            validator = self._layer_validator
        else:
            try:
                from tokenmoulds.dtcg.layer_validator import LayerValidator

                validator = LayerValidator(strict=strict)
            except ImportError:
                audit.warnings.append(
                    "Layer validation skipped: layer_validator not available"
                )
                return results

        for layer in layers:
            result = validator.validate(layer.payload, layer_name=layer.name)
            results[layer.name] = result

            # Add validation issues to audit
            for issue in result.issues:
                if issue.severity == "error":
                    audit.errors.append(f"[{layer.name}] {issue.path}: {issue.message}")
                elif issue.severity == "warning":
                    audit.warnings.append(
                        f"[{layer.name}] {issue.path}: {issue.message}"
                    )

            if not result.is_valid:
                logger.warning(
                    f"Layer '{layer.name}' validation failed with {len(result.errors)} errors"
                )

        return results

    def _log_audit(self, audit: ResolutionAudit) -> None:
        """Log audit trail for compliance."""
        log_entry = {
            "event": "token.resolution",
            "timestamp": audit.timestamp.isoformat(),
            "layers": audit.layers_used,
            "license_id": audit.license_id,
            "overrides_count": len(audit.overrides),
            "formulas_count": len(audit.formulas_evaluated),
            "conflicts_count": len(audit.conflicts),
            "warnings_count": len(audit.warnings),
            "errors_count": len(audit.errors),
        }
        logger.info(json.dumps(log_entry, sort_keys=True))

        if audit.warnings:
            for warning in audit.warnings:
                logger.warning(f"Token resolution warning: {warning}")

        if audit.errors:
            for error in audit.errors:
                logger.error(f"Token resolution error: {error}")
