"""MCP server management."""

import copy
import json
import shutil

from abc import ABC, abstractmethod
from collections.abc import Mapping
from enum import Enum
from pathlib import Path
from typing import Any, ClassVar, cast

import yaml

from .config import Config, write_file_atomic
from .utils import deep_merge


class OperationResult(Enum):
    CREATED = "created"
    UPDATED = "updated"
    REMOVED = "removed"
    ALREADY_INSTALLED = "already_installed"
    NOT_FOUND = "not_found"
    ERROR = "error"


class MCPStatus:
    def __init__(self) -> None:
        self.managed_mcps: dict[str, dict[str, Any]] = {}
        self.unmanaged_mcps: dict[str, dict[str, Any]] = {}
        self.pending_mcps: dict[str, dict[str, Any]] = {}
        self.stale_mcps: dict[str, dict[str, Any]] = {}
        self.installed: dict[str, bool] = {}
        self.has_overrides: dict[str, bool] = {}


MANAGED_BY_VALUE = "ai-agent-rules"
_MANAGED_BY_VALUE_LEGACY = "ai-rules"


def is_managed_value(v: str | None) -> bool:
    """Check if a value matches the current or legacy managed-by marker."""
    return v in (MANAGED_BY_VALUE, _MANAGED_BY_VALUE_LEGACY)


