import logging
import sys
from enum import Enum
from pathlib import Path
from typing import Any

from pydantic import BaseModel, PrivateAttr

from agentmux.core.component import (
    Agent,
    Command,
    Component,
    ComponentType,
    Instructions,
    MCPConfig,
    Rule,
    Skill,
)

__all__ = [
    "InstallStrategy",
    "ToolConfig",
    "ComponentsConfig",
    "ExtendsConfig",
    "Setup",
]

logger = logging.getLogger(__name__)

if sys.version_info >= (3, 11):
    import tomllib
else:
    import tomli as tomllib


class InstallStrategy(str, Enum):
    SYMLINK = "symlink"
    COPY = "copy"
    GENERATE = "generate"


class ToolConfig(BaseModel):
    enabled: bool = True
    strategy: InstallStrategy = InstallStrategy.SYMLINK


class ComponentsConfig(BaseModel):
    agents: list[str] = []
    commands: list[str] = []
    skills: list[str] = []
    rules: list[str] = []
    instructions: str | None = None
    mcps: str | dict[str, str] | None = None


class ExtendsConfig(BaseModel):
    """Configuration for a single base setup referenced in ``[extends]``.

    Attributes:
        base: A git URL or local path pointing to the base setup's
            ``agentmux.toml`` directory.  When a git URL is provided and the
            setup is not already in the local registry it will be cloned
            automatically on the first ``load_components()`` call.
    """

    base: str


def _resolve_extended_setup(source: str, _visited: frozenset[str] | None = None) -> "Setup":
    """Locate (and optionally auto-install) the extended setup.

    Resolution order:

    1. If *source* is a local path, load directly from that path.
    2. If *source* is a git URL:
       a. Derive the setup name from the URL.
       b. Check whether it is already in ``~/.agentmux/setups/<name>/``.
       c. If not present, shallow-clone it there and register it.
       d. Load ``agentmux.toml`` from the cloned directory.

    Args:
        source: A git URL or local filesystem path.
        _visited: Internal set of normalised sources already on the resolution
            stack.  Used to detect circular ``[extends]`` chains.  Callers
            should not pass this argument directly; it is threaded through
            recursive calls automatically.

    Returns:
        A fully-parsed :class:`Setup` for the base setup, with the visited set
        stored on ``_extends_visited`` for use by :meth:`Setup.load_components`.

    Raises:
        ValueError: If the source cannot be resolved to a valid agentmux setup,
            or if a circular ``[extends]`` dependency is detected.
    """
    visited: frozenset[str] = _visited or frozenset()

    from agentmux.utils.git import derive_name_from_source as _derive_name_from_source
    from agentmux.utils.git import is_git_url as _is_git_url

    if not _is_git_url(source):
        # Local path — resolve to an absolute canonical path first so that
        # different relative spellings of the same directory are caught.
        local_path = Path(source).expanduser().resolve()
        canonical_key = str(local_path)
        if canonical_key in visited:
            raise ValueError(
                f"Circular [extends] dependency detected: '{source}' "
                f"(resolved: '{local_path}') is already in the resolution chain."
            )
        visited = visited | {canonical_key}
        toml_path = local_path / "agentmux.toml"
        if not toml_path.exists():
            raise ValueError(f"Cannot extend '{source}': no agentmux.toml found at '{toml_path}'.")
        try:
            loaded = Setup.from_toml(toml_path)
        except Exception as exc:
            raise ValueError(f"Failed to load extended setup from '{toml_path}': {exc}") from exc
        loaded._extends_visited = visited
        return loaded

    # Git URL — normalise for cycle detection.
    canonical_key = source.rstrip("/")
    if canonical_key in visited:
        raise ValueError(
            f"Circular [extends] dependency detected: '{source}' "
            "is already in the resolution chain."
        )
    visited = visited | {canonical_key}

    # Git URL — use the registry and auto-clone if needed.
    from agentmux.core.registry import REGISTRY_DIR, Registry, RegistryEntry
    from agentmux.utils.git import clone_repo, get_head_sha

    name = _derive_name_from_source(source)
    registry = Registry.load()
    entry = registry.get(name)

    if entry is not None and entry.local_path.is_dir():
        # Already installed — just load it.
        setup_dir = entry.local_path
    else:
        # Not installed yet — auto-clone into the registry directory.
        REGISTRY_DIR.mkdir(parents=True, exist_ok=True)
        target_dir = REGISTRY_DIR / name
        if not target_dir.exists():
            try:
                clone_repo(source, target_dir)
            except RuntimeError as exc:
                raise ValueError(
                    f"Failed to auto-install extended setup '{source}': {exc}"
                ) from exc
        setup_dir = target_dir

        # Register the newly-cloned setup so subsequent calls skip the clone.
        try:
            sha = get_head_sha(setup_dir)
        except RuntimeError:
            sha = ""

        from datetime import datetime, timezone

        new_entry = RegistryEntry(
            name=name,
            source_url=source,
            installed_sha=sha,
            installed_at=datetime.now(tz=timezone.utc),
            local_path=setup_dir,
        )
        registry.add(new_entry)

    toml_path = setup_dir / "agentmux.toml"
    if not toml_path.exists():
        raise ValueError(f"Extended setup '{source}' (at '{setup_dir}') has no agentmux.toml.")
    try:
        loaded = Setup.from_toml(toml_path)
    except Exception as exc:
        raise ValueError(f"Failed to load extended setup from '{toml_path}': {exc}") from exc
    loaded._extends_visited = visited
    return loaded


