"""
Oprel Daemon Server - Persistent model caching for fast inference.

This server keeps models loaded in memory, eliminating the 2-minute
startup time on every script execution. It also manages conversation
history and model discovery.

Usage:
    oprel serve                  # Start server on default port 11434
    oprel serve --port 8080      # Start on custom port

API Endpoints:
    POST /load      - Load a model into cache
    POST /generate  - Generate text (supports chat history)
    GET /models     - List all loaded and available cached models
    DELETE /unload/{model_id} - Unload a specific model
    GET /health     - Health check
    
    # Conversation APIs
    GET /conversations - List active conversations
    GET /conversations/{id} - Get conversation history
    DELETE /conversations/{id} - Delete conversation
    POST /conversations/{id}/reset - Reset conversation
"""

import os
import sys
import signal
import atexit
import uuid
import json
import time as time_module
import asyncio
from datetime import datetime
from typing import Dict, Optional, Any, List
from contextlib import asynccontextmanager
from pathlib import Path
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from starlette.middleware.base import BaseHTTPMiddleware

from oprel.core.config import Config
from oprel.utils.logging import get_logger

# Initialize logger
logger = get_logger(__name__)

# Import Model with use_server=False to avoid circular dependency
# The server itself uses direct mode internally

# Initialize Config
CONFIG = Config()
CONFIG.ensure_dirs()

# ANSI color codes for terminal output
class Colors:
    RESET = "\033[0m"
    BOLD = "\033[1m"
    GREEN = "\033[92m"      # 2xx success
    CYAN = "\033[96m"       # 3xx redirect
    YELLOW = "\033[93m"     # 4xx client error
    RED = "\033[91m"        # 5xx server error
    BLUE = "\033[94m"       # Method
    MAGENTA = "\033[95m"    # Path
    WHITE = "\033[97m"      # IP
    GRAY = "\033[90m"       # Timestamp prefix


def get_status_color(status_code: int) -> str:
    """Get color based on HTTP status code"""
    if 200 <= status_code < 300: return Colors.GREEN
    elif 300 <= status_code < 400: return Colors.CYAN
    elif 400 <= status_code < 500: return Colors.YELLOW
    else: return Colors.RED


def format_duration(duration_ms: float) -> str:
    """Format duration in human-readable format"""
    if duration_ms < 1: return f"{duration_ms * 1000:.2f}µs"
    elif duration_ms < 1000: return f"{duration_ms:.2f}ms"
    else: return f"{duration_ms / 1000:.2f}s"


class GinStyleLoggingMiddleware(BaseHTTPMiddleware):
    """Middleware that logs requests in Gin-like format with colors"""
    async def dispatch(self, request: Request, call_next):
        start_time = time_module.perf_counter()
        response = await call_next(request)
        duration = (time_module.perf_counter() - start_time) * 1000  # ms
        
        status_code = response.status_code
        method = request.method
        path = request.url.path
        client_ip = request.client.host if request.client else "unknown"
        timestamp = datetime.now().strftime("%Y/%m/%d - %H:%M:%S")
        status_color = get_status_color(status_code)
        
        log_line = (
            f"{Colors.GRAY}[OPREL]{Colors.RESET} "
            f"{timestamp} "
            f"{Colors.GRAY}|{Colors.RESET} "
            f"{status_color}{Colors.BOLD}{status_code}{Colors.RESET} "
            f"{Colors.GRAY}|{Colors.RESET} "
            f"{format_duration(duration):>12} "
            f"{Colors.GRAY}|{Colors.RESET} "
            f"{Colors.WHITE}{client_ip:>15}{Colors.RESET} "
            f"{Colors.GRAY}|{Colors.RESET} "
            f"{Colors.BLUE}{method:<8}{Colors.RESET} "
            f"{Colors.MAGENTA}\"{path}\"{Colors.RESET}"
        )
        print(log_line)
        return response


# --- Data Models ---

