#!/usr/bin/env python3
"""
TITAN CACHE MANAGER
===================
Manages Gemini Context Cache (Titan Memory) for codebase caching.

VERIFICATION_STAMP
Story: TITAN-002 (Cache Manager)
Verified By: Claude
Verified At: 2026-01-23
"""

import os
import sys
import time
import logging
from pathlib import Path
from dataclasses import dataclass, field
from typing import List, Optional, Generator
from datetime import datetime, timedelta

# Add genesis root for imports
sys.path.insert(0, str(Path(__file__).parent.parent.parent))

logger = logging.getLogger(__name__)

# Try to import Google GenAI
try:
    import google.generativeai as genai
    GENAI_AVAILABLE = True
except ImportError:
    GENAI_AVAILABLE = False
    logger.warning("google-generativeai not installed. Titan Memory unavailable.")


@dataclass
class TitanMemory:
    """
    Represents a Titan Memory (cached content) block.
    """
    name: str  # Cache resource name (e.g., 'cachedContents/12345')
    display_name: str
    model: str
    expire_time: str
    token_count: int
    created_time: str = field(default_factory=lambda: datetime.now().isoformat())
    ttl_seconds: int = 3600

    def time_until_expiry(self) -> timedelta:
        """Get time remaining until cache expires."""
        try:
            expire = datetime.fromisoformat(self.expire_time.replace('Z', '+00:00'))
            return expire - datetime.now(expire.tzinfo)
        except:
            return timedelta(seconds=self.ttl_seconds)

    def is_expired(self) -> bool:
        """Check if cache has expired."""
        return self.time_until_expiry().total_seconds() <= 0


