import json
import os
import re
import shutil
import subprocess
from pathlib import Path
from typing import Dict, List, Optional, Any, Callable, Set
from dataclasses import dataclass, asdict

@dataclass
class Skill:
    name: str
    description: str
    config: Dict[str, Any]  # MCP server configuration (command, args, env)


@dataclass
class SkillCheckResult:
    name: str
    ready: bool
    missing_commands: List[str]
    missing_env: List[str]
    notes: List[str]

class SkillManager:
    """Manages the registry of reusable MCP Skills (The Armory)."""
    
    SKILLS_FILE = Path.home() / ".armorer" / "skills.json"

    def __init__(self):
        self.SKILLS_FILE.parent.mkdir(parents=True, exist_ok=True)
        self._ensure_registry()

    def _ensure_registry(self):
        if not self.SKILLS_FILE.exists():
            # Seed with some default skills
            defaults = {
                "filesystem": {
                    "name": "filesystem",
                    "description": "Access local files via MCP",
                    "config": {
                        "command": "npx",
                        "args": ["-y", "@modelcontextprotocol/server-filesystem", "/Users/Shared"]
                    }
                },
                "brave-search": {
                    "name": "brave-search",
                    "description": "Web search using Brave",
                    "config": {
                        "command": "npx",
                        "args": ["-y", "@modelcontextprotocol/server-brave-search"],
                        "env": {"BRAVE_API_KEY": "${BRAVE_API_KEY}"}
                    }
                },
                "git": {
                    "name": "git",
                    "description": "Git repository operations",
                    "config": {
                        "command": "docker",
                        "args": ["run", "-i", "--rm", "mcp/git"] 
                    }
                }
            }
            self._save_registry(defaults)

    def _load_registry(self) -> Dict[str, dict]:
        try:
            with open(self.SKILLS_FILE, "r") as f:
                return json.load(f)
        except json.JSONDecodeError:
            return {}

    def _save_registry(self, data: Dict[str, dict]):
        with open(self.SKILLS_FILE, "w") as f:
            json.dump(data, f, indent=4)

    def add_skill(self, skill: Skill):
        """Add or update a skill."""
        registry = self._load_registry()
        registry[skill.name] = asdict(skill)
        self._save_registry(registry)

    def get_skill(self, name: str) -> Optional[Skill]:
        """Retrieve a skill by name."""
        registry = self._load_registry()
        data = registry.get(name)
        if data:
            return Skill(**data)
        return None

    def list_skills(self) -> List[Skill]:
        """List all available skills."""
        registry = self._load_registry()
        return [Skill(**data) for data in registry.values()]

    def remove_skill(self, name: str):
        """Remove a skill from the registry."""
        registry = self._load_registry()
        if name in registry:
            del registry[name]
            self._save_registry(registry)

    def required_env_vars(self, skill: Skill) -> List[str]:
        env = skill.config.get("env", {})
        if not isinstance(env, dict):
            return []
        required: Set[str] = set()
        pattern = re.compile(r"\$\{([A-Z0-9_]+)\}")
        for value in env.values():
            if not isinstance(value, str):
                continue
            for match in pattern.findall(value):
                required.add(match)
        return sorted(required)

    def check_skill(
        self,
        name: str,
        env_lookup: Optional[Callable[[str], Optional[str]]] = None,
    ) -> SkillCheckResult:
        skill = self.get_skill(name)
        if not skill:
            return SkillCheckResult(
                name=name,
                ready=False,
                missing_commands=[],
                missing_env=[],
                notes=["Skill is not in the Armorer registry."],
            )

        command = str(skill.config.get("command", "")).strip()
        missing_commands: List[str] = []
        missing_env: List[str] = []
        notes: List[str] = []

        if command and shutil.which(command) is None:
            missing_commands.append(command)

        resolver = env_lookup or (lambda key: os.getenv(key))
        for env_key in self.required_env_vars(skill):
            if not resolver(env_key):
                missing_env.append(env_key)

        return SkillCheckResult(
            name=name,
            ready=not missing_commands and not missing_env,
            missing_commands=missing_commands,
            missing_env=missing_env,
            notes=notes,
        )

    def ensure_skill_ready(
        self,
        name: str,
        env_lookup: Optional[Callable[[str], Optional[str]]] = None,
    ) -> SkillCheckResult:
        result = self.check_skill(name, env_lookup=env_lookup)
        if not result.ready:
            return result

        skill = self.get_skill(name)
        if not skill:
            return result

        command = str(skill.config.get("command", "")).strip()
        args = [str(arg) for arg in skill.config.get("args", []) if str(arg).strip()]
        env = os.environ.copy()
        resolver = env_lookup or (lambda key: os.getenv(key))
        for env_key in self.required_env_vars(skill):
            resolved = resolver(env_key)
            if resolved:
                env[env_key] = resolved

        try:
            if command == "npx":
                package = self._extract_npx_package(args)
                if package:
                    probe = subprocess.run(
                        ["npx", "-y", package, "--help"],
                        capture_output=True,
                        text=True,
                        timeout=30,
                        env=env,
                    )
                    if probe.returncode != 0:
                        result.notes.append((probe.stderr or probe.stdout or "").strip())
                        result.ready = False
                        return result
                    result.notes.append(f"Validated npm package {package}.")
            elif command == "docker":
                image = self._extract_docker_image(args)
                if image:
                    inspect_result = subprocess.run(
                        ["docker", "image", "inspect", image],
                        capture_output=True,
                        text=True,
                        timeout=30,
                        env=env,
                    )
                    if inspect_result.returncode != 0:
                        pull_result = subprocess.run(
                            ["docker", "pull", image],
                            capture_output=True,
                            text=True,
                            timeout=120,
                            env=env,
                        )
                        if pull_result.returncode != 0:
                            result.notes.append((pull_result.stderr or pull_result.stdout or "").strip())
                            result.ready = False
                            return result
                        result.notes.append(f"Pulled Docker image {image}.")
                    else:
                        result.notes.append(f"Docker image {image} is available.")
        except Exception as exc:
            result.ready = False
            result.notes.append(str(exc))
            return result

        result.ready = True
        return result

    def _extract_npx_package(self, args: List[str]) -> Optional[str]:
        skip_next = False
        for idx, arg in enumerate(args):
            if skip_next:
                skip_next = False
                continue
            if arg in {"-y", "--yes"}:
                continue
            if arg in {"-p", "--package"}:
                if idx + 1 < len(args):
                    return args[idx + 1]
                skip_next = True
                continue
            if arg.startswith("-"):
                continue
            return arg
        return None

    def _extract_docker_image(self, args: List[str]) -> Optional[str]:
        seen_run = False
        skip_next = False
        for arg in args:
            if skip_next:
                skip_next = False
                continue
            if not seen_run:
                if arg == "run":
                    seen_run = True
                continue
            if arg in {"--name", "--platform", "--entrypoint", "-e", "--env", "-v", "--volume", "-w", "--workdir"}:
                skip_next = True
                continue
            if arg.startswith("-"):
                continue
            return arg
        return None