class LoadRequest(BaseModel):
    """Request to load a model"""
    model_id: str
    quantization: Optional[str] = None
    max_memory_mb: Optional[int] = None
    backend: str = "llama.cpp"


class GenerateRequest(BaseModel):
    """Request to generate text"""
    model_id: str
    prompt: str
    max_tokens: int = 512
    temperature: float = 0.7
    stream: bool = False
    
    # New fields for conversational API
    conversation_id: Optional[str] = None
    system_prompt: Optional[str] = None
    reset_conversation: bool = False


class GenerateResponse(BaseModel):
    """Response from text generation"""
    text: str
    model_id: str
    conversation_id: str
    message_count: int


class ModelInfo(BaseModel):
    """Information about a model"""
    model_id: str
    quantization: Optional[str]
    backend: str
    loaded: bool
    size_gb: Optional[float] = None
    name: Optional[str] = None


class ChatMessage(BaseModel):
    role: str
    content: str


class ConversationInfo(BaseModel):
    id: str
    created_at: str
    last_updated: str
    message_count: int
    model_id: str


class HealthResponse(BaseModel):
    status: str
    models_loaded: int
    active_conversations: int


class LoadResponse(BaseModel):
    success: bool
    model_id: str
    message: str


class UnloadResponse(BaseModel):
    success: bool
    model_id: str
    message: str


class UnloadRequest(BaseModel):
    """Request to unload a model"""
    model_id: str





# --- Global State ---

_models: Dict[str, Any] = {}  # model_id -> Model instance
_model_configs: Dict[str, dict] = {}  # model_id -> config info
_model_last_used: Dict[str, float] = {}  # model_id -> timestamp of last use
_conversations: Dict[str, List[Dict[str, str]]] = {} # conversation_id -> list of messages
_conversation_meta: Dict[str, Dict] = {} # conversation_id -> metadata
MAX_CONVERSATIONS = 100
MAX_HISTORY_MSGS = 50
IDLE_TIMEOUT_SECONDS = 15 * 60  # 15 minutes
_cleanup_task: Optional[asyncio.Task] = None


# --- Helper Functions ---

def _cleanup_models():
    """Unload all models on shutdown"""
    global _models, _model_last_used
    for model_id, model in list(_models.items()):
        try:
            model.unload()
            print(f"Unloaded model: {model_id}")
        except Exception as e:
            print(f"Error unloading {model_id}: {e}")
    _models.clear()
    _model_configs.clear()
    _model_last_used.clear()


def _unload_idle_model(model_id: str):
    """
    Unload idle model and force-clear GPU memory.
    This is more aggressive than regular unload() - it terminates the backend process.
    """
    global _models, _model_configs, _model_last_used
    
    if model_id not in _models:
        return
    
    try:
        model = _models[model_id]
        logger.info(f"Unloading idle model: {model_id} (forcing backend termination)")
        
        # CRITICAL: Force-stop backend process to free GPU memory
        # model.unload() in server mode keeps the process running,
        # but idle cleanup needs to actually free GPU memory
        
        if hasattr(model, '_process') and model._process:
            logger.debug(f"Stopping backend process for {model_id}")
            model._process.stop(force=False)  # Graceful stop
            model._process = None
        
        # Stop monitor if exists
        if hasattr(model, '_monitor') and model._monitor:
            model._monitor.stop()
            model._monitor = None
        
        # Cleanup PyTorch backend if used
        if hasattr(model, '_pytorch_backend') and model._pytorch_backend:
            logger.debug("Unloading PyTorch backend...")
            model._pytorch_backend.unload()
            model._pytorch_backend = None
        
        # Mark as unloaded
        model._loaded = False
        
        # Remove from caches
        del _models[model_id]
        if model_id in _model_configs:
            del _model_configs[model_id]
        if model_id in _model_last_used:
            del _model_last_used[model_id]
        
        # Force Python garbage collection
        import gc
        gc.collect()
        
        # Clear CUDA cache if available (for PyTorch models)
        try:
            import torch
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
                torch.cuda.synchronize()
                logger.debug("CUDA cache cleared")
        except ImportError:
            pass
        except Exception as e:
            logger.debug(f"Could not clear CUDA cache: {e}")
        
        logger.info(f"✓ Idle model {model_id} unloaded, GPU memory freed")
        
    except Exception as e:
        logger.error(f"Error unloading idle model {model_id}: {e}")
        import traceback
        logger.error(traceback.format_exc())