class Setup(BaseModel):
    name: str
    version: str = "0.1.0"
    description: str = ""
    authors: list[str] = []
    tools: dict[str, ToolConfig] = {}
    components_config: ComponentsConfig = ComponentsConfig()
    extends: ExtendsConfig | None = None
    root_dir: Path

    # Private cache: avoids re-parsing the same MCP file on every get_mcp_for_tool() call
    # (e.g. when status or sync loops over multiple adapters with a shared config).
    _mcp_cache: dict[str, MCPConfig | None] = PrivateAttr(default_factory=dict)

    # Cached resolved base Setup instance for [extends] — avoids re-cloning / re-parsing
    # the base setup on every get_mcp_for_tool() call (called once per adapter).
    _resolved_base: "Setup | None" = PrivateAttr(default=None)

    # Visited set threaded through [extends] resolution to detect cycles.
    # Populated by _resolve_extended_setup before returning a Setup instance.
    _extends_visited: frozenset[str] = PrivateAttr(default_factory=frozenset)

    @classmethod
    def from_toml(cls, path: Path) -> "Setup":
        with open(path, "rb") as f:
            data = tomllib.load(f)

        setup_meta = data.get("setup", {})

        tools: dict[str, ToolConfig] = {}
        for tool_name, tool_data in data.get("tools", {}).items():
            if isinstance(tool_data, dict):
                tools[tool_name] = ToolConfig(**tool_data)
            else:
                tools[tool_name] = ToolConfig()

        components_data: dict[str, Any] = dict(data.get("components", {}))
        for key in ("agents", "commands", "skills", "rules"):
            if isinstance(components_data.get(key), str):
                components_data[key] = [components_data[key]]

        components_config = ComponentsConfig(**components_data)

        # Parse optional [extends] section.
        extends_config: ExtendsConfig | None = None
        extends_data = data.get("extends")
        if isinstance(extends_data, dict) and extends_data:
            try:
                extends_config = ExtendsConfig(**extends_data)
            except Exception as exc:
                raise ValueError(f"Invalid [extends] section in '{path}': {exc}") from exc

        return cls(
            name=setup_meta.get("name", path.parent.name),
            version=setup_meta.get("version", "0.1.0"),
            description=setup_meta.get("description", ""),
            authors=setup_meta.get("authors", []),
            tools=tools,
            components_config=components_config,
            extends=extends_config,
            root_dir=path.parent,
        )

    def load_components(self) -> dict[ComponentType, list[Component]]:
        """Load all components for this setup, merging any extended base setup.

        Merge rules:
        - Extended base components are loaded first (lower precedence).
        - Local components with the same name override base components.
        - :attr:`ComponentType.MCP` servers are merged additively; local
          server definitions win when the same server name appears in both.
        - :attr:`ComponentType.INSTRUCTIONS` follows the same last-wins rule
          as other named components (there is typically at most one per setup).
        """
        local_result = self._load_local_components()

        if self.extends is None:
            return local_result

        # Resolve and load the base setup (may trigger auto-clone).
        # Pass the visited set so circular [extends] chains are detected.
        # Use the cached instance if already resolved (avoids re-cloning/re-parsing
        # when load_components and get_mcp_for_tool both resolve the same base).
        try:
            if self._resolved_base is None:
                self._resolved_base = _resolve_extended_setup(
                    self.extends.base, self._extends_visited
                )
            base_setup = self._resolved_base
        except ValueError:
            raise
        base_result = base_setup.load_components()

        return _merge_component_maps(base_result, local_result)

    def _load_local_components(self) -> dict[ComponentType, list[Component]]:
        """Load only this setup's own components (no extends resolution)."""
        result: dict[ComponentType, list[Component]] = {
            ComponentType.AGENT: [],
            ComponentType.COMMAND: [],
            ComponentType.SKILL: [],
            ComponentType.RULE: [],
            ComponentType.INSTRUCTIONS: [],
            ComponentType.MCP: [],
        }

        for pattern in self.components_config.agents:
            for p in sorted(self.root_dir.glob(pattern)):
                if p.is_file():
                    try:
                        result[ComponentType.AGENT].append(Agent.from_file(p))
                    except Exception as exc:
                        raise ValueError(f"Failed to load agent '{p}': {exc}") from exc

        for pattern in self.components_config.commands:
            for p in sorted(self.root_dir.glob(pattern)):
                if p.is_file():
                    try:
                        result[ComponentType.COMMAND].append(Command.from_file(p))
                    except Exception as exc:
                        raise ValueError(f"Failed to load command '{p}': {exc}") from exc

        for pattern in self.components_config.skills:
            for p in sorted(self.root_dir.glob(pattern)):
                if p.is_file() and p.name == "SKILL.md":
                    try:
                        result[ComponentType.SKILL].append(Skill.from_dir(p.parent))
                    except Exception as exc:
                        raise ValueError(f"Failed to load skill '{p.parent}': {exc}") from exc
                elif p.is_dir():
                    try:
                        result[ComponentType.SKILL].append(Skill.from_dir(p))
                    except Exception as exc:
                        raise ValueError(f"Failed to load skill '{p}': {exc}") from exc

        for pattern in self.components_config.rules:
            for p in sorted(self.root_dir.glob(pattern)):
                if p.is_file():
                    try:
                        result[ComponentType.RULE].append(Rule.from_file(p))
                    except Exception as exc:
                        raise ValueError(f"Failed to load rule '{p}': {exc}") from exc

        if self.components_config.instructions:
            instr_path = self.root_dir / self.components_config.instructions
            if instr_path.exists():
                try:
                    result[ComponentType.INSTRUCTIONS].append(Instructions.from_file(instr_path))
                except Exception as exc:
                    raise ValueError(f"Failed to load instructions '{instr_path}': {exc}") from exc

        # Collect unique MCP files for validation/status purposes
        mcps = self.components_config.mcps
        if isinstance(mcps, str):
            mcp_path = self.root_dir / mcps
            if mcp_path.exists():
                try:
                    result[ComponentType.MCP].append(MCPConfig.from_file(mcp_path))
                except Exception as exc:
                    raise ValueError(f"Failed to load MCP config '{mcp_path}': {exc}") from exc
        elif isinstance(mcps, dict):
            seen: set[Path] = set()
            for mcp_file in mcps.values():
                mcp_path = self.root_dir / mcp_file
                if mcp_path.exists() and mcp_path not in seen:
                    seen.add(mcp_path)
                    try:
                        result[ComponentType.MCP].append(MCPConfig.from_file(mcp_path))
                    except Exception as exc:
                        raise ValueError(f"Failed to load MCP config '{mcp_path}': {exc}") from exc

        return result

    def get_mcp_for_tool(self, tool_name: str) -> MCPConfig | None:
        """Return the MCP config file designated for a specific tool.

        Results are cached per tool_name so repeated calls (e.g. when iterating
        over adapters in sync or status) do not re-parse the same file.

        When ``[extends]`` is configured the base setup's MCP servers are merged
        additively with the local setup's servers (local servers win on name
        collision).  The merged result is cached and returned.
        """
        if tool_name in self._mcp_cache:
            return self._mcp_cache[tool_name]

        local_mcp = self._get_local_mcp_for_tool(tool_name)

        if self.extends is None:
            self._mcp_cache[tool_name] = local_mcp
            return local_mcp

        # Resolve base setup and get its MCP for the same tool.
        # Pass the visited set so circular [extends] chains are detected.
        # Reuse the cached base instance to avoid re-cloning/re-parsing on each call.
        try:
            if self._resolved_base is None:
                self._resolved_base = _resolve_extended_setup(
                    self.extends.base, self._extends_visited
                )
            base_setup = self._resolved_base
        except ValueError:
            # If the base setup can't be resolved, just use the local MCP.
            self._mcp_cache[tool_name] = local_mcp
            return local_mcp

        base_mcp = base_setup.get_mcp_for_tool(tool_name)

        if base_mcp is None and local_mcp is None:
            self._mcp_cache[tool_name] = None
            return None

        if base_mcp is None:
            self._mcp_cache[tool_name] = local_mcp
            return local_mcp

        if local_mcp is None:
            self._mcp_cache[tool_name] = base_mcp
            return base_mcp

        # Both have MCP configs — merge additively (local wins on name conflict).
        merged_servers: dict[str, Any] = {**base_mcp.servers, **local_mcp.servers}
        merged = MCPConfig(servers=merged_servers, source_path=local_mcp.source_path)
        self._mcp_cache[tool_name] = merged
        return merged

    def _get_local_mcp_for_tool(self, tool_name: str) -> MCPConfig | None:
        """Return only this setup's (non-extended) MCP config for *tool_name*."""
        mcps = self.components_config.mcps
        if mcps is None:
            return None
        if isinstance(mcps, str):
            mcp_path = self.root_dir / mcps
            return MCPConfig.from_file(mcp_path) if mcp_path.exists() else None
        if isinstance(mcps, dict) and tool_name in mcps:
            mcp_path = self.root_dir / mcps[tool_name]
            return MCPConfig.from_file(mcp_path) if mcp_path.exists() else None
        return None


