"""Snapshot testing engine for MCP tool responses.

Stores golden-file snapshots as JSON files in __snapshots__/ directories
alongside test files. On first run, saves the response. On subsequent runs,
diffs against the saved snapshot.
"""

from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Any

from deepdiff import DeepDiff

from mcp_forge._types import ToolInfo, ToolResult


class SnapshotMismatchError(AssertionError):
    """Raised when a snapshot doesn't match the actual result."""

    def __init__(self, message: str, diff: dict[str, Any] | None = None) -> None:
        self.diff = diff
        super().__init__(message)


class SnapshotManager:
    """Manages golden-file snapshots for MCP tool responses.

    Snapshots are stored as JSON files in a __snapshots__/ directory
    next to the test file.

    Usage in tests:
        async def test_something(client, forge_snapshot):
            result = await client.call_tool("my_tool", {"x": 1})
            forge_snapshot.assert_match(result)
    """

    def __init__(
        self,
        snap_dir: str | Path,
        test_name: str,
        update: bool = False,
    ) -> None:
        self.snap_dir = Path(snap_dir)
        self.test_name = test_name
        self.update = update
        self._counter = 0

    def _get_snap_path(self, suffix: str = "") -> Path:
        """Generate the snapshot file path for the current assertion."""
        self._counter += 1
        safe_name = self.test_name.replace("[", "_").replace("]", "_").replace("/", "_")
        filename = f"{safe_name}"
        if self._counter > 1 or suffix:
            filename += f"_{suffix or self._counter}"
        filename += ".json"
        return self.snap_dir / filename

    def _save_snapshot(self, path: Path, data: dict[str, Any]) -> None:
        """Save a snapshot to disk."""
        path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=2, ensure_ascii=False, default=str)
            f.write("\n")

    def _load_snapshot(self, path: Path) -> dict[str, Any]:
        """Load a snapshot from disk."""
        with open(path, "r", encoding="utf-8") as f:
            return json.load(f)

    def _filter_keys(
        self, data: dict[str, Any], ignore_keys: list[str]
    ) -> dict[str, Any]:
        """Recursively remove ignored keys from a dict."""
        if not isinstance(data, dict):
            return data

        filtered = {}
        for key, value in data.items():
            if key in ignore_keys:
                continue
            if isinstance(value, dict):
                filtered[key] = self._filter_keys(value, ignore_keys)
            elif isinstance(value, list):
                filtered[key] = [
                    self._filter_keys(item, ignore_keys)
                    if isinstance(item, dict)
                    else item
                    for item in value
                ]
            else:
                filtered[key] = value
        return filtered

    def assert_match(
        self,
        result: ToolResult,
        *,
        ignore_keys: list[str] | None = None,
        approximate_numbers: bool = False,
        suffix: str = "",
    ) -> None:
        """Compare a tool result against a saved snapshot.

        First run: saves the snapshot and passes.
        Subsequent runs: diffs against the saved snapshot.

        Args:
            result: The ToolResult to snapshot.
            ignore_keys: Keys to exclude from comparison (e.g., ["timestamp"]).
            approximate_numbers: Allow ±1% float drift.
            suffix: Optional suffix for the snapshot filename.

        Raises:
            SnapshotMismatchError: If the result doesn't match the snapshot.
        """
        snap_path = self._get_snap_path(suffix)
        actual_data = result.to_dict()

        if ignore_keys:
            actual_data = self._filter_keys(actual_data, ignore_keys)

        # If updating or first run, save and pass
        if self.update or not snap_path.exists():
            self._save_snapshot(snap_path, actual_data)
            return

        # Load saved snapshot
        expected_data = self._load_snapshot(snap_path)

        if ignore_keys:
            expected_data = self._filter_keys(expected_data, ignore_keys)

        # Compare
        diff_kwargs: dict[str, Any] = {}
        if approximate_numbers:
            diff_kwargs["significant_digits"] = 2
            diff_kwargs["math_epsilon"] = 0.01

        diff = DeepDiff(expected_data, actual_data, **diff_kwargs)

        if diff:
            diff_dict = dict(diff)
            lines = [
                f"Snapshot mismatch for '{self.test_name}'.",
                f"Snapshot file: {snap_path}",
                "",
                "Differences:",
            ]

            if "values_changed" in diff_dict:
                lines.append("  Changed values:")
                for path, change in diff_dict["values_changed"].items():
                    lines.append(
                        f"    {path}: {change['old_value']!r} → {change['new_value']!r}"
                    )

            if "dictionary_item_added" in diff_dict:
                lines.append("  Added keys:")
                for item in diff_dict["dictionary_item_added"]:
                    lines.append(f"    + {item}")

            if "dictionary_item_removed" in diff_dict:
                lines.append("  Removed keys:")
                for item in diff_dict["dictionary_item_removed"]:
                    lines.append(f"    - {item}")

            if "iterable_item_added" in diff_dict:
                lines.append("  Added items:")
                for path, val in diff_dict["iterable_item_added"].items():
                    lines.append(f"    + {path}: {val!r}")

            if "iterable_item_removed" in diff_dict:
                lines.append("  Removed items:")
                for path, val in diff_dict["iterable_item_removed"].items():
                    lines.append(f"    - {path}: {val!r}")

            if "type_changes" in diff_dict:
                lines.append("  Type changes:")
                for path, change in diff_dict["type_changes"].items():
                    lines.append(
                        f"    {path}: {change['old_type'].__name__} → {change['new_type'].__name__}"
                    )

            lines.append("")
            lines.append("Run with --update-snapshots to accept the new output.")

            raise SnapshotMismatchError("\n".join(lines), diff=diff_dict)

    def assert_schema_unchanged(
        self, tools: list[ToolInfo], *, suffix: str = "schema"
    ) -> None:
        """Snapshot the entire tool schema surface.

        Detects: new tools added, tools removed, parameter changes,
        description changes, type changes.
        """
        snap_path = self._get_snap_path(suffix)

        schema_data = {
            "tools": [t.to_dict() for t in sorted(tools, key=lambda t: t.name)]
        }

        if self.update or not snap_path.exists():
            self._save_snapshot(snap_path, schema_data)
            return

        expected_data = self._load_snapshot(snap_path)
        diff = DeepDiff(expected_data, schema_data)

        if diff:
            diff_dict = dict(diff)

            # Generate human-friendly summary
            lines = ["Tool schema has changed:"]

            expected_names = {
                t["name"] for t in expected_data.get("tools", [])
            }
            actual_names = {t["name"] for t in schema_data.get("tools", [])}

            added = actual_names - expected_names
            removed = expected_names - actual_names

            if added:
                lines.append(f"  New tools: {added}")
            if removed:
                lines.append(f"  Removed tools: {removed}")

            if "values_changed" in diff_dict:
                lines.append("  Changed values:")
                for path, change in diff_dict["values_changed"].items():
                    lines.append(
                        f"    {path}: {change['old_value']!r} → {change['new_value']!r}"
                    )

            lines.append("")
            lines.append("Run with --update-snapshots to accept the new schema.")

            raise SnapshotMismatchError("\n".join(lines), diff=diff_dict)

    def assert_raw_match(
        self,
        data: dict[str, Any],
        *,
        ignore_keys: list[str] | None = None,
        suffix: str = "",
    ) -> None:
        """Snapshot any arbitrary dict (not just ToolResult).

        Useful for testing resources, prompts, or custom structures.
        """
        snap_path = self._get_snap_path(suffix or "raw")

        actual_data = data
        if ignore_keys:
            actual_data = self._filter_keys(actual_data, ignore_keys)

        if self.update or not snap_path.exists():
            self._save_snapshot(snap_path, actual_data)
            return

        expected_data = self._load_snapshot(snap_path)
        if ignore_keys:
            expected_data = self._filter_keys(expected_data, ignore_keys)

        diff = DeepDiff(expected_data, actual_data)
        if diff:
            raise SnapshotMismatchError(
                f"Raw snapshot mismatch.\n"
                f"Snapshot: {snap_path}\n"
                f"Diff: {json.dumps(dict(diff), indent=2, default=str)}\n\n"
                f"Run with --update-snapshots to accept.",
                diff=dict(diff),
            )
