"""FastAPI server for Z-Explorer GUI mode.

This module provides a REST API with Server-Sent Events (SSE) for real-time
progress streaming. It serves as the backend for the Tauri/browser GUI.

Default port: 8345

Usage:
    # As module
    from z_explorer.server import serve
    serve(port=8345)
    
    # Or via CLI
    z-explorer --server --port 8345
"""

import asyncio
import json
import sys
from pathlib import Path
from typing import Optional
from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import JSONResponse, FileResponse
from pydantic import BaseModel
from sse_starlette.sse import EventSourceResponse
from loguru import logger

from z_explorer.core.types import (
    GenerationRequest,
    GenerationResult,
    ProgressEvent,
    GpuInfo,
    VariableInfo,
)
from z_explorer.core.generator import generate
from z_explorer.cli import __version__

# Configure loguru for server
logger.remove()  # Remove default handler
logger.add(
    sys.stderr,
    format="<green>{time:HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>",
    level="DEBUG",
    colorize=True,
)
logger.add(
    "logs/server_{time:YYYY-MM-DD}.log",
    rotation="1 day",
    retention="7 days",
    level="DEBUG",
    format="{time:YYYY-MM-DD HH:mm:ss} | {level: <8} | {name}:{function}:{line} - {message}",
)

# Default configuration
DEFAULT_PORT = 8345
DEFAULT_HOST = "127.0.0.1"


def _get_output_dir() -> Path:
    """Get the output directory for generated images."""
    import os
    output_dir = Path(os.getenv("LOCAL_OUTPUT_DIR", "./output"))
    output_dir.mkdir(parents=True, exist_ok=True)
    return output_dir