async def _monitor_idle_models():
    """Background task to monitor and unload idle models"""
    global _models, _model_last_used
    
    while True:
        try:
            await asyncio.sleep(60)  # Check every minute
            
            current_time = time_module.time()
            models_to_unload = []
            
            # Find idle models
            for model_id in list(_models.keys()):
                last_used = _model_last_used.get(model_id, 0)
                idle_time = current_time - last_used
                
                if idle_time > IDLE_TIMEOUT_SECONDS:
                    models_to_unload.append(model_id)
                    logger.info(
                        f"Model {model_id} has been idle for "
                        f"{idle_time / 60:.1f} minutes (>15 min threshold)"
                    )
            
            # Unload idle models
            for model_id in models_to_unload:
                _unload_idle_model(model_id)
                
        except asyncio.CancelledError:
            logger.info("Idle model monitoring task cancelled")
            break
        except Exception as e:
            logger.error(f"Error in idle model monitor: {e}")


def _mark_model_used(model_id: str):
    """Update the last used timestamp for a model"""
    global _model_last_used
    _model_last_used[model_id] = time_module.time()


def _cleanup_conversations():
    """LRU cleanup for conversational memory"""
    global _conversations, _conversation_meta
    if len(_conversations) > MAX_CONVERSATIONS:
        # Sort by last updated (approximation, Python dicts preserve insertion order mostly)
        # Using _conversation_meta['last_updated'] would be better but expensive to sort
        # Simple FIFO removal
        to_remove = len(_conversations) - MAX_CONVERSATIONS
        keys = list(_conversations.keys())[:to_remove]
        for k in keys:
            del _conversations[k]
            if k in _conversation_meta:
                del _conversation_meta[k]


def _scan_cached_models() -> List[ModelInfo]:
    """Scan cache directory for available models"""
    available = []
    
    # First add loaded models
    for model_id, config in _model_configs.items():
        if model_id not in [m.model_id for m in available]:
             available.append(ModelInfo(
                model_id=model_id,
                quantization=config.get("quantization"),
                backend=config.get("backend", "llama.cpp"),
                loaded=True,
                name=model_id.split("/")[-1] if "/" in model_id else model_id
            ))

    # Scan cache dir
    cache_dir = CONFIG.cache_dir
    if not cache_dir.exists():
        return available
        
    try:
        # Files are stored as 'models--Author--Name'
        for model_dir in cache_dir.iterdir():
            if model_dir.is_dir() and model_dir.name.startswith("models--"):
                try:
                    # Parse model_id: models--TheBloke--Llama-2 -> TheBloke/Llama-2
                    parts = model_dir.name.split("--")
                    if len(parts) >= 3:
                        model_id = f"{parts[1]}/{parts[2]}"
                        
                        # Check snapshots
                        snapshots_dir = model_dir / "snapshots"
                        if snapshots_dir.exists():
                            for snapshot in snapshots_dir.iterdir():
                                if snapshot.is_dir():
                                    # Find GGUF files
                                    gguf_files = list(snapshot.glob("*.gguf"))
                                    for file in gguf_files:
                                        # Deduplicate if already loaded
                                        if any(m.model_id == model_id for m in available):
                                            continue
                                            
                                        size_gb = file.stat().st_size / (1024**3)
                                        
                                        # Attempt to detect quantization from filename
                                        quant = "Unknown"
                                        name_upper = file.name.upper()
                                        for q in ["Q2_K", "Q3_K_M", "Q4_K_M", "Q5_K_M", "Q6_K", "Q8_0", "F16"]:
                                             if q in name_upper:
                                                 quant = q
                                                 break
                                        
                                        available.append(ModelInfo(
                                            model_id=model_id,
                                            quantization=quant,
                                            backend="llama.cpp",
                                            loaded=False,
                                            size_gb=round(size_gb, 2),
                                            name=model_id.split("/")[-1]
                                        ))
                except Exception as e:
                    continue
    except Exception as e:
        print(f"Error scanning cache: {e}")
        
    return available