class MCPManager(ABC):
    """Abstract base for agent-specific MCP config managers."""

    # Default translation: copy these keys when present, then apply extras.
    # Managers with a genuinely different native shape override _translate.
    _translate_keys: ClassVar[tuple[str, ...]] = ("command", "args", "env")
    _translate_extras: ClassVar[Mapping[str, Any]] = {}

    # --- abstract interface ---------------------------------------------------

    @property
    @abstractmethod
    def _marker_field(self) -> str:
        """Field name used to mark managed MCP entries."""

    @abstractmethod
    def _read_installed(self) -> dict[str, Any]:
        """Return the currently installed MCP entries as {name: config}."""

    @abstractmethod
    def _write_installed(self, mcps: dict[str, Any]) -> None:
        """Persist the full set of MCP entries, merging with non-MCP config."""

    def _translate(self, shared_config: dict[str, Any]) -> dict[str, Any]:
        """Convert shared MCP format to agent-native format."""
        result: dict[str, Any] = {
            k: shared_config[k] for k in self._translate_keys if k in shared_config
        }
        result.update(self._translate_extras)
        return result

    @property
    def mcp_settings_key(self) -> str | None:
        """Key in the settings dict where MCPs are stored.

        Returns None if this manager stores MCPs in a separate file
        (e.g. Claude stores MCPs in ~/.claude.json, not in settings.json).
        """
        return None

    @property
    def mcp_tracking_key(self) -> str | None:
        """Key for the managed-names tracking section, if any."""
        return None

    @property
    def _previously_managed_names(self) -> frozenset[str]:
        """Names of MCPs historically managed by ai-agent-rules.

        Entries matching these names are eligible for cleanup even without
        the managed-by marker — handles cases where entries were added
        externally or predate marker-stamping.
        """
        return frozenset()

    def get_native_mcps(self, config_dir: Path, config: Config) -> dict[str, Any]:
        """Load managed MCPs, translate to native format, and stamp marker."""
        managed = self.load_managed_mcps(config_dir, config)
        result: dict[str, Any] = {}
        for name, shared_cfg in managed.items():
            native = self._translate(shared_cfg)
            native[self._marker_field] = MANAGED_BY_VALUE
            result[name] = native
        return result

    # --- shared logic --------------------------------------------------------

    def load_managed_mcps(self, config_dir: Path, config: Config) -> dict[str, Any]:
        """Load managed MCP definitions and apply user overrides.

        Prefers config/mcps.json (shared); falls back to config/claude/mcps.json
        for backward compatibility.
        """
        shared_file = config_dir / "mcps.json"
        legacy_file = config_dir / "claude" / "mcps.json"

        if shared_file.exists():
            with open(shared_file, encoding="utf-8") as f:
                base_mcps: dict[str, Any] = json.load(f)
        elif legacy_file.exists():
            with open(legacy_file, encoding="utf-8") as f:
                base_mcps = json.load(f)
        else:
            return {}

        mcp_overrides = config.mcp_overrides if hasattr(config, "mcp_overrides") else {}

        merged_mcps: dict[str, Any] = {}
        for name in {**base_mcps, **mcp_overrides}:
            if name in base_mcps and name in mcp_overrides:
                merged_mcps[name] = deep_merge(base_mcps[name], mcp_overrides[name])
            elif name in base_mcps:
                merged_mcps[name] = copy.deepcopy(base_mcps[name])
            else:
                merged_mcps[name] = copy.deepcopy(mcp_overrides[name])

        return merged_mcps

    def detect_conflicts(
        self, expected: dict[str, Any], installed: dict[str, Any]
    ) -> list[str]:
        """Detect MCPs that have been modified locally."""
        conflicts = []
        marker = self._marker_field
        for name, expected_config in expected.items():
            if name in installed:
                expected_clean = {
                    k: v for k, v in expected_config.items() if k != marker
                }
                installed_clean = {
                    k: v for k, v in installed[name].items() if k != marker
                }
                if expected_clean != installed_clean:
                    conflicts.append(name)
        return conflicts

    def format_diff(
        self, name: str, expected: dict[str, Any], installed: dict[str, Any]
    ) -> str | None:
        """Format a unified diff between expected and installed MCP config."""
        from ai_rules.symlinks import format_unified_diff

        expected_json = json.dumps(expected, indent=2, sort_keys=True)
        installed_json = json.dumps(installed, indent=2, sort_keys=True)

        diff = format_unified_diff(
            installed_json.splitlines(keepends=True),
            expected_json.splitlines(keepends=True),
            "Installed (local)",
            "Expected (repo)",
        )
        if diff is None:
            return None
        return f"    MCP '{name}' has been modified locally:\n{diff}"

    def format_pending(self, name: str, expected: dict[str, Any]) -> str:
        """Format expected MCP config for pending installation."""
        from ai_rules.cli.display import dim

        marker = self._marker_field
        display_config = {k: v for k, v in expected.items() if k != marker}
        config_json = json.dumps(display_config, indent=2, sort_keys=True)

        lines = [f"    {dim('Will be installed with:')}"]
        for line in config_json.splitlines():
            lines.append(f"[green]    {line}[/green]")
        return "\n".join(lines)

    def _create_backup(self, target: Path) -> Path | None:
        """Create a timestamped backup of a config file."""
        if not target.exists():
            return None

        from .symlinks import create_backup_path

        backup_path = create_backup_path(target)
        shutil.copy2(target, backup_path)
        return backup_path

    def install_mcps(
        self,
        config_dir: Path,
        config: Config,
        force: bool = False,
        dry_run: bool = False,
    ) -> tuple[OperationResult, str, list[str]]:
        """Install managed MCPs into the agent's config file."""
        native_mcps = self.get_native_mcps(config_dir, config)

        current_mcps = self._read_installed()
        original_mcps = copy.deepcopy(current_mcps)

        tracked_mcps = {
            name
            for name, cfg in current_mcps.items()
            if is_managed_value(cfg.get(self._marker_field))
        }
        unmarked_orphans = {
            name
            for name in self._previously_managed_names
            if name in current_mcps and name not in native_mcps
        }
        removed_mcps = (tracked_mcps - set(native_mcps.keys())) | unmarked_orphans

        if not native_mcps and not removed_mcps:
            return (OperationResult.NOT_FOUND, "No MCPs to install or remove", [])

        conflicts = self.detect_conflicts(native_mcps, current_mcps)

        if conflicts and not force:
            return (
                OperationResult.ERROR,
                "Conflicts detected (use -y to override)",
                conflicts,
            )

        if dry_run:
            msg = f"Would update {len(native_mcps)} MCPs, remove {len(removed_mcps)}"
            return (OperationResult.UPDATED, msg, conflicts)

        for name in removed_mcps:
            current_mcps.pop(name, None)

        current_mcps.update(native_mcps)

        if current_mcps == original_mcps:
            return (OperationResult.ALREADY_INSTALLED, "MCPs already up to date", [])

        self._write_installed(current_mcps)

        parts = []
        if native_mcps:
            parts.append(f"installed {len(native_mcps)}")
        if removed_mcps:
            parts.append(f"removed {len(removed_mcps)}")
        msg = f"MCPs {', '.join(parts)}"

        return (OperationResult.UPDATED, msg, [])

    def uninstall_mcps(self) -> tuple[OperationResult, str]:
        """Uninstall managed MCPs from the agent's config file."""
        current_mcps = self._read_installed()

        tracked_mcps = {
            name
            for name, cfg in current_mcps.items()
            if is_managed_value(cfg.get(self._marker_field))
        }
        previously_managed = {
            name for name in self._previously_managed_names if name in current_mcps
        }
        all_to_remove = tracked_mcps | previously_managed

        if not all_to_remove:
            return (OperationResult.NOT_FOUND, "No tracked MCPs found")

        for name in all_to_remove:
            current_mcps.pop(name, None)

        self._write_installed(current_mcps)

        return (OperationResult.REMOVED, f"Removed {len(all_to_remove)} MCPs")

    def get_status(self, config_dir: Path, config: Config) -> MCPStatus:
        """Get status of managed and unmanaged MCPs."""
        status = MCPStatus()
        managed_mcps = self.load_managed_mcps(config_dir, config)

        native_mcps: dict[str, Any] = {}
        for name, shared_cfg in managed_mcps.items():
            translated = self._translate(shared_cfg)
            translated[self._marker_field] = MANAGED_BY_VALUE
            native_mcps[name] = translated

        installed_mcps = self._read_installed()
        mcp_overrides = config.mcp_overrides if hasattr(config, "mcp_overrides") else {}
        marker = self._marker_field

        for name, mcp_config in installed_mcps.items():
            if is_managed_value(mcp_config.get(marker)):
                if name in native_mcps:
                    status.managed_mcps[name] = mcp_config
                    expected_clean = {
                        k: v for k, v in native_mcps[name].items() if k != marker
                    }
                    installed_clean = {
                        k: v for k, v in mcp_config.items() if k != marker
                    }
                    status.installed[name] = expected_clean == installed_clean
                    status.has_overrides[name] = name in mcp_overrides
                else:
                    status.stale_mcps[name] = mcp_config
            else:
                if name in self._previously_managed_names and name not in native_mcps:
                    status.stale_mcps[name] = mcp_config
                else:
                    status.unmanaged_mcps[name] = mcp_config

        for name, mcp_config in native_mcps.items():
            if name not in installed_mcps:
                status.pending_mcps[name] = mcp_config
                status.has_overrides[name] = name in mcp_overrides

        return status


