import os
import uvicorn
from fastapi import FastAPI, UploadFile, File, Header, HTTPException, BackgroundTasks
from fastapi.responses import StreamingResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from typing import List, Optional, Any, Dict, Union
import json
import asyncio
import logging

from xorfice.engine import XoronEngine
from xorfice.memory import FlatFileMemoryManager
from xorfice.logger import JSONLLogger
from fastapi import BackgroundTasks, UploadFile, File

app = FastAPI(title="Xoron-Dev Elite SOTA Inference Server")

# Output directory for media
MEDIA_DIR = "./outputs"
os.makedirs(MEDIA_DIR, exist_ok=True)
app.mount("/outputs", StaticFiles(directory=MEDIA_DIR), name="outputs")

# Global instances
engine = None
logger = JSONLLogger(log_dir="./logs")

@app.on_event("startup")
async def startup_event():
    global engine
    model_path = os.environ.get("XORON_MODEL_PATH", "./weights")
    vram_capacity = int(os.environ.get("XORON_VRAM_CAPACITY", 4))
    kv_blocks = int(os.environ.get("XORON_KV_BLOCKS", 256))
    expert_offload_threshold = int(os.environ.get("XORON_EXPERT_OFFLOAD_THRESHOLD", 0))
    lazy_load = os.environ.get("XORON_LAZY_LOAD", "false").lower() == "true"
    
    if lazy_load:
        print(f"🚀 Starting FastAPI Web UI Server. Model loading deferred (Lazy Load mode)...")
        engine = None
    else:
        print(f"🚀 Starting Elite SOTA XoronEngine with Agentic MCP Integration...")
        engine = XoronEngine(
            model_path=model_path, 
            max_vram_experts=vram_capacity,
            kv_cache_blocks=kv_blocks,
            expert_offload_threshold=expert_offload_threshold
        )

class MCPConnectRequest(BaseModel):
    server_id: str
    command: List[str]

@app.post("/v1/mcp/connect")
async def connect_mcp_server(request: MCPConnectRequest):
    """
    Dynamically connect a standard Model Context Protocol (MCP) server.
    This enables the SOTA S2S model to natively execute tools from external clients.
    """
    if not engine: raise HTTPException(status_code=503, detail="Engine not ready")
    
    try:
        tools_found = engine.agent.connect_mcp(request.server_id, request.command)
        return {"status": "success", "server_id": request.server_id, "tools_discovered": tools_found}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to connect MCP server: {str(e)}")

class ChatMessage(BaseModel):
    role: str
    content: str
    image_url: Optional[str] = None
    video_url: Optional[str] = None
    audio_url: Optional[str] = None

class ChatCompletionRequest(BaseModel):
    model: str
    messages: List[ChatMessage]
    stream: bool = False
    max_tokens: int = 1024
    temperature: float = 0.7
    voice_id: Optional[str] = None

@app.get("/v1/models")
async def list_models():
    if not engine: return {"object": "list", "data": []}
    return {"object": "list", "data": [{"id": "xoron-dev", "object": "model"}]}

@app.get("/v1/audio/voices")
async def list_voices():
    if not engine: raise HTTPException(status_code=503, detail="Engine not ready")
    voices = engine.voice.list_available_voices()
    return {"voices": voices}

@app.post("/v1/audio/voice-clone")
async def voice_clone(file: UploadFile = File(...), x_user_id: str = Header(None)):
    if not x_user_id: raise HTTPException(status_code=400, detail="Missing x-user-id header")
    if not engine: raise HTTPException(status_code=503, detail="Engine not ready")
    
    # Save temp audio file
    temp_path = f"outputs/temp_clone_{x_user_id}.wav"
    with open(temp_path, "wb") as buffer:
        buffer.write(await file.read())
        
    # SOTA: extract and save embedding
    import torch
    import torchaudio
    waveform, _ = torchaudio.load(temp_path)
    embedding = engine.voice.extract_speaker_embedding(engine.model, waveform.to(engine.device))
    engine.voice.set_user_voice(x_user_id, embedding)
    
    return {"status": "success", "voice_id": x_user_id}

