"""
Main FastAPI application for Unsloth UI Backend
"""
import secrets
import shutil
from contextlib import asynccontextmanager

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from pathlib import Path
from datetime import datetime

# Import routers
from routes import training_router, models_router, inference_router, datasets_router, auth_router
from auth import storage

UNSLOTH_CACHE_DIR = Path(__file__).parent / "unsloth_compiled_cache"


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Startup: print setup token if needed. Shutdown: clean up compiled cache."""
    if not storage.is_initialized():
        setup_token = secrets.token_urlsafe(32)
        storage.save_setup_token(setup_token)
        print("\n" + "=" * 60)
        print("FIRST-TIME SETUP REQUIRED")
        print("Use this one-time setup token to create your admin account:\n")
        print(f"    {setup_token}\n")
        print("This token can only be used once.")
        print("=" * 60 + "\n")
    yield
    # Cleanup: remove Unsloth compiled cache on shutdown
    shutil.rmtree(UNSLOTH_CACHE_DIR, ignore_errors=True)


# Create FastAPI app
app = FastAPI(
    title="Unsloth UI Backend",
    version="1.0.0",
    description="Backend API for Unsloth UI - Training and Model Management",
    lifespan=lifespan,
)

# CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # In production, specify allowed origins
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ============ Register API Routes ============

# Register routers
app.include_router(auth_router, prefix="/api/auth", tags=["auth"])
app.include_router(training_router, prefix="/api/train", tags=["training"])
app.include_router(models_router, prefix="/api/models", tags=["models"])
app.include_router(inference_router, prefix="/api/inference", tags=["inference"])
app.include_router(datasets_router, prefix="/api/datasets", tags=["datasets"])


# ============ Health and System Endpoints ============

@app.get("/api/health")
async def health_check():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "timestamp": datetime.now().isoformat(),
        "service": "Unsloth UI Backend"
    }


@app.get("/api/system")
async def get_system_info():
    """Get system information"""
    import torch
    import platform
    import psutil

    # GPU Info
    gpu_info = {"available": False, "devices": []}
    if torch.cuda.is_available():
        gpu_info["available"] = True
        for i in range(torch.cuda.device_count()):
            props = torch.cuda.get_device_properties(i)
            gpu_info["devices"].append(
                {
                    "index": i,
                    "name": props.name,
                    "memory_total_gb": round(props.total_memory / 1e9, 2),
                }
            )

    # CPU & Memory
    memory = psutil.virtual_memory()

    return {
        "platform": platform.platform(),
        "python_version": platform.python_version(),
        "cpu_count": psutil.cpu_count(),
        "memory": {
            "total_gb": round(memory.total / 1e9, 2),
            "available_gb": round(memory.available / 1e9, 2),
            "percent_used": memory.percent,
        },
        "gpu": gpu_info,
    }


# ============ Serve Frontend (Optional) ============

def setup_frontend(app: FastAPI, build_path: Path):
    """Mount frontend static files (optional)"""
    if build_path.exists():
        # Mount assets
        assets_dir = build_path / "assets"
        if assets_dir.exists():
            app.mount("/assets", StaticFiles(directory=assets_dir), name="assets")

        @app.get("/")
        async def serve_root():
            return FileResponse(build_path / "index.html")

        @app.get("/{full_path:path}")
        async def serve_frontend(full_path: str):
            if full_path.startswith("api"):
                return {"error": "API endpoint not found"}

            file_path = build_path / full_path
            if file_path.is_file():
                return FileResponse(file_path)

            return FileResponse(build_path / "index.html")

        return True
    return False