def _build_chat_prompt(model_id: str, history: List[Dict[str, str]], system_prompt: Optional[str] = None, new_user_msg: str = "") -> str:
    """Build a prompt based on model type and chat history"""
    from ..utils.chat_templates import format_chat_prompt
    
    # Build full conversation history including new message
    conversation_history = []
    
    # Add existing history
    conversation_history.extend(history)
    
    # Use the comprehensive chat template system
    return format_chat_prompt(
        model_id=model_id,
        user_message=new_user_msg,
        system_prompt=system_prompt,
        conversation_history=conversation_history
    )


# --- Startup/Shutdown ---

def _signal_handler(signum, frame):
    """Handle shutdown signals gracefully"""
    print("\nReceived shutdown signal, cleaning up...")
    _cleanup_models()
    sys.exit(0)


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Lifespan context manager for startup/shutdown"""
    global _cleanup_task
    
    print(f"{Colors.GREEN}Oprel daemon starting...{Colors.RESET}")
    print(f"Cache Dir: {CONFIG.cache_dir}")
    print(f"Idle model timeout: {IDLE_TIMEOUT_SECONDS / 60:.0f} minutes")
    
    # Start background task to monitor idle models
    _cleanup_task = asyncio.create_task(_monitor_idle_models())
    logger.info("Started idle model monitoring task")
    
    yield
    
    # Shutdown
    print("\nReceived shutdown signal, cleaning up...")
    _cleanup_models()
    
    # Cancel background task
    if _cleanup_task:
        _cleanup_task.cancel()
        try:
            await _cleanup_task
        except asyncio.CancelledError:
            pass
    
    print(f"{Colors.YELLOW}Oprel daemon shutting down...{Colors.RESET}")
    _cleanup_models()


app = FastAPI(
    title="Oprel Daemon",
    description="Persistent model server for fast LLM inference",
    version="0.3.0",
    lifespan=lifespan,
)

app.add_middleware(GinStyleLoggingMiddleware)


# --- Endpoints ---

@app.get("/health", response_model=HealthResponse)
async def health_check():
    """Health check endpoint"""
    return HealthResponse(
        status="ok",
        models_loaded=len(_models),
        active_conversations=len(_conversations)
    )


@app.get("/models", response_model=List[ModelInfo])
async def list_models():
    """List all loaded and cached models"""
    return _scan_cached_models()


@app.post("/load", response_model=LoadResponse)
async def load_model(request: LoadRequest):
    """Load a model into the cache"""
    global _models, _model_configs
    
    if request.model_id in _models:
        # Check if backend process is still alive
        model = _models[request.model_id]
        if hasattr(model, '_process') and model._process is not None:
            if not model._process.is_running():
                logger.warning(f"Backend process for {request.model_id} died, reloading...")
                # Remove from cache and reload
                del _models[request.model_id]
                del _model_configs[request.model_id]
            else:
                return LoadResponse(
                    success=True,
                    model_id=request.model_id,
                    message="Model already loaded"
                )
        else:
            return LoadResponse(
                success=True,
                model_id=request.model_id,
                message="Model already loaded"
            )
    
    try:
        from oprel.core.model import Model
        
        logger.info(f"Loading model: {request.model_id} (quant={request.quantization}, backend={request.backend})")
        
        # Create model with use_server=False (direct mode inside server)
        model = Model(
            model_id=request.model_id,
            quantization=request.quantization,
            max_memory_mb=request.max_memory_mb,
            backend=request.backend,
            use_server=False,
        )
        
        # Load the model
        model.load()
        
        # Cache it
        _models[request.model_id] = model
        _model_configs[request.model_id] = {
            "quantization": request.quantization,
            "max_memory_mb": request.max_memory_mb,
            "backend": request.backend,
        }
        
        # Mark as just loaded (used)
        _mark_model_used(request.model_id)
        
        logger.info(f"Model loaded successfully: {request.model_id}")
        
        return LoadResponse(
            success=True,
            model_id=request.model_id,
            message="Model loaded successfully"
        )
    
    except Exception as e:
        logger.error(f"Failed to load model {request.model_id}: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"Failed to load model: {str(e)}")


@app.post("/generate")
async def generate_text(request: GenerateRequest):
    """Generate text (Conversational)"""
    
    # Auto-load logic
    if request.model_id not in _models:
        load_req = LoadRequest(model_id=request.model_id)
        await load_model(load_req)
    
    model = _models[request.model_id]
    
    # Check if backend process is still alive
    if hasattr(model, '_process') and model._process is not None:
        if not model._process.is_running():
            logger.warning(f"Backend process for {request.model_id} died, reloading...")
            # Remove from cache and reload
            del _models[request.model_id]
            del _model_configs[request.model_id]
            load_req = LoadRequest(model_id=request.model_id)
            await load_model(load_req)
            model = _models[request.model_id]
    
    model._loaded = True # Fix reload bug
    
    if not hasattr(model, '_client') or model._client is None:
        raise HTTPException(status_code=500, detail="Model client not available")

    # Mark model as used (for idle timeout tracking)
    _mark_model_used(request.model_id)

    # --- Conversation Management ---
    conv_id = request.conversation_id
    if not conv_id:
        conv_id = str(uuid.uuid4())
        
    if request.reset_conversation or conv_id not in _conversations:
        _conversations[conv_id] = []
        _conversation_meta[conv_id] = {
            "created_at": str(datetime.now()),
            "model_id": request.model_id,
        }
        
    history = _conversations[conv_id]
    _cleanup_conversations() # Prune if needed
    
    # Update metadata
    _conversation_meta[conv_id]["last_updated"] = str(datetime.now())
    
    # Build prompt with history
    # If conversation_id was NOT passed explicitly (one-off), we might still want to use template
    # but the request.prompt is the "new user message"
    
    # If specific system prompt requested, use it, otherwise use stored or None
    sys_prompt = request.system_prompt
    
    full_prompt = _build_chat_prompt(
        request.model_id, 
        history, 
        sys_prompt, 
        request.prompt
    )
    
    # If raw prompt mode requested? (Future feature). For now, always use Chat template if model detection works.
    
    try:
        final_text = ""
        
        if request.stream:
            def generate_stream():
                full_resp = ""
                try:
                    for token in model._client.generate(
                        prompt=full_prompt,
                        max_tokens=request.max_tokens,
                        temperature=request.temperature,
                        stream=True,
                    ):
                        full_resp += token
                        yield f"data: {token}\n\n"
                    
                    # Update history after full generation
                    # We can't update per-token in the generator safely due to async/yield
                    # But Python generators run in the worker... 
                    # We'll update history when stream is DONE
                    if len(history) >= MAX_HISTORY_MSGS:
                        history.pop(0) # Remove oldest pair? Ideally remove 0 and 1.
                        if len(history) > 0: history.pop(0)

                    history.append({"role": "user", "content": request.prompt})
                    history.append({"role": "assistant", "content": full_resp})
                    
                    yield "data: [DONE]\n\n"
                except Exception as e:
                    yield f"data: [ERROR] {str(e)}\n\n"
            
            # Note: The conversation update happens inside the generator which might be tricky if it fails mid-stream.
            # But for a basic implementation this works.
            
            return StreamingResponse(
                generate_stream(),
                media_type="text/event-stream",
                headers={
                    "Cache-Control": "no-cache",
                    "Connection": "keep-alive",
                    "X-Conversation-ID": conv_id,
                }
            )
        else:
            # Non-streaming
            text = model._client.generate(
                prompt=full_prompt,
                max_tokens=request.max_tokens,
                temperature=request.temperature,
                stream=False,
            )
            
            # Update history
            if len(history) >= MAX_HISTORY_MSGS:
                history.pop(0)
                if len(history) > 0: history.pop(0)
                
            history.append({"role": "user", "content": request.prompt})
            history.append({"role": "assistant", "content": text})
            
            return GenerateResponse(
                text=text,
                model_id=request.model_id,
                conversation_id=conv_id,
                message_count=len(history)
            )
            
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Generation failed: {str(e)}")


# --- Conversation Endpoints ---

@app.get("/conversations", response_model=List[ConversationInfo])
async def list_conversations():
    """List active conversations"""
    results = []
    for cid, meta in _conversation_meta.items():
        results.append(ConversationInfo(
            id=cid,
            created_at=meta.get("created_at", ""),
            last_updated=meta.get("last_updated", ""),
            message_count=len(_conversations.get(cid, [])),
            model_id=meta.get("model_id", "unknown")
        ))
    return results


@app.get("/conversations/{conversation_id}")
async def get_conversation(conversation_id: str):
    """Get conversation history"""
    if conversation_id not in _conversations:
        raise HTTPException(status_code=404, detail="Conversation not found")
    return _conversations[conversation_id]


@app.delete("/conversations/{conversation_id}")
async def delete_conversation(conversation_id: str):
    """Delete a conversation"""
    if conversation_id in _conversations:
        del _conversations[conversation_id]
    if conversation_id in _conversation_meta:
        del _conversation_meta[conversation_id]
    return {"success": True}


@app.post("/conversations/{conversation_id}/reset")
async def reset_conversation(conversation_id: str):
    """Reset a conversation history"""
    if conversation_id in _conversations:
        _conversations[conversation_id] = []
        return {"success": True}
    raise HTTPException(status_code=404, detail="Conversation not found")


@app.post("/unload", response_model=UnloadResponse)
async def unload_model_post(request: UnloadRequest):
    return await unload_model(request.model_id) # Reuse logic


@app.delete("/unload/{model_id}", response_model=UnloadResponse)
async def unload_model(model_id: str):
    global _models, _model_configs
    if model_id not in _models:
        raise HTTPException(status_code=404, detail=f"Model '{model_id}' not loaded")
    
    try:
        model = _models[model_id]
        model.unload()
        del _models[model_id]
        if model_id in _model_configs:
            del _model_configs[model_id]
        return UnloadResponse(success=True, model_id=model_id, message="Unloaded")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))



# ============================================================================
# OpenAI & Ollama API Compatibility (Week 14)
# ============================================================================

class OpenAIChatMessage(BaseModel):
    role: str
    content: str


class OpenAIChatRequest(BaseModel):
    model: str
    messages: List[OpenAIChatMessage]
    temperature: float = 0.7
    max_tokens: Optional[int] = 512
    stream: bool = False


class OpenAICompletionRequest(BaseModel):
    model: str
    prompt: str
    temperature: float = 0.7
    max_tokens: Optional[int] = 512
    stream: bool = False


@app.post("/v1/chat/completions")
async def v1_chat_completions(request: OpenAIChatRequest):
    """OpenAI-compatible chat completions endpoint"""
    # Extract prompt from messages
    prompt = request.messages[-1].content if request.messages else ""
    
    # Build conversation history
    conversation_history = [{"role": msg.role, "content": msg.content} for msg in request.messages[:-1]]
    
    # Create GenerateRequest
    gen_request = GenerateRequest(
        model_id=request.model,
        prompt=prompt,
        max_tokens=request.max_tokens or 512,
        temperature=request.temperature,
        stream=request.stream
    )
    
    # Set up conversation
    conv_id = str(uuid.uuid4())
    _conversations[conv_id] = conversation_history
    _conversation_meta[conv_id] = {
        "created_at": str(datetime.now()),
        "model_id": request.model,
        "last_updated": str(datetime.now())
    }
    gen_request.conversation_id = conv_id
    
    response = await generate_text(gen_request)
    
    if request.stream:
        # Convert daemon SSE format to OpenAI format
        async def openai_stream_wrapper():
            request_id = f"chatcmpl-{int(time_module.time() * 1000)}"
            buffer = ""
            async for chunk in response.body_iterator:
                chunk_str = chunk.decode('utf-8') if isinstance(chunk, bytes) else chunk
                buffer += chunk_str
                while "\n\n" in buffer:
                    line, buffer = buffer.split("\n\n", 1)
                    if line.startswith("data: "):
                        token = line[6:]
                        if token and token != "[DONE]" and not token.startswith("[ERROR]"):
                            chunk = {
                                "id": request_id,
                                "object": "chat.completion.chunk",
                                "created": int(time_module.time()),
                                "model": request.model,
                                "choices": [{
                                    "index": 0,
                                    "delta": {"content": token},
                                    "finish_reason": None
                                }]
                            }
                            yield f"data: {json.dumps(chunk)}\n\n"
            
            # Send final chunk
            final_chunk = {
                "id": request_id,
                "object": "chat.completion.chunk",
                "created": int(time_module.time()),
                "model": request.model,
                "choices": [{
                    "index": 0,
                    "delta": {},
                    "finish_reason": "stop"
                }]
            }
            yield f"data: {json.dumps(final_chunk)}\n\n"
            yield "data: [DONE]\n\n"
        
        return StreamingResponse(
            openai_stream_wrapper(),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "X-Conversation-ID": conv_id,
            }
        )
    else:
        return {
            "id": f"chatcmpl-{int(time_module.time() * 1000)}",
            "object": "chat.completion",
            "created": int(time_module.time()),
            "model": request.model,
            "choices": [{
                "index": 0,
                "message": {"role": "assistant", "content": response.text},
                "finish_reason": "stop"
            }],
            "usage": {
                "prompt_tokens": len(prompt.split()),
                "completion_tokens": len(response.text.split()),
                "total_tokens": len(prompt.split()) + len(response.text.split())
            }
        }


@app.post("/v1/completions")
async def v1_completions(request: OpenAICompletionRequest):
    """OpenAI-compatible text completions endpoint"""
    gen_request = GenerateRequest(
        model_id=request.model,
        prompt=request.prompt,
        max_tokens=request.max_tokens or 512,
        temperature=request.temperature,
        stream=request.stream
    )
    
    response = await generate_text(gen_request)
    
    if request.stream:
        # Convert daemon SSE format to OpenAI format
        async def openai_stream_wrapper():
            request_id = f"cmpl-{int(time_module.time() * 1000)}"
            buffer = ""
            async for chunk in response.body_iterator:
                chunk_str = chunk.decode('utf-8') if isinstance(chunk, bytes) else chunk
                buffer += chunk_str
                while "\n\n" in buffer:
                    line, buffer = buffer.split("\n\n", 1)
                    if line.startswith("data: "):
                        token = line[6:]
                        if token and token != "[DONE]" and not token.startswith("[ERROR]"):
                            chunk = {
                                "id": request_id,
                                "object": "text_completion",
                                "created": int(time_module.time()),
                                "model": request.model,
                                "choices": [{
                                    "text": token,
                                    "index": 0,
                                    "finish_reason": None
                                }]
                            }
                            yield f"data: {json.dumps(chunk)}\n\n"
            
            # Send final chunk
            final_chunk = {
                "id": request_id,
                "object": "text_completion",
                "created": int(time_module.time()),
                "model": request.model,
                "choices": [{
                    "text": "",
                    "index": 0,
                    "finish_reason": "stop"
                }]
            }
            yield f"data: {json.dumps(final_chunk)}\n\n"
            yield "data: [DONE]\n\n"
        
        return StreamingResponse(
            openai_stream_wrapper(),
            media_type="text/event-stream"
        )
    else:
        return {
            "id": f"cmpl-{int(time_module.time() * 1000)}",
            "object": "text_completion",
            "created": int(time_module.time()),
            "model": request.model,
            "choices": [{
                "text": response.text,
                "index": 0,
                "finish_reason": "stop"
            }],
            "usage": {
                "prompt_tokens": len(request.prompt.split()),
                "completion_tokens": len(response.text.split()),
                "total_tokens": len(request.prompt.split()) + len(response.text.split())
            }
        }


@app.get("/v1/models")
async def v1_list_models():
    """OpenAI-compatible model list endpoint"""
    from oprel.downloader.cache import list_cached_models
    from oprel.downloader.aliases import MODEL_ALIASES
    
    models = []
    
    try:
        cached = list_cached_models()
        for model_info in cached:
            model_name = model_info.get('name', '')
            if model_name:
                models.append({
                    "id": model_name,
                    "object": "model",
                    "created": int(model_info.get('modified', datetime.now()).timestamp()),
                    "owned_by": "oprel"
                })
    except Exception as e:
        logger.warning(f"Could not list cached models: {e}")
    
    try:
        for alias in MODEL_ALIASES.keys():
            if not any(m["id"] == alias for m in models):
                models.append({
                    "id": alias,
                    "object": "model",
                    "created": int(time_module.time()),
                    "owned_by": "oprel"
                })
    except Exception as e:
        logger.warning(f"Could not list aliases: {e}")
    
    return {"object": "list", "data": models}


@app.get("/v1/health")
async def v1_health():
    """OpenAI-style health check"""
    return await health_check()


@app.post("/api/chat")
async def api_chat(request: dict):
    """Ollama-compatible chat endpoint"""
    openai_req = OpenAIChatRequest(
        model=request.get("model", ""),
        messages=[OpenAIChatMessage(**msg) for msg in request.get("messages", [])],
        stream=request.get("stream", False),
        temperature=request.get("options", {}).get("temperature", 0.7),
        max_tokens=request.get("options", {}).get("num_predict", 512)
    )
    return await v1_chat_completions(openai_req)


@app.post("/api/generate")
async def api_generate(request: dict):
    """Ollama-compatible generate endpoint"""
    openai_req = OpenAICompletionRequest(
        model=request.get("model", ""),
        prompt=request.get("prompt", ""),
        stream=request.get("stream", False),
        temperature=request.get("options", {}).get("temperature", 0.7),
        max_tokens=request.get("options", {}).get("num_predict", 512)
    )
    return await v1_completions(openai_req)


@app.get("/api/tags")
async def api_tags():
    """Ollama-compatible model list endpoint"""
    models_response = await v1_list_models()
    return {
        "models": [
            {
                "name": model["id"],
                "modified_at": time_module.strftime("%Y-%m-%dT%H:%M:%SZ", time_module.gmtime(model["created"])),
                "size": 0,
                "digest": ""
            }
            for model in models_response["data"]
        ]
    }


@app.post("/shutdown")
async def shutdown_server():
    """Gracefully shutdown the server"""
    import asyncio
    
    async def shutdown():
        # Give response time to be sent
        await asyncio.sleep(0.5)
        # Cleanup models
        _cleanup_models()
        # Exit the process
        import os
        os._exit(0)
    
    # Start shutdown in background
    asyncio.create_task(shutdown())
    return {"status": "shutting down"}


def run_server(host: str = "127.0.0.1", port: int = 11434):
    """Run the daemon server"""
    import uvicorn
    print(f"{Colors.GREEN}{Colors.BOLD}Oprel Daemon v0.3.0{Colors.RESET}")
    print(f"  Listening on: {Colors.CYAN}http://{host}:{port}{Colors.RESET}")
    print(f"  Press {Colors.YELLOW}Ctrl+C{Colors.RESET} to stop\n")
    uvicorn.run(app, host=host, port=port, log_level="warning", access_log=False)


if __name__ == "__main__":
    run_server()
