"""Skill registry with on-demand loading.

This module provides the SkillRegistry class that manages skills
and loads them only when their tools are invoked.
"""

from __future__ import annotations

import contextlib
import os
from pathlib import Path
from typing import Any

from oclawma.skills.base import (
    SkillError,
    SkillMetadata,
    SkillNotFoundError,
)
from oclawma.skills.discovery import SkillDiscovery
from oclawma.skills.manifest import SkillManifest
from oclawma.skills.skill import LazySkill, ManifestSkill
from oclawma.skills.token_cache import TokenBudgetTracker, TokenCache


class SkillRegistry:
    """Registry for managing skills with lazy loading.

    The registry discovers skills from the filesystem but only loads
    them when a tool is actually invoked. This keeps startup time
    fast and memory usage low.

    Example:
        >>> registry = SkillRegistry()
        >>> registry.discover_skills("/path/to/skills")
        >>> # Skills are discovered but not loaded
        >>> print(registry.list_available())
        ['docker', 'git', 'kubernetes']

        >>> # Skill is loaded only when tool is invoked
        >>> result = await registry.execute_tool("docker", "ps")
    """

    def __init__(
        self,
        token_cache: TokenCache | None = None,
        budget_tracker: TokenBudgetTracker | None = None,
    ) -> None:
        """Initialize the skill registry.

        Args:
            token_cache: Optional token cache for skill token counts
            budget_tracker: Optional budget tracker for token management
        """
        self._skills: dict[str, LazySkill] = {}
        self._manifests: dict[str, SkillManifest] = {}
        self._manifest_paths: dict[str, Path] = {}
        self._discovery = SkillDiscovery()
        self._token_cache = token_cache or TokenCache()
        self._budget_tracker = budget_tracker or TokenBudgetTracker()

    # Discovery methods

    def discover_skills(self, path: Path | str, recursive: bool = True) -> list[str]:
        """Discover skills from a directory.

        This scans the directory for skill manifests and registers
        them for lazy loading. The skills are not actually loaded
        until a tool is invoked.

        Args:
            path: Directory to scan
            recursive: Whether to scan subdirectories

        Returns:
            List of discovered skill names
        """
        discovered = []

        for manifest in self._discovery.discover(path, recursive):
            self._register_manifest(manifest)
            discovered.append(manifest.name)

        # Also track manifest paths
        scanned = self._discovery.scan(path, recursive)
        self._manifest_paths.update(scanned)

        return discovered

    def register_skill(self, skill: LazySkill, metadata: SkillMetadata | None = None) -> None:
        """Register a skill instance directly.

        Args:
            skill: Skill instance to register
            metadata: Optional metadata (extracted from skill if not provided)
        """
        name = skill.name
        self._skills[name] = skill

        # Use provided metadata or get from skill
        meta = metadata or skill.metadata

        # Create a minimal manifest for the skill
        self._manifests[name] = SkillManifest(
            name=name,
            version=meta.version,
            description=meta.description,
            author=meta.author,
            category=meta.category,
            tags=meta.tags,
            requirements=meta.requirements,
            entry_point=meta.entry_point,
        )

    def register_skill_class(
        self, name: str, skill_class: type[LazySkill], metadata: SkillMetadata
    ) -> None:
        """Register a skill class for lazy instantiation.

        Args:
            name: Skill name
            skill_class: Skill class to instantiate on demand
            metadata: Skill metadata
        """
        # Create a lazy wrapper that instantiates on first access
        self._manifests[name] = SkillManifest(
            name=name,
            version=metadata.version,
            description=metadata.description,
            entry_point=f"{skill_class.__module__}:{skill_class.__name__}",
        )

        # Store metadata for later instantiation
        self._skills[name] = skill_class(metadata)

    def _register_manifest(self, manifest: SkillManifest) -> None:
        """Register a skill from its manifest."""
        name = manifest.name

        # Don't overwrite existing skills
        if name in self._manifests:
            return

        self._manifests[name] = manifest

        # Create a lazy skill wrapper
        manifest_path = self._manifest_paths.get(name)
        lazy_skill = ManifestSkill(manifest, manifest_path)
        self._skills[name] = lazy_skill

    # Query methods

    def list_available(self) -> list[str]:
        """List all available skill names.

        Returns:
            List of skill names that are registered
        """
        return list(self._manifests.keys())

    def list_loaded(self) -> list[str]:
        """List skills that have been loaded.

        Returns:
            List of skill names that are currently loaded
        """
        return [name for name, skill in self._skills.items() if skill.is_loaded]

    def has_skill(self, name: str) -> bool:
        """Check if a skill is available.

        Args:
            name: Skill name

        Returns:
            True if the skill is registered
        """
        return name in self._manifests

    def get_skill(self, name: str) -> LazySkill:
        """Get a skill by name.

        Args:
            name: Skill name

        Returns:
            The skill instance

        Raises:
            SkillNotFoundError: If the skill is not found
        """
        if name not in self._skills:
            raise SkillNotFoundError(
                f"Skill '{name}' not found. " f"Available: {', '.join(self.list_available())}",
                skill_name=name,
            )

        return self._skills[name]

    def get_manifest(self, name: str) -> SkillManifest:
        """Get a skill's manifest.

        Args:
            name: Skill name

        Returns:
            The skill manifest

        Raises:
            SkillNotFoundError: If the skill is not found
        """
        if name not in self._manifests:
            raise SkillNotFoundError(f"Skill '{name}' not found", skill_name=name)

        return self._manifests[name]

    def get_metadata(self, name: str) -> SkillMetadata:
        """Get a skill's metadata.

        Args:
            name: Skill name

        Returns:
            The skill metadata
        """
        skill = self.get_skill(name)
        return skill.metadata

    # Tool execution methods

    def has_tool(self, skill_name: str, tool_name: str) -> bool:
        """Check if a skill provides a specific tool.

        Args:
            skill_name: Skill name
            tool_name: Tool name

        Returns:
            True if the skill provides this tool
        """
        if not self.has_skill(skill_name):
            return False

        skill = self._skills[skill_name]
        return skill.has_tool(tool_name)

    def get_tool(self, skill_name: str, tool_name: str) -> Any:
        """Get a tool from a skill.

        This will load the skill if not already loaded.

        Args:
            skill_name: Skill name
            tool_name: Tool name

        Returns:
            The tool instance

        Raises:
            SkillNotFoundError: If the skill is not found
            SkillToolError: If the tool is not found
        """
        skill = self.get_skill(skill_name)
        return skill.get_tool(tool_name)

    async def execute_tool(self, skill_name: str, tool_name: str, **params: Any) -> Any:
        """Execute a tool from a skill.

        This will load the skill if not already loaded, then
        execute the specified tool.

        Args:
            skill_name: Skill name
            tool_name: Tool name
            **params: Parameters to pass to the tool

        Returns:
            Tool execution result

        Raises:
            SkillNotFoundError: If the skill is not found
            SkillToolError: If the tool is not found or fails
        """
        skill = self.get_skill(skill_name)

        # Track token usage for this skill
        tokens = skill.estimate_tokens()
        if not self._budget_tracker.allocate(skill_name, tokens):
            raise SkillError(
                f"Token budget exceeded for skill '{skill_name}'", skill_name=skill_name
            )

        try:
            return await skill.execute_tool(tool_name, **params)
        finally:
            # Release budget after execution
            self._budget_tracker.release(skill_name, tokens)

    def list_tools(self, skill_name: str | None = None) -> list[str] | dict[str, list[str]]:
        """List available tools.

        Args:
            skill_name: Optional skill name to filter by

        Returns:
            If skill_name is provided, returns list of tool names.
            Otherwise, returns dict mapping skill names to tool lists.
        """
        if skill_name:
            skill = self.get_skill(skill_name)
            return skill.list_tools()

        result = {}
        for name, skill in self._skills.items():
            result[name] = skill.list_tools()
        return result

    def get_all_schemas(self) -> list[dict[str, Any]]:
        """Get schemas for all tools in all skills.

        Returns:
            List of all tool schemas
        """
        schemas = []
        for skill in self._skills.values():
            schemas.extend(skill.get_tool_schemas())
        return schemas

    # Loading and unloading

    def load_skill(self, name: str) -> None:
        """Explicitly load a skill.

        Args:
            name: Skill name

        Raises:
            SkillNotFoundError: If the skill is not found
        """
        skill = self.get_skill(name)
        skill.load()

    def unload_skill(self, name: str) -> None:
        """Unload a skill to free resources.

        Args:
            name: Skill name

        Raises:
            SkillNotFoundError: If the skill is not found
        """
        skill = self.get_skill(name)
        skill.unload()

    def unload_all(self) -> None:
        """Unload all skills."""
        for skill in self._skills.values():
            skill.unload()

    def preload_skills(self, names: list[str] | None = None) -> None:
        """Preload skills.

        Args:
            names: Skill names to preload (None for all)
        """
        if names is None:
            names = self.list_available()

        for name in names:
            with contextlib.suppress(SkillError):
                self.load_skill(name)

    # Token management

    def estimate_tokens(self, skill_name: str | None = None) -> int:
        """Estimate token count for skills.

        Args:
            skill_name: Optional skill name (None for all)

        Returns:
            Estimated token count
        """
        if skill_name:
            skill = self.get_skill(skill_name)
            return skill.estimate_tokens()

        total = 0
        for skill in self._skills.values():
            total += skill.estimate_tokens()
        return total

    def get_cached_token_count(self, skill_name: str) -> int | None:
        """Get cached token count for a skill.

        Args:
            skill_name: Skill name

        Returns:
            Cached token count or None
        """
        return self._token_cache.get(skill_name)

    def cache_token_count(self, skill_name: str, tokens: int) -> None:
        """Cache token count for a skill.

        Args:
            skill_name: Skill name
            tokens: Token count to cache
        """
        self._token_cache.set(skill_name, tokens)

    # Registry management

    def clear(self) -> None:
        """Clear all skills from the registry."""
        self.unload_all()
        self._skills.clear()
        self._manifests.clear()
        self._manifest_paths.clear()
        self._discovery.clear_cache()

    def __len__(self) -> int:
        """Return the number of registered skills."""
        return len(self._skills)

    def __contains__(self, name: str) -> bool:
        """Check if a skill is registered."""
        return name in self._skills


def create_registry(
    skills_path: str | None = None, preload: list[str] | None = None
) -> SkillRegistry:
    """Create and configure a skill registry.

    This is a convenience function for creating a fully configured
    skill registry with lazy loading enabled.

    Args:
        skills_path: Path to skills directory (uses env var or default if None)
        preload: Skill names to preload (None for lazy loading only)

    Returns:
        Configured SkillRegistry
    """
    if skills_path is None:
        skills_path = os.environ.get("OCLAWMA_SKILLS_PATH", os.path.expanduser("~/.oclawma/skills"))

    registry = SkillRegistry()

    # Discover skills but don't load them yet
    if os.path.isdir(skills_path):
        registry.discover_skills(skills_path)

    # Preload specified skills
    if preload:
        registry.preload_skills(preload)

    return registry