def _get_gui_dist_dir() -> Optional[Path]:
    """Get the GUI dist directory if it exists."""
    # Check relative to this module (for installed package)
    module_dir = Path(__file__).parent
    gui_dist = module_dir / "gui" / "dist"
    if gui_dist.exists():
        return gui_dist
    
    # Check for development mode (running from source)
    dev_dist = module_dir.parent.parent / "gui" / "dist"
    if dev_dist.exists():
        return dev_dist
    
    return None


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifespan handler."""
    # Startup
    output_dir = _get_output_dir()
    print(f"[output] Output directory: {output_dir.absolute()}")
    
    # Mount output directory for serving generated images
    if output_dir.exists():
        app.mount("/output", StaticFiles(directory=str(output_dir)), name="output")
    
    # Mount GUI dist directory if it exists
    gui_dist = _get_gui_dist_dir()
    if gui_dist:
        print(f"[gui] Serving GUI from: {gui_dist}")
        # Mount assets directory
        assets_dir = gui_dist / "assets"
        if assets_dir.exists():
            app.mount("/assets", StaticFiles(directory=str(assets_dir)), name="assets")
    else:
        print("[gui] No GUI dist found - API only mode")
    
    yield
    # Shutdown
    print("[shutdown] Server shutting down...")


# Create FastAPI app
app = FastAPI(
    title="Z-Explorer API",
    description="Local AI image generation with Z-Image-Turbo + Qwen3-4B",
    version=__version__,
    lifespan=lifespan,
)

# Enable CORS for Tauri and browser frontends
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allow all origins for local desktop app
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# ============================================================================
# Health & Status Endpoints
# ============================================================================

@app.get("/", include_in_schema=False)
async def serve_gui():
    """Serve the GUI index.html."""
    gui_dist = _get_gui_dist_dir()
    if gui_dist:
        index_file = gui_dist / "index.html"
        if index_file.exists():
            return FileResponse(index_file)
    return JSONResponse(
        {"error": "GUI not found. Run 'npm run build' in src/z_explorer/gui first."},
        status_code=404
    )


@app.get("/api/health")
async def health_check():
    """Health check endpoint."""
    return {
        "status": "ok",
        "version": __version__,
        "service": "z-explorer",
    }


@app.get("/api/gpu", response_model=GpuInfo)
async def get_gpu_info():
    """Get GPU memory information."""
    try:
        from z_explorer.image_generator import get_gpu_memory_info
        info = get_gpu_memory_info()
        
        if "error" in info:
            return GpuInfo(available=False, error=info.get("error"))
        
        return GpuInfo(
            available=True,
            device_name=info.get("device_name"),
            allocated_gb=info.get("allocated_gb"),
            reserved_gb=info.get("reserved_gb"),
            total_gb=info.get("total_gb"),
            free_gb=info.get("free_gb"),
        )
    except Exception as e:
        return GpuInfo(available=False, error=str(e))


# ============================================================================
# Variables Endpoints
# ============================================================================

class VariablesResponse(BaseModel):
    """Response model for variables listing."""
    success: bool
    count: int
    variables: list[VariableInfo]
    error: Optional[str] = None


@app.get("/api/variables", response_model=VariablesResponse)
async def list_variables():
    """List all available prompt variables."""
    try:
        from z_explorer.models.prompt_vars import load_prompt_vars
        
        prompt_vars = load_prompt_vars()
        
        variables = []
        for var_id, var in prompt_vars.items():
            variables.append(VariableInfo(
                id=var_id,
                description=var.description,
                count=len(var.values),
                sample=var.values[:5] if var.values else [],
                file_path=var.file_path or "",
            ))
        
        # Sort by ID
        variables.sort(key=lambda v: v.id)
        
        return VariablesResponse(
            success=True,
            count=len(variables),
            variables=variables,
        )
    except Exception as e:
        return VariablesResponse(
            success=False,
            count=0,
            variables=[],
            error=str(e),
        )


# ============================================================================
# Model Management Endpoints
# ============================================================================

class UnloadResponse(BaseModel):
    """Response model for model unloading."""
    success: bool
    llm_unloaded: bool
    image_model_unloaded: bool
    cuda_cache_cleared: bool
    errors: list[str]


@app.post("/api/unload", response_model=UnloadResponse)
async def unload_models():
    """Unload all models to free GPU memory."""
    result = UnloadResponse(
        success=False,
        llm_unloaded=False,
        image_model_unloaded=False,
        cuda_cache_cleared=False,
        errors=[],
    )
    
    # Unload LLM
    try:
        from z_explorer.llm_provider import unload_model
        unload_model()
        result.llm_unloaded = True
    except Exception as e:
        result.errors.append(f"LLM unload: {e}")
    
    # Unload Image model
    try:
        from z_explorer.image_generator import unload_pipeline
        unload_pipeline()
        result.image_model_unloaded = True
    except Exception as e:
        result.errors.append(f"Image model unload: {e}")
    
    # Clear CUDA cache
    try:
        import gc
        gc.collect()
        
        import torch
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.synchronize()
            result.cuda_cache_cleared = True
    except Exception as e:
        result.errors.append(f"CUDA cache: {e}")
    
    result.success = result.llm_unloaded or result.image_model_unloaded
    return result


# ============================================================================
# Image Listing Endpoint
# ============================================================================

class ImageInfo(BaseModel):
    """Information about a generated image."""
    path: str
    name: str
    url: str
    modified: float
    prompt: Optional[str] = None


class ImagesResponse(BaseModel):
    """Response model for images listing."""
    images: list[ImageInfo]
    count: int


@app.get("/api/images", response_model=ImagesResponse)
async def list_images():
    """List all generated images, newest first."""
    output_dir = _get_output_dir()
    
    images = []
    for f in output_dir.glob("*.png"):
        # Try to read the prompt from the corresponding .txt file
        prompt = None
        prompt_file = f.with_suffix('.txt')
        if prompt_file.exists():
            try:
                prompt = prompt_file.read_text(encoding='utf-8').strip()
            except Exception:
                pass
        
        images.append(ImageInfo(
            path=str(f),
            name=f.name,
            url=f"/output/{f.name}",
            modified=f.stat().st_mtime,
            prompt=prompt,
        ))
    
    # Sort by modification time, newest first
    images.sort(key=lambda x: x.modified, reverse=True)
    
    return ImagesResponse(images=images, count=len(images))


# ============================================================================
# Generation Endpoints with SSE
# ============================================================================

async def _generate_event_stream(request: GenerationRequest):
    """
    Shared event stream generator for image generation.
    
    Yields SSE events with JSON data for browser-friendly format.
    """
    logger.info(f"🎬 Starting SSE stream for prompt: {request.prompt[:50]}...")
    
    progress_queue: asyncio.Queue[ProgressEvent] = asyncio.Queue()
    result_holder: list[GenerationResult] = []
    error_holder: list[str] = []
    events_sent = 0
    events_queued = 0
    
    def on_progress(event: ProgressEvent):
        """Callback to queue progress events."""
        nonlocal events_queued
        try:
            progress_queue.put_nowait(event)
            events_queued += 1
            logger.debug(f"📤 [{events_queued}] Queued: stage={event.stage}, msg={event.message}, progress={event.progress}")
        except Exception as e:
            logger.error(f"❌ Failed to queue progress event: {e}")
    
    def run_generation():
        """Run generation in thread pool."""
        logger.info("🔧 Generation thread started")
        try:
            result = generate(request, on_progress=on_progress)
            result_holder.append(result)
            logger.info(f"✅ Generation complete: {len(result.images)} images, prompts: {result.final_prompts}")
        except Exception as e:
            logger.error(f"❌ Generation failed: {e}")
            error_holder.append(str(e))
    
    # Start generation in background thread
    loop = asyncio.get_event_loop()
    gen_task = loop.run_in_executor(None, run_generation)
    logger.debug("Generation task submitted to executor")
    
    def event_to_json(event: ProgressEvent) -> str:
        """Convert ProgressEvent to JSON for SSE."""
        path = event.data.get("path") if event.data else None
        return json.dumps({
            "stage": event.stage,
            "message": event.message,
            "progress": event.progress,
            "path": path,
            "data": event.data,
        })
    
    # Stream progress events while generation is running
    last_percent = -1
    poll_count = 0
    while not gen_task.done():
        try:
            event = await asyncio.wait_for(progress_queue.get(), timeout=0.2)
            current_percent = event.progress if event.progress is not None else 0
            
            events_sent += 1
            json_data = event_to_json(event)
            logger.debug(f"📨 [{events_sent}] Sending SSE: {event.stage} ({event.progress}%)")
            yield {"data": json_data}
            last_percent = current_percent
            await asyncio.sleep(0)
                
        except asyncio.TimeoutError:
            poll_count += 1
            if poll_count % 10 == 0:
                logger.trace(f"⏳ Polling... (queue size: {progress_queue.qsize()}, task done: {gen_task.done()})")
            continue
        except Exception as e:
            logger.error(f"❌ Error in event loop: {e}")
            break
    
    logger.debug(f"🏁 Generation task done. Draining queue (size: {progress_queue.qsize()})")
    
    # Drain remaining events
    drain_count = 0
    while not progress_queue.empty():
        try:
            event = progress_queue.get_nowait()
            drain_count += 1
            events_sent += 1
            logger.debug(f"📨 [DRAIN {drain_count}] Sending: {event.stage}")
            yield {"data": event_to_json(event)}
        except Exception as e:
            logger.error(f"❌ Error draining queue: {e}")
            break
    
    logger.info(f"📊 SSE Stats: queued={events_queued}, sent={events_sent}, drained={drain_count}")
    
    # Send final result or error
    if error_holder:
        logger.error(f"❌ Sending error: {error_holder[0]}")
        yield {"data": json.dumps({"stage": "error", "message": error_holder[0]})}
    elif result_holder:
        result = result_holder[0]
        logger.info(f"✅ Sending complete event with {len(result.images)} images")
        yield {
            "data": json.dumps({
                "stage": "complete",
                "message": f"Generated {len(result.images)} image(s)",
                "images": result.images,
                "final_prompts": result.final_prompts,
            }),
        }
    else:
        yield {"data": json.dumps({"stage": "error", "message": "Unknown error"})}


@app.get("/api/generate/stream")
async def generate_images_stream(
    prompt: str,
    count: int = 1,
    width: int = 1024,
    height: int = 1024,
    seed: Optional[int] = None,
    enhance: bool = False,
    enhance_instruction: str = "",
):
    """
    Generate images via SSE stream (GET for browser EventSource).
    
    Query Parameters:
        prompt: The image prompt
        count: Number of images to generate (default: 1)
        width: Image width (default: 1024)
        height: Image height (default: 1024)
        seed: Random seed (optional)
        enhance: Whether to enhance the prompt (default: false)
        enhance_instruction: Optional enhancement instruction
    
    Returns Server-Sent Events with JSON data.
    """
    request = GenerationRequest(
        prompt=prompt,
        count=count,
        width=width,
        height=height,
        seed=seed,
        enhance=enhance,
        enhance_instruction=enhance_instruction,
    )
    return EventSourceResponse(
        _generate_event_stream(request),
        ping=1,  # Send keepalive ping every second to prevent buffering
    )


@app.post("/api/generate")
async def generate_images_post(request: GenerationRequest):
    """
    Generate images via SSE stream (POST for programmatic use).
    
    Request Body:
        prompt: The image prompt
        count: Number of images to generate (default: 1)
        width: Image width (default: 1024)
        height: Image height (default: 1024)
        seed: Random seed (optional)
        enhance: Whether to enhance the prompt (default: false)
        enhance_instruction: Optional enhancement instruction
    
    Returns Server-Sent Events with JSON data.
    """
    return EventSourceResponse(
        _generate_event_stream(request),
        ping=1,  # Send keepalive ping every second to prevent buffering
    )


# ============================================================================
# Server Entry Point
# ============================================================================

def serve(host: str = DEFAULT_HOST, port: int = DEFAULT_PORT):
    """Start the Z-Explorer API server.
    
    Args:
        host: Host to bind to (default: 127.0.0.1)
        port: Port to bind to (default: 8345)
    """
    import uvicorn
    
    print(f"[startup] Starting Z-Explorer server on http://{host}:{port}")
    print(f"[docs] API docs at http://{host}:{port}/docs")
    
    uvicorn.run(
        app,
        host=host,
        port=port,
        log_level="info",
    )


if __name__ == "__main__":
    serve()