# ---------------------------------------------------------------------------
# Claude
# ---------------------------------------------------------------------------


class ClaudeMCPManager(MCPManager):
    """Manages MCPs in ~/.claude.json under the mcpServers key."""

    @property
    def _marker_field(self) -> str:
        return "_managedBy"

    @property
    def _claude_json_path(self) -> Path:
        return Path.home() / ".claude.json"

    def _read_installed(self) -> dict[str, Any]:
        if not self._claude_json_path.exists():
            return {}
        with open(self._claude_json_path, encoding="utf-8") as f:
            data: dict[str, Any] = json.load(f)
        return cast(dict[str, Any], data.get("mcpServers", {}))

    def _write_installed(self, mcps: dict[str, Any]) -> None:
        self._claude_json_path.parent.mkdir(parents=True, exist_ok=True)

        if self._claude_json_path.exists():
            with open(self._claude_json_path, encoding="utf-8") as f:
                data: dict[str, Any] = json.load(f)
        else:
            data = {}

        data.setdefault("mcpServers", {})
        data["mcpServers"] = mcps

        self._write_json_atomic(self._claude_json_path, data)

    def _translate(self, shared_config: dict[str, Any]) -> dict[str, Any]:
        return {
            k: v for k, v in shared_config.items() if k not in ("description", "name")
        }

    def _write_json_atomic(self, path: Path, data: dict[str, Any]) -> None:
        write_file_atomic(path, lambda f: json.dump(data, f, indent=2, sort_keys=True))

    @property
    def CLAUDE_JSON(self) -> Path:
        return self._claude_json_path

    def install_mcps(
        self,
        config_dir: Path,
        config: Config,
        force: bool = False,
        dry_run: bool = False,
    ) -> tuple[OperationResult, str, list[str]]:
        """Install MCPs, creating a timestamped backup of ~/.claude.json first."""
        if not dry_run and self._claude_json_path.exists():
            self._create_backup(self._claude_json_path)
        result, msg, conflicts = super().install_mcps(
            config_dir, config, force, dry_run
        )
        return result, msg, conflicts

    def uninstall_mcps(self) -> tuple[OperationResult, str]:
        if self._claude_json_path.exists():
            backup = self._create_backup(self._claude_json_path)
        else:
            backup = None
        result, msg = super().uninstall_mcps()
        if backup:
            msg += f" (backup: {backup})"
        return result, msg


