#!/usr/bin/env python3
"""
Run isolation module for multi-user safety in web deployment.
Each run gets its own isolated directory based on UUID.
"""

import os
import uuid
import datetime
from pathlib import Path
from typing import Optional, Dict, Any, TYPE_CHECKING

if TYPE_CHECKING:
    # Avoid circular imports, only used for type hints
    from fastapi import Request
import json


def generate_run_id() -> str:
    """Generate a unique run ID with timestamp prefix for easier identification."""
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    short_uuid = str(uuid.uuid4()).replace("-", "")[:12]
    return f"run_{timestamp}_{short_uuid}"


def get_workdir_base() -> Path:
    """Get the base working directory from environment variable or use default."""
    workdir = os.getenv("UHT_WORKDIR", "/var/lib/uht-discovery")
    return Path(workdir)


def get_current_user(request: Optional[Any] = None) -> str:
    """
    Get the current user identifier.
    
    Checks in order:
    1. Request object (if provided, checks request.username from auth)
    2. Thread-local authenticated username (set by user_auth.validate_user())
    3. UHT_CURRENT_USER environment variable
    4. "anonymous" if in dev mode (UHT_DEV=1)
    5. "shared" as final fallback
    
    Args:
        request: Optional Request object. If provided, will extract username from authenticated session.
    
    Returns:
        str: User identifier for cache isolation and run ownership
    """
    # Check Request object if provided (most reliable for async context)
    if request is not None:
        try:
            # NiceGUI/FastAPI may store authenticated username in request.username when using custom auth
            if hasattr(request, 'username') and request.username:
                return request.username
            # Alternative: check request headers or session
            if hasattr(request, 'client') and hasattr(request.client, 'host'):
                # Request is available, try to extract from auth headers if needed
                pass
        except Exception:
            # Request parsing failed, continue to other checks
            pass
    
    # Check thread-local storage for authenticated username (from custom auth function)
    # Note: Thread-local works well for synchronous flows, but request-based is more reliable
    # but request-based is more reliable for async operations
    try:
        from uht_discovery.user_auth import get_authenticated_username
        user = get_authenticated_username()
        if user:
            return user
    except ImportError:
        # user_auth module not available (e.g., in CLI context), continue to other checks
        pass
    except Exception:
        # Thread-local access failed, continue to other checks
        pass
    
    # Check for explicit user setting
    user = os.getenv("UHT_CURRENT_USER")
    if user:
        return user
    
    # In dev mode, use anonymous
    if os.getenv("UHT_DEV") == "1":
        return "anonymous"
    
    # Default to shared for backward compatibility
    return "shared"


def get_cache_dir(user_id: Optional[str] = None) -> Path:
    """
    Get the cache directory path for embeddings.
    
    The cache is stored in persistent storage and shared across runs globally.
    
    Args:
        user_id: Optional user identifier. Ignored to keep cache global.
    
    Returns:
        Path: Path to the cache directory (e.g., /persistent/uht-discovery/cache/embeddings/esm2)
    
    For local development (UHT_WORKDIR not set), returns relative path for backward compatibility.
    """
    _ = user_id
    
    base = get_workdir_base()
    
    # If using default workdir (local dev), use relative path for backward compatibility
    if str(base) == "/var/lib/uht-discovery" and not os.getenv("UHT_WORKDIR"):
        # Local development - use relative path
        return Path("embeddings") / "esm2"
    
    # Production/Modal - use persistent storage
    cache_dir = base / "cache" / "embeddings" / "esm2"
    cache_dir.mkdir(parents=True, exist_ok=True)
    return cache_dir


def create_run_directory(run_id: str) -> Path:
    """
    Create a run-scoped directory structure.
    Returns the run directory path.
    """
    base = get_workdir_base()
    run_dir = base / "runs" / run_id
    run_dir.mkdir(parents=True, exist_ok=True)
    
    # Create standard subdirectories
    (run_dir / "inputs").mkdir(exist_ok=True)
    (run_dir / "results").mkdir(exist_ok=True)
    (run_dir / "metadata").mkdir(exist_ok=True)
    
    return run_dir


def save_run_metadata(run_id: str, metadata: Dict[str, Any]) -> None:
    """
    Save run metadata to the run directory.
    
    Automatically adds username to metadata if not already present.
    """
    run_dir = get_workdir_base() / "runs" / run_id
    metadata_file = run_dir / "metadata" / "run_info.json"
    
    metadata["run_id"] = run_id
    metadata["created_at"] = datetime.datetime.now().isoformat()
    
    # Add username if not already present
    if "username" not in metadata:
        metadata["username"] = get_current_user()
    
    with open(metadata_file, 'w') as f:
        json.dump(metadata, f, indent=2)


