app/routers/system.py

"""
System API Router - System monitoring and resource management
"""

from fastapi import APIRouter, HTTPException
from pydantic import BaseModel, Field
from typing import Dict, Any, List, Optional
from datetime import datetime
import logging
import psutil
import torch
import platform
import os
import shutil

from app.config import settings

logger = logging.getLogger(__name__)
router = APIRouter(prefix="/api/system", tags=["System"])


class SystemInfo(BaseModel):
    """System information."""
    platform: str
    python_version: str
    torch_version: str
    transformers_version: str
    cuda_available: bool
    cuda_version: Optional[str]
    gpu_count: int
    gpu_names: List[str]
    cpu_count: int
    total_memory_gb: float
    

class GPUInfo(BaseModel):
    """GPU information."""
    available: bool
    count: int = 0
    names: List[str] = []
    memory_used_gb: Optional[float] = None
    memory_total_gb: Optional[float] = None
    utilization: Optional[float] = None


class ResourceUsage(BaseModel):
    """Current resource usage."""
    cpu: Dict[str, float]
    memory: Dict[str, float]
    disk: Dict[str, float]
    gpu: GPUInfo
    cache: Dict[str, Any]


class StorageInfo(BaseModel):
    """Storage information."""
    path: str
    exists: bool
    size_gb: float
    file_count: int
    

class ConfigResponse(BaseModel):
    """Configuration response."""
    max_concurrent_jobs: int
    max_upload_size_mb: int
    supported_tasks: List[str]
    supported_dataset_sources: List[str]
    default_hardware: str
    available_hardware: List[str]


@router.get("/info", response_model=SystemInfo)
async def get_system_info():
    """Get system information."""
    import transformers
    
    cuda_version = None
    if torch.cuda.is_available():
        cuda_version = torch.version.cuda
    
    gpu_names = []
    if torch.cuda.is_available():
        for i in range(torch.cuda.device_count()):
            gpu_names.append(torch.cuda.get_device_name(i))
    
    return SystemInfo(
        platform=f"{platform.system()} {platform.release()}",
        python_version=platform.python_version(),
        torch_version=torch.__version__,
        transformers_version=transformers.__version__,
        cuda_available=torch.cuda.is_available(),
        cuda_version=cuda_version,
        gpu_count=torch.cuda.device_count(),
        gpu_names=gpu_names,
        cpu_count=os.cpu_count() or 1,
        total_memory_gb=round(psutil.virtual_memory().total / (1024**3), 2)
    )


@router.get("/resources", response_model=ResourceUsage)
async def get_resource_usage():
    """Get current resource usage."""
    # CPU and memory
    cpu_percent = psutil.cpu_percent(interval=1)
    memory = psutil.virtual_memory()
    
    # Disk
    disk = shutil.disk_usage('/')
    
    # GPU info

    gpu_available = torch.cuda.is_available()
    gpu_info = GPUInfo(available=gpu_available, count=0, names=[])
    
    if torch.cuda.is_available():
        try:
            gpu_names = []
            for i in range(torch.cuda.device_count()):
                gpu_names.append(torch.cuda.get_device_name(i))
            
            gpu_memory_used = round(torch.cuda.memory_allocated() / (1024**3), 2)
            gpu_memory_total = round(torch.cuda.get_device_properties(0).total_memory / (1024**3), 2)
            
            gpu_info = GPUInfo(
                available=True,
                count=torch.cuda.device_count(),
                names=gpu_names,
                memory_used_gb=gpu_memory_used,
                memory_total_gb=gpu_memory_total,
                utilization=None
            )
        except Exception as e:
            logger.error(f"Error getting GPU info: {e}")
    
    # Cache info
    cache_total_bytes = 0
    cache_dirs = [settings.CACHE_DIR, settings.HF_CACHE_DIR]
    for cache_path in cache_dirs:
        if os.path.exists(cache_path):
            for root, dirs, files in os.walk(cache_path):
                for f in files:
                    try:
                        cache_total_bytes += os.path.getsize(os.path.join(root, f))
                    except:
                        pass
    
    return ResourceUsage(
        cpu={
            "percent": round(cpu_percent, 1)
        },
        memory={
            "percent": round(memory.percent, 1),
            "used_gb": round(memory.used / (1024**3), 2),
            "total_gb": round(memory.total / (1024**3), 2)
        },
        disk={
            "percent": round((disk.used / disk.total) * 100, 1),
            "used_gb": round(disk.used / (1024**3), 2),
            "total_gb": round(disk.total / (1024**3), 2)
        },
        gpu=gpu_info,
        cache={
            "total_bytes": cache_total_bytes
        }
    )


@router.get("/storage", response_model=List[StorageInfo])
async def get_storage_info():
    """Get storage information for important directories."""
    paths = [
        ("Models", settings.MODELS_DIR),
        ("Cache", settings.CACHE_DIR),
        ("Logs", settings.LOGS_DIR),
        ("Uploads", settings.UPLOAD_DIR),
        ("Outputs", settings.OUTPUT_DIR)
    ]
    
    result = []
    for name, path in paths:
        exists = os.path.exists(path)
        size = 0
        file_count = 0
        
        if exists:
            for root, dirs, files in os.walk(path):
                file_count += len(files)
                for f in files:
                    try:
                        size += os.path.getsize(os.path.join(root, f))
                    except:
                        pass
        
        result.append(StorageInfo(
            path=f"{name} ({path})",
            exists=exists,
            size_gb=round(size / (1024**3), 4),
            file_count=file_count
        ))
    
    return result