def _merge_component_maps(
    base: dict[ComponentType, list[Component]],
    local: dict[ComponentType, list[Component]],
) -> dict[ComponentType, list[Component]]:
    """Merge *base* component map with *local* component map.

    Merge strategy:
    - For named components (:class:`Agent`, :class:`Command`, :class:`Skill`,
      :class:`Rule`): local components whose ``name`` matches a base component
      override the base entry.  Base-only and local-only components are both
      included.
    - For :class:`Instructions`: the local instructions replace the base
      instructions entirely (there is typically at most one per setup).
    - For :class:`MCPConfig`: servers are merged additively — local server
      definitions win on name collision.

    Args:
        base: Component map from the extended (base) setup.
        local: Component map from the current (child) setup.

    Returns:
        A new merged component map following the rules above.
    """
    result: dict[ComponentType, list[Component]] = {}

    _NAMED_TYPES = (
        ComponentType.AGENT,
        ComponentType.COMMAND,
        ComponentType.SKILL,
        ComponentType.RULE,
    )

    for comp_type in _NAMED_TYPES:
        base_components = base.get(comp_type, [])
        local_components = local.get(comp_type, [])

        local_names: set[str] = {c.name for c in local_components if hasattr(c, "name")}
        # Keep base components whose name is NOT overridden locally.
        merged: list[Component] = [
            c for c in base_components if getattr(c, "name", None) not in local_names
        ]
        # Log any base components that are shadowed by a local component with the same name.
        for shadowed in base_components:
            shadowed_name = getattr(shadowed, "name", None)
            if shadowed_name in local_names:
                logger.debug(
                    "Component %r (%s) from base setup shadowed by local component with same name.",
                    shadowed_name,
                    comp_type.value,
                )
        merged.extend(local_components)
        result[comp_type] = merged

    # Instructions: local replaces base entirely if present.
    local_instructions = local.get(ComponentType.INSTRUCTIONS, [])
    if local_instructions:
        result[ComponentType.INSTRUCTIONS] = local_instructions
    else:
        result[ComponentType.INSTRUCTIONS] = base.get(ComponentType.INSTRUCTIONS, [])

    # MCP: additive merge across all MCPConfig objects (local wins on server
    # name collision).
    base_mcps: list[MCPConfig] = [
        c for c in base.get(ComponentType.MCP, []) if isinstance(c, MCPConfig)
    ]
    local_mcps: list[MCPConfig] = [
        c for c in local.get(ComponentType.MCP, []) if isinstance(c, MCPConfig)
    ]

    if not base_mcps and not local_mcps:
        result[ComponentType.MCP] = []
    elif not base_mcps:
        result[ComponentType.MCP] = list(local_mcps)
    elif not local_mcps:
        result[ComponentType.MCP] = list(base_mcps)
    else:
        # Merge all base servers, then all local servers (local wins on conflict).
        merged_servers: dict[str, Any] = {}
        for mcp in base_mcps:
            merged_servers.update(mcp.servers)
        # Track combined local source path (last local path wins, matching
        # how `get_mcp_for_tool` works).
        local_source: Path | None = None
        for mcp in local_mcps:
            merged_servers.update(mcp.servers)
            if mcp.source_path is not None:
                local_source = mcp.source_path
        result[ComponentType.MCP] = [MCPConfig(servers=merged_servers, source_path=local_source)]

    return result
