"""Token count caching for skills.

This module provides caching for skill token counts to avoid
recomputing them on every request.
"""

from __future__ import annotations

import hashlib
import json
import os
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any


@dataclass
class TokenCacheEntry:
    """A single cache entry for token counts."""

    skill_name: str
    token_count: int
    timestamp: float
    source_hash: str = ""  # Hash of source files for invalidation
    metadata: dict[str, Any] = field(default_factory=dict)

    def is_valid(self, max_age_seconds: float = 3600) -> bool:
        """Check if the cache entry is still valid.

        Args:
            max_age_seconds: Maximum age in seconds

        Returns:
            True if the entry is still valid
        """
        return (time.time() - self.timestamp) < max_age_seconds

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "skill_name": self.skill_name,
            "token_count": self.token_count,
            "timestamp": self.timestamp,
            "source_hash": self.source_hash,
            "metadata": self.metadata,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> TokenCacheEntry:
        """Create from dictionary."""
        return cls(
            skill_name=data["skill_name"],
            token_count=data["token_count"],
            timestamp=data["timestamp"],
            source_hash=data.get("source_hash", ""),
            metadata=data.get("metadata", {}),
        )


class TokenCache:
    """Cache for skill token counts.

    This provides both in-memory and persistent caching of token
    counts to improve performance.

    Example:
        >>> cache = TokenCache()
        >>> cache.set("docker", 1500, source_hash="abc123")
        >>> cache.get("docker")
        1500
        >>> cache.get("unknown")  # Returns None
        None
    """

    DEFAULT_MAX_AGE = 3600  # 1 hour

    def __init__(
        self, cache_dir: Path | str | None = None, max_age_seconds: float = DEFAULT_MAX_AGE
    ) -> None:
        """Initialize the token cache.

        Args:
            cache_dir: Directory for persistent cache (None for no persistence)
            max_age_seconds: Default maximum age for cache entries
        """
        self._memory_cache: dict[str, TokenCacheEntry] = {}
        self._max_age = max_age_seconds

        if cache_dir is None:
            cache_dir = os.environ.get("OCLAWMA_CACHE_DIR", os.path.expanduser("~/.cache/oclawma"))

        self._cache_dir = Path(cache_dir)
        self._cache_file = self._cache_dir / "token_cache.json"

        # Load persistent cache
        self._load_persistent_cache()

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

        Args:
            skill_name: Name of the skill
            source_hash: Optional hash to validate against

        Returns:
            Cached token count or None if not found/invalid
        """
        entry = self._memory_cache.get(skill_name)

        if entry is None:
            return None

        # Check if entry is expired
        if not entry.is_valid(self._max_age):
            del self._memory_cache[skill_name]
            return None

        # Validate source hash if provided
        if source_hash is not None and entry.source_hash != source_hash:
            return None

        return entry.token_count

    def set(
        self,
        skill_name: str,
        token_count: int,
        source_hash: str = "",
        metadata: dict[str, Any] | None = None,
    ) -> None:
        """Set cached token count for a skill.

        Args:
            skill_name: Name of the skill
            token_count: Token count to cache
            source_hash: Hash of source files for invalidation
            metadata: Optional metadata
        """
        entry = TokenCacheEntry(
            skill_name=skill_name,
            token_count=token_count,
            timestamp=time.time(),
            source_hash=source_hash,
            metadata=metadata or {},
        )

        self._memory_cache[skill_name] = entry
        self._save_persistent_cache()

    def invalidate(self, skill_name: str) -> bool:
        """Invalidate a cached entry.

        Args:
            skill_name: Name of the skill to invalidate

        Returns:
            True if an entry was removed
        """
        if skill_name in self._memory_cache:
            del self._memory_cache[skill_name]
            self._save_persistent_cache()
            return True
        return False

    def clear(self) -> None:
        """Clear all cached entries."""
        self._memory_cache.clear()
        self._save_persistent_cache()

    def get_all(self) -> dict[str, TokenCacheEntry]:
        """Get all valid cache entries.

        Returns:
            Dictionary of skill name to cache entry
        """
        # Clean up expired entries
        now = time.time()
        expired = [
            name
            for name, entry in self._memory_cache.items()
            if (now - entry.timestamp) >= self._max_age
        ]
        for name in expired:
            del self._memory_cache[name]

        return dict(self._memory_cache)

    def estimate_tokens(self, text: str) -> int:
        """Estimate token count for text.

        Uses a rough approximation of 1 token per 4 characters
        for English text.

        Args:
            text: Text to estimate

        Returns:
            Estimated token count
        """
        # Rough estimate: 1 token per 4 characters for English
        return len(text) // 4

    def estimate_from_schema(self, schema: dict[str, Any]) -> int:
        """Estimate token count from a tool schema.

        Args:
            schema: Tool schema dictionary

        Returns:
            Estimated token count
        """
        # Convert schema to string and estimate
        schema_str = json.dumps(schema, sort_keys=True)
        return self.estimate_tokens(schema_str)

    def _load_persistent_cache(self) -> None:
        """Load cache from disk."""
        if not self._cache_file.exists():
            return

        try:
            with open(self._cache_file, encoding="utf-8") as f:
                data = json.load(f)

            for skill_name, entry_data in data.items():
                entry = TokenCacheEntry.from_dict(entry_data)

                # Only load valid entries
                if entry.is_valid(self._max_age):
                    self._memory_cache[skill_name] = entry

        except (json.JSONDecodeError, KeyError, TypeError):
            # Ignore corrupt cache files
            pass

    def _save_persistent_cache(self) -> None:
        """Save cache to disk."""
        try:
            self._cache_dir.mkdir(parents=True, exist_ok=True)

            data = {name: entry.to_dict() for name, entry in self._memory_cache.items()}

            with open(self._cache_file, "w", encoding="utf-8") as f:
                json.dump(data, f, indent=2)

        except (OSError, PermissionError):
            # Ignore write errors (cache is optional)
            pass

    @staticmethod
    def compute_hash(*files: Path | str) -> str:
        """Compute a hash of source files for cache invalidation.

        Args:
            *files: Paths to files to hash

        Returns:
            Hash string
        """
        hasher = hashlib.sha256()

        for file_path in sorted(files):
            try:
                with open(file_path, "rb") as f:
                    hasher.update(f.read())
            except (OSError, FileNotFoundError):
                pass

        return hasher.hexdigest()[:16]


class TokenBudgetTracker:
    """Track token usage across skills with budget management."""

    def __init__(self, budget: int = 128000) -> None:
        """Initialize the budget tracker.

        Args:
            budget: Maximum token budget
        """
        self._budget = budget
        self._used = 0
        self._skill_usage: dict[str, int] = {}

    @property
    def budget(self) -> int:
        """Get the total budget."""
        return self._budget

    @property
    def used(self) -> int:
        """Get the used tokens."""
        return self._used

    @property
    def remaining(self) -> int:
        """Get the remaining tokens."""
        return self._budget - self._used

    @property
    def percent_used(self) -> float:
        """Get the percentage of budget used."""
        if self._budget == 0:
            return 100.0
        return (self._used / self._budget) * 100

    def allocate(self, skill_name: str, tokens: int) -> bool:
        """Allocate tokens to a skill.

        Args:
            skill_name: Name of the skill
            tokens: Number of tokens to allocate

        Returns:
            True if allocation succeeded
        """
        if tokens > self.remaining:
            return False

        self._used += tokens
        self._skill_usage[skill_name] = self._skill_usage.get(skill_name, 0) + tokens
        return True

    def release(self, skill_name: str, tokens: int | None = None) -> None:
        """Release tokens allocated to a skill.

        Args:
            skill_name: Name of the skill
            tokens: Tokens to release (None for all)
        """
        if tokens is None:
            tokens = self._skill_usage.get(skill_name, 0)

        self._used = max(0, self._used - tokens)
        self._skill_usage[skill_name] = max(0, self._skill_usage.get(skill_name, 0) - tokens)

    def get_skill_usage(self, skill_name: str) -> int:
        """Get token usage for a specific skill."""
        return self._skill_usage.get(skill_name, 0)

    def reset(self) -> None:
        """Reset all usage tracking."""
        self._used = 0
        self._skill_usage.clear()