def get_run_metadata(run_id: str) -> Optional[Dict[str, Any]]:
    """Load run metadata from the run directory."""
    run_dir = get_workdir_base() / "runs" / run_id
    metadata_file = run_dir / "metadata" / "run_info.json"
    
    if not metadata_file.exists():
        return None
    
    with open(metadata_file, 'r') as f:
        return json.load(f)


def validate_run_path(path: str, run_id: str) -> bool:
    """
    Validate that a path is within the run directory.
    Prevents path traversal attacks.
    """
    run_dir = get_workdir_base() / "runs" / run_id
    try:
        resolved_path = Path(path).resolve()
        run_dir_resolved = run_dir.resolve()
        # Check if resolved path is within run directory
        return str(resolved_path).startswith(str(run_dir_resolved))
    except (ValueError, OSError):
        return False


def get_run_scoped_project_dir(module: str, run_id: str) -> str:
    """
    Get a project directory name that maps to the run-scoped directory.
    For backward compatibility with existing code that expects project names,
    we use a consistent naming scheme based on the run_id.
    """
    # Use run_id as the project name, but ensure paths resolve to run directory
    return run_id


def setup_run_environment(run_id: str) -> Dict[str, str]:
    """
    Set up environment variables for a run to redirect all I/O to run-scoped directories.
    Returns a dict of environment variables to set.
    """
    run_dir = get_workdir_base() / "runs" / run_id
    
    # Set environment variables that the modules will use
    # The modules use project_dir() which checks env vars first, then config
    env_vars = {
        "UHT_RUN_ID": run_id,
        "UHT_RUN_DIR": str(run_dir),
    }
    
    # Set module-specific project directories to use run_id
    # This causes project_dir() to return run_id, which then gets used in paths
    # But we need to ensure the paths resolve to run_dir
    env_vars["BLASTER_PROJECT_ID"] = run_id
    env_vars["TRIM_PROJECT_ID"] = run_id
    env_vars["PLMCLUSTV2_PROJECT_ID"] = run_id
    
    # Also override the base paths by monkey-patching or using a custom path resolver
    # Actually, we'll modify the GUI functions to use absolute paths from run_ctx
    # The env vars ensure project_dir() returns run_id
    
    return env_vars


class RunIsolationContext:
    """
    Context manager for run isolation.
    Sets up environment variables and validates paths during execution.
    """
    
    def __init__(self, run_id: str):
        self.run_id = run_id
        self.run_dir = create_run_directory(run_id)
        self.original_env = {}
        self.env_vars = setup_run_environment(run_id)
    
    def __enter__(self):
        # Save original environment
        for key in self.env_vars:
            self.original_env[key] = os.environ.get(key)
            os.environ[key] = self.env_vars[key]
        
        # Save original working directory
        self.original_cwd = os.getcwd()
        
        # Create directory structure within run directory that matches expected layout
        # The modules expect inputs/<module>/<project>/ and results/<module>/<project>/
        # We'll create these within the run directory
        for module in ["blaster", "trim", "plmclustv2"]:
            # Create inputs structure
            module_input = self.run_dir / "inputs" / module / self.run_id
            module_input.mkdir(parents=True, exist_ok=True)
            # Create results structure
            module_results = self.run_dir / "results" / module / self.run_id
            module_results.mkdir(parents=True, exist_ok=True)
        
        # Change to run directory so relative paths resolve correctly
        # Since we use queuing with concurrency_limit=1, this is thread-safe
        os.chdir(str(self.run_dir))
        
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        # Restore original working directory first
        os.chdir(self.original_cwd)
        
        # Restore original environment
        for key, value in self.original_env.items():
            if value is None:
                os.environ.pop(key, None)
            else:
                os.environ[key] = value
        
        return False  # Don't suppress exceptions
    
    def get_input_path(self, module: str, filename: str = "") -> Path:
        """Get input path within the run directory."""
        if filename:
            return self.run_dir / "inputs" / module / self.run_id / filename
        return self.run_dir / "inputs" / module / self.run_id
    
    def get_result_path(self, module: str, filename: str = "") -> Path:
        """Get result path within the run directory."""
        if filename:
            return self.run_dir / "results" / module / self.run_id / filename
        return self.run_dir / "results" / module / self.run_id
    
    def validate_and_relocate_path(self, path: str, module: str, is_input: bool = True) -> Path:
        """
        Validate a path and relocate it to the run-scoped directory.
        Used when copying files into the run directory.
        """
        source_path = Path(path)
        if is_input:
            dest = self.get_input_path(module, source_path.name)
        else:
            dest = self.get_result_path(module, source_path.name)
        return dest