# ---------------------------------------------------------------------------
# Goose
# ---------------------------------------------------------------------------


class GooseMCPManager(MCPManager):
    """Manages MCPs in ~/.config/goose/config.yaml under the extensions key."""

    @property
    def _marker_field(self) -> str:
        return "_managed_by"

    @property
    def _previously_managed_names(self) -> frozenset[str]:
        from ai_rules.bootstrap.registry import get_deprecated_mcp_names

        return get_deprecated_mcp_names()

    @property
    def mcp_settings_key(self) -> str | None:
        return "extensions"

    @property
    def _config_path(self) -> Path:
        from ai_rules.platform import get_goose_config_dir

        return get_goose_config_dir().expanduser() / "config.yaml"

    def _load_full_config(self) -> dict[str, Any]:
        if not self._config_path.exists():
            return {}
        with open(self._config_path, encoding="utf-8") as f:
            result = yaml.safe_load(f)
        return cast(dict[str, Any], result) if result else {}

    def _read_installed(self) -> dict[str, Any]:
        full = self._load_full_config()
        extensions = full.get("extensions", {})
        if not isinstance(extensions, dict):
            return {}
        return cast(dict[str, Any], extensions)

    def _write_installed(self, mcps: dict[str, Any]) -> None:
        self._config_path.parent.mkdir(parents=True, exist_ok=True)
        full = self._load_full_config()
        full["extensions"] = mcps
        with open(self._config_path, "w", encoding="utf-8") as f:
            yaml.safe_dump(full, f, default_flow_style=False, sort_keys=True)

    def _translate(self, shared_config: dict[str, Any]) -> dict[str, Any]:
        return {
            "type": "stdio",
            "cmd": shared_config.get("command", ""),
            "args": shared_config.get("args", []),
            "envs": shared_config.get("env", {}),
            "env_keys": [],
            "enabled": True,
            "name": shared_config.get("name", ""),
            "description": shared_config.get("description", ""),
            "timeout": 300,
            "bundled": False,
            "available_tools": [],
        }


# ---------------------------------------------------------------------------
# Codex
# ---------------------------------------------------------------------------


