#!/usr/bin/env python3
"""
Resource manager for MCP server to track and serve generated files.
"""

import os
import tempfile
from pathlib import Path
from typing import Dict, Optional
from datetime import datetime


class ResourceManager:
    """Manages resources (files) that can be accessed via MCP resources API."""

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

        Args:
            base_dir: Base directory for storing resources. If None, uses a temp directory.
        """
        if base_dir:
            self.base_dir = Path(base_dir)
            self.base_dir.mkdir(parents=True, exist_ok=True)
        else:
            # Use a temp directory in the user's home or system temp
            temp_base = Path(tempfile.gettempdir()) / "comet-mcp-resources"
            temp_base.mkdir(parents=True, exist_ok=True)
            self.base_dir = temp_base

        # Track resources: uri -> file_path
        self._resources: Dict[str, Path] = {}

    def register_file(
        self,
        file_path: Path,
        uri: Optional[str] = None,
        description: Optional[str] = None,
    ) -> str:
        """
        Register a file as a resource.

        Args:
            file_path: Path to the file to register
            uri: Optional custom URI. If not provided, generates one based on filename.
            description: Optional description of the resource

        Returns:
            The URI for the resource
        """
        if uri is None:
            # Generate a URI based on the filename
            # Use the filename directly without adding a timestamp prefix
            # The filename already contains a timestamp if it was generated by our tools
            filename = file_path.name
            uri = f"file://comet-mcp/{filename}"

        # Store absolute path
        abs_path = file_path.resolve()
        self._resources[uri] = abs_path

        return uri

    def create_file(
        self, filename: str, content: bytes, uri: Optional[str] = None
    ) -> str:
        """
        Create a file and register it as a resource.

        Args:
            filename: Name of the file to create
            content: File content as bytes
            uri: Optional custom URI. If not provided, uses file://comet-mcp/{filename}

        Returns:
            The URI for the resource
        """
        file_path = self.base_dir / filename
        file_path.write_bytes(content)

        # If URI not provided, generate one based on filename (no timestamp prefix)
        if uri is None:
            uri = f"file://comet-mcp/{filename}"

        return self.register_file(file_path, uri)

    def get_file_path(self, uri: str) -> Optional[Path]:
        """
        Get the file path for a resource URI.

        Args:
            uri: The resource URI

        Returns:
            Path to the file, or None if not found
        """
        return self._resources.get(uri)

    def list_resources(self) -> Dict[str, Dict[str, str]]:
        """
        List all registered resources.

        Also scans the base directory for files that might have been created
        but not registered (e.g., across process boundaries).

        Returns:
            Dictionary mapping URIs to resource metadata
        """
        resources = {}

        # First, add all registered resources
        for uri, file_path in self._resources.items():
            if file_path.exists():
                resources[uri] = {
                    "uri": uri,
                    "name": file_path.name,
                    "mimeType": self.guess_mime_type(file_path),
                    "description": f"Generated file: {file_path.name}",
                }

        # Also scan the base directory for files that match our URI pattern
        # This handles cases where files were created but not registered (e.g., across processes)
        if self.base_dir.exists():
            for file_path in self.base_dir.glob("*"):
                if file_path.is_file():
                    # Generate URI for this file if it matches our pattern
                    # Files created by our tools should match: experiments_*.csv or similar
                    if (
                        file_path.name.startswith("experiments_")
                        or file_path.suffix == ".csv"
                    ):
                        # Try to find existing URI or create one
                        found_uri = None
                        for existing_uri, existing_path in self._resources.items():
                            if existing_path == file_path:
                                found_uri = existing_uri
                                break

                        # If not found, generate URI and register it
                        if found_uri is None:
                            # Check if a URI ending with this filename already exists
                            # This handles the case where the URI was created with a timestamp prefix
                            # Format: file://comet-mcp/timestamp_experiments_project_timestamp.csv
                            matching_uri = None
                            for existing_uri in self._resources.keys():
                                if existing_uri.endswith(f"/{file_path.name}"):
                                    matching_uri = existing_uri
                                    # Update the path in case it changed
                                    self._resources[matching_uri] = file_path
                                    found_uri = matching_uri
                                    break

                            # If still not found, try to extract timestamp from filename
                            # Files created by tools: experiments_project_timestamp.csv
                            # URI format: file://comet-mcp/timestamp_experiments_project_timestamp.csv
                            # Try to find a pattern match
                            if found_uri is None:
                                # Try simple format first
                                simple_uri = f"file://comet-mcp/{file_path.name}"
                                # But also check if filename contains a timestamp pattern
                                # If filename is like "experiments_project_20251206_085836.csv"
                                # and we're looking for "20251206_085836_experiments_project_20251206_085836.csv"
                                # we need to try different URI patterns

                                # For now, use simple format and let the lookup handle matching
                                found_uri = simple_uri
                                self._resources[found_uri] = file_path

                        # Add to resources if not already there
                        if found_uri not in resources:
                            resources[found_uri] = {
                                "uri": found_uri,
                                "name": file_path.name,
                                "mimeType": self.guess_mime_type(file_path),
                                "description": f"Generated file: {file_path.name}",
                            }

        return resources

    def guess_mime_type(self, file_path: Path) -> str:
        """Guess MIME type from file extension."""
        ext = file_path.suffix.lower()
        mime_types = {
            ".csv": "text/csv",
            ".json": "application/json",
            ".txt": "text/plain",
            ".md": "text/markdown",
            ".html": "text/html",
            ".xml": "application/xml",
            ".pdf": "application/pdf",
            ".png": "image/png",
            ".jpg": "image/jpeg",
            ".jpeg": "image/jpeg",
        }
        return mime_types.get(ext, "application/octet-stream")

    def _guess_mime_type(self, file_path: Path) -> str:
        """Alias for guess_mime_type for backward compatibility."""
        return self.guess_mime_type(file_path)

    def cleanup_old_files(self, max_age_hours: int = 24):
        """
        Clean up files older than max_age_hours.

        Args:
            max_age_hours: Maximum age in hours for files to keep
        """
        import time

        current_time = time.time()
        max_age_seconds = max_age_hours * 3600

        to_remove = []
        for uri, file_path in self._resources.items():
            if file_path.exists():
                file_age = current_time - file_path.stat().st_mtime
                if file_age > max_age_seconds:
                    try:
                        file_path.unlink()
                        to_remove.append(uri)
                    except Exception:
                        pass  # Ignore errors during cleanup

        for uri in to_remove:
            del self._resources[uri]


# Global resource manager instance
_resource_manager: Optional[ResourceManager] = None


def get_resource_manager() -> ResourceManager:
    """Get or create the global resource manager instance."""
    global _resource_manager
    if _resource_manager is None:
        # Try to use workspace directory if available, otherwise use temp
        workspace_dir = os.environ.get("COMET_MCP_RESOURCE_DIR")
        _resource_manager = ResourceManager(base_dir=workspace_dir)
    return _resource_manager