@app.post("/v1/chat/completions")
async def chat_completions(request: ChatCompletionRequest, background_tasks: BackgroundTasks, x_user_id: str = Header(None)):
    if not x_user_id: raise HTTPException(status_code=400, detail="Missing x-user-id header")
    if not engine: raise HTTPException(status_code=503, detail="Engine offline (Started in Lazy Load / UI-Only Mode)")

    last_msg = request.messages[-1]
    full_prompt = last_msg.content
    
    # 1. Omnimodal Direct Routing: If an audio url is provided, seamlessly route to Speech-to-Speech
    if last_msg.audio_url:
        import torch
        import torchaudio
        import tempfile
        import requests
        import io
        try:
            is_url = last_msg.audio_url.startswith("http")
            if is_url:
                response = requests.get(last_msg.audio_url, timeout=10)
                response.raise_for_status()
                with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
                    tmp.write(response.content)
                    tmp_path = tmp.name
                waveform, _ = torchaudio.load(tmp_path)
                os.remove(tmp_path)
            else:
                waveform, _ = torchaudio.load(last_msg.audio_url)
            
            # Use the actual loaded waveform
            waveform = waveform[0].unsqueeze(0).to(engine.device) # Extract single channel (B, T)
            
        except Exception as e:
            # Revert to silent fallback only if the file literally doesn't load/doesn't exist
            print(f"[API] Error loading audio: {e}")
            waveform = torch.zeros((1, 16000), device=engine.device)
            
        result = engine.speech_to_speech(waveform, full_prompt, user_id=x_user_id)
        return {"id": "chatcmpl-s2s", "choices": [{"message": {"role": "assistant", "content": result.get("text"), "audio": result.get("audio_url")}, "index": 0}]}

    # 2. Heuristic T2T / Multimodal: If image or video keywords appear and no audio URL, format for generation
    # Note: Xoron is natively decoding <|generate_x|> in the architecture.
    if "generate" in full_prompt.lower() and "image" in full_prompt.lower(): 
        full_prompt += " <|generate_image|>"
    elif "generate" in full_prompt.lower() and "video" in full_prompt.lower(): 
        full_prompt += " <|generate_video|>"
    
    # 4. Agentic / Text Execution
    if not request.stream:
        # Use simple direct blocking call if not streaming
        result = engine.generate(
            full_prompt, 
            user_id=x_user_id, 
            stream=False, 
            max_new_tokens=request.max_tokens, 
            temperature=request.temperature
        )
        
        background_tasks.add_task(engine.train_user_on_interaction, x_user_id, last_msg.content, result.get("text", ""))
        logger.log_interaction(x_user_id, "/v1/chat/completions", {"prompt": last_msg.content}, {"response": result})
        
        return {"id": "chatcmpl-xoron", "choices": [{"message": {"role": "assistant", "content": result.get("text"), "media": result}}]}

    # 5. Streaming Path
    async def stream_generator():
        full_response = ""
        # engine.generate returns a generator if stream=True
        generator = engine.generate(
            full_prompt, 
            user_id=x_user_id, 
            stream=True, 
            max_new_tokens=request.max_tokens, 
            temperature=request.temperature
        )
        for chunk in generator:
            if isinstance(chunk, dict) and "usage" in chunk:
                # Transmit final usage chunk natively as an OpenAI spec stream message
                yield f"data: {json.dumps({'choices': [], 'usage': chunk['usage']})}\n\n"
            else:
                full_response += chunk
                yield f"data: {json.dumps({'choices': [{'delta': {'content': chunk}}]})}\n\n"
                await asyncio.sleep(0.01)
        
        background_tasks.add_task(engine.train_user_on_interaction, x_user_id, last_msg.content, full_response)
        logger.log_interaction(x_user_id, "/v1/chat/completions", {"prompt": last_msg.content}, {"response": full_response})
        yield "data: [DONE]\n\n"
        
    return StreamingResponse(stream_generator(), media_type="text/event-stream")

@app.get("/health")
async def health_check():
    return {"status": "ok", "engine": "running" if engine else "initializing"}

@app.get("/v1/metrics")
async def get_metrics():
    if not engine: raise HTTPException(status_code=503, detail="Engine not ready")
    import torch
    if torch.cuda.is_available():
        allocated = torch.cuda.memory_allocated() / (1024 ** 3)
        reserved = torch.cuda.memory_reserved() / (1024 ** 3)
    else:
        allocated = 0
        reserved = 0
    return {
        "status": "online",
        "device": str(engine.device),
        "vram_allocated_gb": round(allocated, 2),
        "vram_reserved_gb": round(reserved, 2)
    }

@app.get("/v1/mcp/tools")
async def get_tools():
    if not engine or not engine.agent: return {"tools": []}
    return {"tools": engine.agent.get_all_tool_definitions()}

class ToolToggleRequest(BaseModel):
    tool_name: str
    enabled: bool

@app.post("/v1/mcp/tools/toggle")
async def toggle_tool(req: ToolToggleRequest):
    if not engine or not engine.agent: raise HTTPException(status_code=503, detail="Agent offline")
    # Basic toggle logic mapped to agent memory config
    engine.agent.options = getattr(engine.agent, 'options', {})
    engine.agent.options[f"disable_{req.tool_name}"] = not req.enabled
    return {"status": "success", "tool": req.tool_name, "enabled": req.enabled}

# --- Serve Vite React UI ---
ui_dist_path = os.path.join(os.path.dirname(__file__), "webui")
if os.path.exists(ui_dist_path):
    app.mount("/assets", StaticFiles(directory=os.path.join(ui_dist_path, "assets")), name="ui-assets")
    
    @app.get("/{full_path:path}")
    async def serve_ui(full_path: str):
        # Allow API routes to pass through natively
        if full_path.startswith("v1/") or full_path.startswith("health") or full_path.startswith("docs") or full_path.startswith("openapi"):
            raise HTTPException(status_code=404, detail="Not Found")
            
        file_path = os.path.join(ui_dist_path, full_path)
        if os.path.isfile(file_path):
            return FileResponse(file_path)
            
        # SPA fallback to index.html
        return FileResponse(os.path.join(ui_dist_path, "index.html"))

def start_server(host: str = "0.0.0.0", port: int = 8000):
    uvicorn.run(app, host=host, port=port)
