#!/usr/bin/env python3
"""
FastAPI server for serving Mistral 7B fine-tuned models
"""

import os
import sys
from typing import Optional, Dict, Any
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import uvicorn
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
from inference.inference_mistral7b import load_local_model, generate_with_local_model, get_device_info
import torch

# Configuration - Resolve model path relative to msp root
_MODEL_BASE = Path(__file__).parent.parent / "mistral7b-finetuned-ahb2apb"
DEFAULT_MODEL_PATH = str(_MODEL_BASE)

# Global model and tokenizer (loaded once at startup)
model = None
tokenizer = None
device_info = None

app = FastAPI(
    title="Mistral 7B AHB2APB API",
    description="API for serving the fine-tuned Mistral 7B model for AHB2APB conversion",
    version="1.0.0"
)

# Enable CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Request/Response models
class GenerateRequest(BaseModel):
    prompt: str
    max_length: Optional[int] = 512
    temperature: Optional[float] = 0.7

class GenerateResponse(BaseModel):
    response: str
    model: str
    max_length: int
    temperature: float

class HealthResponse(BaseModel):
    status: str
    model_loaded: bool
    device: str
    model_path: str

@app.on_event("startup")
async def load_model():
    """Load the model when the server starts"""
    global model, tokenizer, device_info
    
    model_path = os.environ.get("MODEL_PATH", DEFAULT_MODEL_PATH)
    
    print(f"Loading model from: {model_path}")
    print("=" * 70)
    
    try:
        device_info = get_device_info()
        model, tokenizer = load_local_model(model_path)
        print(f"\n✓ Model loaded successfully on {device_info['device']}!")
        print(f"✓ Server ready to accept requests")
        print("=" * 70)
    except Exception as e:
        print(f"\n✗ Error loading model: {e}")
        print("=" * 70)
        sys.exit(1)

@app.get("/health", response_model=HealthResponse)
async def health_check():
    """Health check endpoint"""
    return HealthResponse(
        status="healthy" if model is not None else "error",
        model_loaded=model is not None,
        device=device_info["device"] if device_info else "unknown",
        model_path=os.environ.get("MODEL_PATH", DEFAULT_MODEL_PATH)
    )

@app.get("/")
async def root():
    """Root endpoint with API information"""
    return {
        "name": "Mistral 7B AHB2APB API",
        "version": "1.0.0",
        "status": "running",
        "model": os.environ.get("MODEL_PATH", DEFAULT_MODEL_PATH),
        "endpoints": {
            "health": "/health",
            "generate": "/api/generate",
            "docs": "/docs"
        }
    }

@app.post("/api/generate", response_model=GenerateResponse)
async def generate(request: GenerateRequest):
    """
    Generate text from a prompt using the fine-tuned model
    """
    if model is None or tokenizer is None:
        raise HTTPException(status_code=503, detail="Model not loaded")
    
    try:
        response = generate_with_local_model(
            model=model,
            tokenizer=tokenizer,
            prompt=request.prompt,
            max_length=request.max_length or 512,
            temperature=request.temperature or 0.7
        )
        
        return GenerateResponse(
            response=response,
            model=os.environ.get("MODEL_PATH", DEFAULT_MODEL_PATH),
            max_length=request.max_length or 512,
            temperature=request.temperature or 0.7
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Generation error: {str(e)}")

@app.post("/api/generate/batch")
async def generate_batch(requests: list[GenerateRequest]):
    """
    Generate text from multiple prompts (batch processing)
    """
    if model is None or tokenizer is None:
        raise HTTPException(status_code=503, detail="Model not loaded")
    
    try:
        responses = []
        for req in requests:
            response = generate_with_local_model(
                model=model,
                tokenizer=tokenizer,
                prompt=req.prompt,
                max_length=req.max_length or 512,
                temperature=req.temperature or 0.7
            )
            responses.append({
                "response": response,
                "prompt": req.prompt
            })
        
        return {"results": responses}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Batch generation error: {str(e)}")

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="Start Mistral 7B API server")
    parser.add_argument(
        "--model-path",
        type=str,
        default=DEFAULT_MODEL_PATH,
        help=f"Path to fine-tuned model (default: {DEFAULT_MODEL_PATH})"
    )
    parser.add_argument(
        "--host",
        type=str,
        default="0.0.0.0",
        help="Host to bind to (default: 0.0.0.0)"
    )
    parser.add_argument(
        "--port",
        type=int,
        default=8000,
        help="Port to bind to (default: 8000)"
    )
    parser.add_argument(
        "--reload",
        action="store_true",
        help="Enable auto-reload (for development)"
    )
    parser.add_argument(
        "--workers",
        type=int,
        default=1,
        help="Number of worker processes (default: 1)"
    )
    
    args = parser.parse_args()
    
    # Set model path as environment variable for the startup event
    os.environ["MODEL_PATH"] = args.model_path
    
    print(f"\n🚀 Starting Mistral 7B AHB2APB API Server")
    print(f"   Model: {args.model_path}")
    print(f"   Host: {args.host}")
    print(f"   Port: {args.port}")
    print(f"   Workers: {args.workers}")
    print(f"   Reload: {args.reload}\n")
    
    # Change to api directory for proper module resolution
    import os
    os.chdir(os.path.dirname(os.path.abspath(__file__)))
    uvicorn.run(
        "api_server:app",
        host=args.host,
        port=args.port,
        reload=args.reload,
        workers=1 if args.reload else args.workers
    )