import os
import logging
from typing import Optional
from datetime import datetime
import time

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# Logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(title="Gemma 4 Inference API")

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

# Config
MODEL_NAME = os.getenv("MODEL_NAME", "google/gemma-4-E4B-it")

# Global model/tokenizer
model = None
tokenizer = None

# Pydantic models
class Message(BaseModel):
    role: str
    content: str

class ChatRequest(BaseModel):
    messages: list[Message]
    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
    max_tokens: int = Field(default=512, ge=1, le=2048)
    top_p: float = Field(default=0.9, ge=0.0, le=1.0)

class ChatChoice(BaseModel):
    index: int
    message: Message
    finish_reason: str

class ChatUsage(BaseModel):
    completion_tokens: int
    total_tokens: int

class ChatResponse(BaseModel):
    model: str
    object: str = "chat.completion"
    created: int
    choices: list[ChatChoice]
    usage: ChatUsage

def load_model():
    """Load model and tokenizer on startup."""
    global model, tokenizer
    
    logger.info(f"Loading {MODEL_NAME}...")
    
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    
    # Load with 4-bit quantization to fit in 16GB
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        device_map="auto",
        torch_dtype=torch.bfloat16,
        load_in_4bit=True,
        low_cpu_mem_usage=True,
    )
    
    logger.info(f"✓ {MODEL_NAME} loaded successfully")

@app.on_event("startup")
async def startup():
    load_model()

@app.get("/health")
def health():
    return {"status": "ok", "model": MODEL_NAME}

@app.get("/v1/models")
def list_models():
    return {
        "object": "list",
        "data": [
            {
                "id": "gemma-4",
                "object": "model",
                "owned_by": "google",
                "created": int(time.time()),
            }
        ]
    }

@app.post("/v1/chat/completions", response_model=ChatResponse)
def chat_completions(request: ChatRequest):
    """OpenAI-compatible chat completions endpoint."""
    
    try:
        # Build prompt from messages
        prompt = ""
        for msg in request.messages:
            if msg.role == "system":
                prompt += f"<|system|>\n{msg.content}<|end_of_turn|>\n"
            elif msg.role == "user":
                prompt += f"<|user|>\n{msg.content}<|end_of_turn|>\n"
            elif msg.role == "assistant":
                prompt += f"<|assistant|>\n{msg.content}<|end_of_turn|>\n"
        
        prompt += "<|assistant|>\n"
        
        # Tokenize
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        input_length = inputs.input_ids.shape[1]
        
        # Generate
        start_time = time.time()
        outputs = model.generate(
            **inputs,
            max_new_tokens=request.max_tokens,
            temperature=request.temperature,
            top_p=request.top_p,
            do_sample=request.temperature > 0,
            pad_token_id=tokenizer.eos_token_id,
        )
        
        # Decode
        full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Extract just the response
        if "<|assistant|>" in full_text:
            response_text = full_text.split("<|assistant|>")[-1].strip()
        else:
            response_text = full_text
        
        tokens_generated = outputs.shape[1] - input_length
        
        return ChatResponse(
            model="gemma-4",
            created=int(time.time()),
            choices=[
                ChatChoice(
                    index=0,
                    message=Message(role="assistant", content=response_text),
                    finish_reason="stop",
                )
            ],
            usage=ChatUsage(
                completion_tokens=tokens_generated,
                total_tokens=tokens_generated,
            ),
        )
    
    except Exception as e:
        logger.error(f"Error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/chat/completions", response_model=ChatResponse)
def chat_completions_no_v1(request: ChatRequest):
    """Alias without /v1/ prefix."""
    return chat_completions(request)

@app.get("/")
def root():
    return {
        "name": "Gemma 4 API",
        "model": MODEL_NAME,
        "docs": "Use /v1/chat/completions for OpenAI compatibility",
        "example": {
            "url": "/v1/chat/completions",
            "method": "POST",
            "body": {
                "messages": [{"role": "user", "content": "Hello"}],
                "temperature": 0.7,
                "max_tokens": 512,
            }
        }
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)