@router.get("/config", response_model=ConfigResponse)
async def get_config():
    """Get current configuration."""
    return ConfigResponse(
        max_concurrent_jobs=settings.MAX_CONCURRENT_JOBS,
        max_upload_size_mb=settings.MAX_UPLOAD_SIZE_MB,
        supported_tasks=settings.SUPPORTED_TASKS,
        supported_dataset_sources=settings.DATASET_SOURCES,
        default_hardware=settings.DEFAULT_HARDWARE,
        available_hardware=settings.AVAILABLE_HARDWARE
    )


@router.post("/clear-cache")
async def clear_cache():
    """Clear the model and dataset cache."""
    cache_paths = [settings.CACHE_DIR, settings.HF_CACHE_DIR]
    cleared = []
    
    for path in cache_paths:
        if os.path.exists(path):
            try:
                shutil.rmtree(path)
                os.makedirs(path, exist_ok=True)
                cleared.append(path)
            except Exception as e:
                logger.error(f"Failed to clear {path}: {e}")
    
    # Also clear torch cache if applicable
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    
    return {
        "message": "Cache cleared",
        "paths_cleared": cleared,
        "timestamp": datetime.utcnow().isoformat()
    }


@router.get("/environment")
async def get_environment():
    """Get relevant environment variables (safe to expose)."""
    safe_vars = [
        "HF_HOME",
        "HF_HUB_CACHE",
        "TRANSFORMERS_CACHE",
        "WANDB_MODE",
        "WANDB_PROJECT"
    ]
    
    env_vars = {}
    for var in safe_vars:
        env_vars[var] = os.getenv(var, "not set")
    
    # Check for tokens (just existence, not value)
    tokens_status = {
        "HF_TOKEN": bool(os.getenv("HF_TOKEN")),
        "WANDB_API_KEY": bool(os.getenv("WANDB_API_KEY"))
    }
    
    return {
        "environment_variables": env_vars,
        "tokens_configured": tokens_status
    }


@router.get("/processes")
async def get_processes():
    """Get running training processes."""
    processes = []
    
    for proc in psutil.process_iter(['pid', 'name', 'cmdline', 'cpu_percent', 'memory_percent']):
        try:
            # Look for training-related processes
            cmdline = ' '.join(proc.info['cmdline'] or [])
            if any(x in cmdline.lower() for x in ['train', 'torch', 'python']):
                processes.append({
                    "pid": proc.info['pid'],
                    "name": proc.info['name'],
                    "cpu_percent": proc.info['cpu_percent'],
                    "memory_percent": round(proc.info['memory_percent'], 2)
                })
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            continue
    
    return {
        "processes": processes[:20],
        "total": len(processes)
    }


@router.get("/health")
async def health_check():
    """Detailed health check."""
    health_status = {
        "status": "healthy",
        "timestamp": datetime.utcnow().isoformat(),
        "components": {}
    }
    
    # Check disk space
    disk = shutil.disk_usage('/')
    disk_percent = (disk.used / disk.total) * 100
    health_status["components"]["disk"] = {
        "status": "ok" if disk_percent < 90 else "warning",
        "used_percent": round(disk_percent, 1)
    }
    
    # Check memory
    memory = psutil.virtual_memory()
    health_status["components"]["memory"] = {
        "status": "ok" if memory.percent < 90 else "warning",
        "used_percent": round(memory.percent, 1)
    }
    
    # Check GPU if available
    if torch.cuda.is_available():
        try:
            gpu_mem = torch.cuda.memory_allocated() / torch.cuda.get_device_properties(0).total_memory
            health_status["components"]["gpu"] = {
                "status": "ok" if gpu_mem < 0.9 else "warning",
                "memory_used_percent": round(gpu_mem * 100, 1)
            }
        except:
            health_status["components"]["gpu"] = {"status": "unavailable"}
    
    # Overall status
    if any(c.get("status") == "warning" for c in health_status["components"].values()):
        health_status["status"] = "warning"
    
    return health_status


@router.get("/logs")
async def get_system_logs(
    lines: int = 100,
    level: str = "INFO"
):
    """Get recent system logs."""
    log_file = os.path.join(settings.LOGS_DIR, "app.log")
    
    if not os.path.exists(log_file):
        return {"message": "No log file found", "logs": []}
    
    try:
        with open(log_file, 'r') as f:
            all_lines = f.readlines()
            
        # Filter by level if specified
        if level.upper() != "ALL":
            filtered = [l for l in all_lines if level.upper() in l]
        else:
            filtered = all_lines
            
        return {
            "total_lines": len(filtered),
            "logs": filtered[-lines:]
        }
    except Exception as e:
        return {"error": str(e), "logs": []}