class CodexMCPManager(MCPManager):
    """Manages MCPs in ~/.codex/config.toml under the mcp_servers section.

    Uses tomlkit for round-trip editing so non-MCP keys (approval_policy,
    trust_level, model) and comments are preserved.

    The managed-names list is stored in a dedicated TOML section
    because TOML inline per-table markers are awkward.
    """

    _MANAGED_SECTION = "_ai_agent_rules_managed"
    _LEGACY_MANAGED_SECTION = "_ai_rules_managed"

    @property
    def _marker_field(self) -> str:
        return "_ai_agent_rules_managed_entry"

    @property
    def mcp_settings_key(self) -> str | None:
        return "mcp_servers"

    @property
    def mcp_tracking_key(self) -> str | None:
        return self._MANAGED_SECTION

    @property
    def _config_path(self) -> Path:
        return Path.home() / ".codex" / "config.toml"

    def _load_doc(self) -> Any:
        import tomlkit

        if not self._config_path.exists():
            return tomlkit.document()
        with open(self._config_path, encoding="utf-8") as f:
            return tomlkit.load(f)

    def _get_managed_names(self, doc: Any) -> set[str]:
        section = doc.get(self._MANAGED_SECTION) or doc.get(
            self._LEGACY_MANAGED_SECTION, {}
        )
        names = section.get("names", [])
        return set(names) if isinstance(names, list) else set()

    def _read_installed(self) -> dict[str, Any]:
        import tomlkit

        doc = self._load_doc()
        managed_names = self._get_managed_names(doc)
        mcp_servers = doc.get("mcp_servers", tomlkit.table())
        if not isinstance(mcp_servers, dict):
            return {}

        result: dict[str, Any] = {}
        for name, cfg in mcp_servers.items():
            entry = dict(cfg) if hasattr(cfg, "items") else {}
            if name in managed_names:
                entry[self._marker_field] = MANAGED_BY_VALUE
            result[name] = entry
        return result

    def _write_installed(self, mcps: dict[str, Any]) -> None:
        import tomlkit

        self._config_path.parent.mkdir(parents=True, exist_ok=True)
        doc = self._load_doc()

        managed_names: list[str] = []
        mcp_table = tomlkit.table()

        for name, cfg in mcps.items():
            is_managed = is_managed_value(cfg.get(self._marker_field))
            entry = tomlkit.table()
            for k, v in cfg.items():
                if k == self._marker_field:
                    continue
                if isinstance(v, dict):
                    sub = tomlkit.table()
                    for sk, sv in v.items():
                        sub[sk] = sv
                    entry[k] = sub
                else:
                    entry[k] = v
            mcp_table[name] = entry
            if is_managed:
                managed_names.append(name)

        doc["mcp_servers"] = mcp_table

        if managed_names:
            mgmt_section = tomlkit.table()
            mgmt_section["names"] = managed_names
            doc[self._MANAGED_SECTION] = mgmt_section
        elif self._MANAGED_SECTION in doc:
            del doc[self._MANAGED_SECTION]

        if self._LEGACY_MANAGED_SECTION in doc:
            del doc[self._LEGACY_MANAGED_SECTION]

        with open(self._config_path, "w", encoding="utf-8") as f:
            f.write(tomlkit.dumps(doc))


# ---------------------------------------------------------------------------
# Keyed-JSON managers (Gemini, Amp)
# ---------------------------------------------------------------------------


class _KeyedJsonMCPManager(MCPManager):
    """Manager for a JSON settings file with MCPs stored under one key."""

    @property
    def _marker_field(self) -> str:
        return "_managedBy"

    @property
    @abstractmethod
    def mcp_settings_key(self) -> str:
        """Key in the JSON file under which MCP entries are stored."""

    @property
    @abstractmethod
    def _config_path(self) -> Path:
        """Path to the agent's JSON settings file."""

    def _load_full_config(self) -> dict[str, Any]:
        if not self._config_path.exists():
            return {}
        with open(self._config_path, encoding="utf-8") as f:
            data: dict[str, Any] = json.load(f)
        return data

    def _read_installed(self) -> dict[str, Any]:
        full = self._load_full_config()
        return cast(dict[str, Any], full.get(self.mcp_settings_key, {}))

    def _write_installed(self, mcps: dict[str, Any]) -> None:
        self._config_path.parent.mkdir(parents=True, exist_ok=True)
        full = self._load_full_config()
        full[self.mcp_settings_key] = mcps
        write_file_atomic(
            self._config_path, lambda f: json.dump(full, f, indent=2, sort_keys=True)
        )


class GeminiMCPManager(_KeyedJsonMCPManager):
    """Manages MCPs in ~/.gemini/settings.json under the mcpServers key."""

    _translate_extras = {"timeout": 30000, "trust": False}

    @property
    def mcp_settings_key(self) -> str:
        return "mcpServers"

    @property
    def _config_path(self) -> Path:
        return Path.home() / ".gemini" / "settings.json"


class AmpMCPManager(_KeyedJsonMCPManager):
    """Manages MCPs in ~/.config/amp/settings.json under the amp.mcpServers key."""

    @property
    def mcp_settings_key(self) -> str:
        return "amp.mcpServers"

    @property
    def _config_path(self) -> Path:
        return Path.home() / ".config" / "amp" / "settings.json"