class TitanCacheManager:
    """
    Manages Titan Memory (Gemini Context Cache) lifecycle.

    Features:
    - Create caches from file lists
    - Batch upload for large file sets
    - Query caches with prompts
    - Automatic cache selection
    - Cache cleanup
    """

    # Configuration
    BATCH_SIZE = 50  # Max files per batch upload
    DEFAULT_MODEL = "gemini-2.0-flash-001"  # Supports context caching
    DEFAULT_TTL_MINUTES = 60
    UPLOAD_TIMEOUT = 300  # 5 minutes max for upload
    PROCESSING_TIMEOUT = 120  # 2 minutes max for processing
    MAX_RETRIES = 3

    def __init__(self, api_key: Optional[str] = None):
        """
        Initialize the cache manager.

        Args:
            api_key: Gemini API key (or uses GEMINI_API_KEY env var)
        """
        self.api_key = api_key or os.environ.get('GEMINI_API_KEY')

        if not self.api_key:
            raise ValueError("GEMINI_API_KEY not set")

        if GENAI_AVAILABLE:
            genai.configure(api_key=self.api_key)

        self._current_cache: Optional[TitanMemory] = None

    def create_cache(
        self,
        files: List[Path],
        display_name: Optional[str] = None,
        model: str = DEFAULT_MODEL,
        ttl_minutes: int = DEFAULT_TTL_MINUTES,
        system_instruction: Optional[str] = None,
    ) -> Optional[TitanMemory]:
        """
        Create a new Titan Memory cache from files.

        Args:
            files: List of file paths to cache
            display_name: Human-readable cache name
            model: Gemini model to use
            ttl_minutes: Time-to-live in minutes
            system_instruction: System prompt for the cache

        Returns:
            TitanMemory object or None on failure
        """
        if not GENAI_AVAILABLE:
            logger.error("google-generativeai SDK not available")
            return None

        display_name = display_name or f"genesis_cache_{int(time.time())}"
        system_instruction = system_instruction or self._default_system_instruction()

        logger.info(f"Creating Titan cache '{display_name}' with {len(files)} files...")

        try:
            # Step 1: Upload files in batches
            uploaded_files = []
            for batch in self._create_batches(files):
                batch_uploaded = self._upload_batch(batch)
                uploaded_files.extend(batch_uploaded)

            if not uploaded_files:
                logger.error("No files were uploaded successfully")
                return None

            logger.info(f"Uploaded {len(uploaded_files)} files, waiting for processing...")

            # Step 2: Wait for all files to finish processing
            ready_files = []
            for uf in uploaded_files:
                if self._wait_for_file(uf):
                    ready_files.append(uf)

            if not ready_files:
                logger.error("No files finished processing")
                return None

            logger.info(f"{len(ready_files)} files ready, creating cache...")

            # Step 3: Create the cache
            cache = genai.caching.CachedContent.create(
                model=model,
                display_name=display_name,
                system_instruction=system_instruction,
                contents=ready_files,
                ttl=timedelta(minutes=ttl_minutes),
            )

            titan_memory = TitanMemory(
                name=cache.name,
                display_name=cache.display_name,
                model=cache.model,
                expire_time=str(cache.expire_time),
                token_count=cache.usage_metadata.total_token_count,
                ttl_seconds=ttl_minutes * 60,
            )

            self._current_cache = titan_memory

            logger.info(f"Cache created: {cache.name} ({titan_memory.token_count:,} tokens)")
            return titan_memory

        except Exception as e:
            logger.error(f"Failed to create cache: {e}")
            return None

    def create_full_stack_cache(
        self,
        ttl_minutes: int = DEFAULT_TTL_MINUTES,
    ) -> Optional[TitanMemory]:
        """
        Create a cache with the full Genesis stack.

        Includes: core/, skills/, tools/, rlm/, swarms/

        Args:
            ttl_minutes: Time-to-live in minutes

        Returns:
            TitanMemory object or None on failure
        """
        from .file_scanner import FileScanner

        scanner = FileScanner()
        files = scanner.scan()

        logger.info(f"Full stack scan: {len(files)} files, ~{scanner.get_estimated_tokens(files):,} tokens")

        return self.create_cache(
            files=files,
            display_name=f"genesis_fullstack_{int(time.time())}",
            ttl_minutes=ttl_minutes,
        )

    def query(
        self,
        prompt: str,
        cache_name: Optional[str] = None,
        temperature: float = 0.7,
    ) -> Optional[str]:
        """
        Query a Titan cache with a prompt.

        Args:
            prompt: The question/prompt
            cache_name: Specific cache to use (or auto-select)
            temperature: Generation temperature

        Returns:
            Response text or None on failure
        """
        if not GENAI_AVAILABLE:
            return None

        try:
            # Get cache
            if cache_name:
                cache = genai.caching.CachedContent.get(cache_name)
            else:
                cache = self._get_best_cache()

            if not cache:
                logger.warning("No cache available for query")
                return None

            # Create model with cache
            model = genai.GenerativeModel.from_cached_content(cache)

            # Generate response
            response = model.generate_content(
                prompt,
                generation_config=genai.GenerationConfig(temperature=temperature)
            )

            return response.text

        except Exception as e:
            logger.error(f"Query failed: {e}")
            return None

    def list_caches(self) -> List[TitanMemory]:
        """
        List all active Titan caches.

        Returns:
            List of TitanMemory objects
        """
        if not GENAI_AVAILABLE:
            return []

        try:
            caches = []
            for cache in genai.caching.CachedContent.list():
                caches.append(TitanMemory(
                    name=cache.name,
                    display_name=cache.display_name,
                    model=cache.model,
                    expire_time=str(cache.expire_time),
                    token_count=cache.usage_metadata.total_token_count if cache.usage_metadata else 0,
                ))
            return caches
        except Exception as e:
            logger.error(f"Failed to list caches: {e}")
            return []

    def delete_cache(self, cache_name: str) -> bool:
        """
        Delete a Titan cache.

        Args:
            cache_name: Cache resource name

        Returns:
            True if deleted successfully
        """
        if not GENAI_AVAILABLE:
            return False

        try:
            cache = genai.caching.CachedContent.get(cache_name)
            cache.delete()
            logger.info(f"Deleted cache: {cache_name}")
            return True
        except Exception as e:
            logger.error(f"Failed to delete cache: {e}")
            return False

    def ensure_cache_exists(self) -> Optional[TitanMemory]:
        """
        Ensure a cache exists, creating if needed.

        Returns:
            Active TitanMemory or None
        """
        caches = self.list_caches()

        if caches:
            # Return the one with most time remaining
            return self._select_best_cache()
        else:
            # Create new full-stack cache
            return self.create_full_stack_cache()

    def _create_batches(self, files: List[Path]) -> Generator[List[Path], None, None]:
        """
        Split files into batches for upload.

        Args:
            files: List of file paths

        Yields:
            Batches of file paths
        """
        for i in range(0, len(files), self.BATCH_SIZE):
            yield files[i:i + self.BATCH_SIZE]

    def _upload_batch(self, files: List[Path]) -> List:
        """
        Upload a batch of files.

        Args:
            files: List of file paths

        Returns:
            List of uploaded file objects
        """
        uploaded = []

        for file_path in files:
            try:
                uf = self._upload_with_retry(file_path)
                if uf:
                    uploaded.append(uf)
            except Exception as e:
                logger.warning(f"Failed to upload {file_path}: {e}")

        return uploaded

    def _upload_with_retry(self, file_path: Path, max_retries: int = None) -> Optional[object]:
        """
        Upload a single file with retry logic.

        Args:
            file_path: Path to file
            max_retries: Number of retries

        Returns:
            Uploaded file object or None
        """
        max_retries = max_retries or self.MAX_RETRIES

        for attempt in range(max_retries):
            try:
                return self._upload_single_file(file_path)
            except Exception as e:
                if attempt < max_retries - 1:
                    time.sleep(2 ** attempt)  # Exponential backoff
                else:
                    raise

        return None

    def _upload_single_file(self, file_path: Path) -> object:
        """
        Upload a single file to Gemini.

        Args:
            file_path: Path to file

        Returns:
            Uploaded file object
        """
        mime_type = self._get_mime_type(file_path)

        return genai.upload_file(
            path=str(file_path),
            mime_type=mime_type,
            display_name=file_path.name,
        )

    def _wait_for_file(self, uploaded_file, timeout: int = None) -> bool:
        """
        Wait for a file to finish processing.

        Args:
            uploaded_file: Uploaded file object
            timeout: Max wait time in seconds

        Returns:
            True if file is ready
        """
        timeout = timeout or self.PROCESSING_TIMEOUT
        start = time.time()

        while time.time() - start < timeout:
            try:
                file = genai.get_file(uploaded_file.name)
                if file.state.name == "ACTIVE":
                    return True
                elif file.state.name == "FAILED":
                    logger.warning(f"File processing failed: {uploaded_file.name}")
                    return False
            except:
                pass
            time.sleep(1)

        logger.warning(f"File processing timeout: {uploaded_file.name}")
        return False

    def _get_file_state(self, file_name: str) -> str:
        """Get the state of an uploaded file."""
        try:
            file = genai.get_file(file_name)
            return file.state.name
        except:
            return "UNKNOWN"

    def _get_mime_type(self, file_path: Path) -> str:
        """
        Determine MIME type for a file.

        Args:
            file_path: Path to file

        Returns:
            MIME type string
        """
        suffix = file_path.suffix.lower()
        mime_map = {
            '.py': 'text/x-python',
            '.md': 'text/markdown',
            '.json': 'application/json',
            '.yaml': 'text/yaml',
            '.yml': 'text/yaml',
            '.html': 'text/html',
            '.txt': 'text/plain',
        }
        return mime_map.get(suffix, 'text/plain')

    def _get_best_cache(self):
        """Get the best available cache (internal)."""
        try:
            caches = list(genai.caching.CachedContent.list())
            if not caches:
                return None
            # Sort by expire time, return the one with most time remaining
            return sorted(caches, key=lambda c: str(c.expire_time), reverse=True)[0]
        except:
            return None

    def _select_best_cache(self) -> Optional[TitanMemory]:
        """
        Select the best cache from available caches.

        Returns:
            Best TitanMemory or None
        """
        caches = self.list_caches()
        if not caches:
            return None

        # Sort by time until expiry, pick the one with most time remaining
        valid_caches = [c for c in caches if not c.is_expired()]
        if not valid_caches:
            return None

        return max(valid_caches, key=lambda c: c.time_until_expiry())

    def _default_system_instruction(self) -> str:
        """Default system instruction for caches."""
        return """You are the Genesis Titan Intelligence.

You have complete knowledge of the Genesis codebase via the cached files.
When answering questions:
1. Reference specific files and line numbers when relevant
2. Explain code patterns and architecture
3. If information isn't in the cache, say so clearly
4. Be concise but thorough

Your knowledge includes:
- core/ - Main execution layer, orchestration, memory systems
- skills/ - Skill modules and capabilities
- tools/ - Utility tools and integrations
- rlm/ - Reasoning Layer (knowledge graph, learning)
- swarms/ - Multi-agent swarm orchestration
"""


if __name__ == '__main__':
    # Quick test
    logging.basicConfig(level=logging.INFO)

    manager = TitanCacheManager()

    print("Listing existing caches...")
    caches = manager.list_caches()

    if caches:
        print(f"Found {len(caches)} caches:")
        for c in caches:
            print(f"  - {c.display_name}: {c.token_count:,} tokens, expires {c.expire_time}")
    else:
        print("No caches found.")
        print("\nCreating full-stack cache...")
        cache = manager.create_full_stack_cache(ttl_minutes=60)
        if cache:
            print(f"Created: {cache.name} ({cache.token_count:,} tokens)")