import os
import json
import hashlib
from pathlib import Path
from typing import Dict, Any, Optional

CACHE_DIR_NAME = ".ast_cache"
INDEX_FILE_NAME = ".cache_index.json"

class CacheManager:
    def __init__(self, project_root: Path):
        self.project_root = project_root
        self.cache_dir = project_root / CACHE_DIR_NAME
        self.index_file = self.cache_dir / INDEX_FILE_NAME
        self.index: Dict[str, str] = {}
        self._load_index()

    def _load_index(self):
        """Loads the cache index from the file."""
        if self.index_file.exists():
            try:
                with open(self.index_file, "r", encoding="utf-8") as f:
                    self.index = json.load(f)
            except (json.JSONDecodeError, IOError):
                self.index = {}
        else:
            self.index = {}

    def _save_index(self):
        """Saves the current index to the file."""
        self.cache_dir.mkdir(parents=True, exist_ok=True)
        with open(self.index_file, "w", encoding="utf-8") as f:
            json.dump(self.index, f, indent=2)

    def get_file_hash(self, file_path: Path) -> str:
        """Calculates the SHA-256 hash of a file."""
        sha256_hash = hashlib.sha256()
        with open(file_path, "rb") as f:
            for byte_block in iter(lambda: f.read(4096), b""):
                sha256_hash.update(byte_block)
        return sha256_hash.hexdigest()

    def is_up_to_date(self, rel_path: str, current_hash: str) -> bool:
        """Checks if the cached AST for a file is up-to-date."""
        cached_hash = self.index.get(rel_path)
        if not cached_hash:
            return False
        
        # Also check if the actual cache file exists
        cache_file = self.cache_dir / f"{rel_path}.json"
        return cached_hash == current_hash and cache_file.exists()

    def update_cache(self, rel_path: str, current_hash: str, ast_data: Any):
        """Updates the cache for a specific file."""
        # Ensure subdirectories exist in cache
        cache_file = self.cache_dir / f"{rel_path}.json"
        cache_file.parent.mkdir(parents=True, exist_ok=True)
        
        # Save AST data
        with open(cache_file, "w", encoding="utf-8") as f:
            json.dump(ast_data, f, separators=(",", ":"))
        
        # Update index
        self.index[rel_path] = current_hash
        self._save_index()

    def get_cached_ast(self, rel_path: str) -> Optional[Any]:
        """Retrieves the cached AST data for a file."""
        cache_file = self.cache_dir / f"{rel_path}.json"
        if cache_file.exists():
            try:
                with open(cache_file, "r", encoding="utf-8") as f:
                    return json.load(f)
            except (json.JSONDecodeError, IOError):
                return None
        return None

    def clear_cache(self):
        """Removes the entire cache directory."""
        if self.cache_dir.exists():
            import shutil
            shutil.rmtree(self.cache_dir)
        self.index = {}

    def prune_index(self, existing_files: set[str]):
        """Removes entries from the index and cache files that no longer exist in the project."""
        to_remove = [rel_path for rel_path in self.index if rel_path not in existing_files]
        for rel_path in to_remove:
            cache_file = self.cache_dir / f"{rel_path}.json"
            if cache_file.exists():
                cache_file.unlink()
            del self.index[rel_path]
        
        if to_remove:
            self._save_index()

