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

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

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

MODEL_NAME = os.getenv("MODEL_NAME", "gemma-4-E4B-it")
MODEL_SIZE = int(os.getenv("MODEL_SIZE", "5"))
MODEL_CONTEXT = int(os.getenv("MODEL_CONTEXT", "128000"))
HF_REPO_ID = os.getenv("HF_REPO_ID", f"google/{MODEL_NAME}")
QUANTIZATION = os.getenv("QUANTIZATION", "Q4_K_M")

DEVICE = "cpu"
if torch.cuda.is_available():
    DEVICE = "cuda"
elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
    DEVICE = "mps"

app = FastAPI(title=f"Gemma Inference API - {MODEL_NAME}")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

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

class InferenceRequest(BaseModel):
    messages: list[ChatMessage]
    model: str = MODEL_NAME
    temperature: float = 0.7
    max_tokens: int = 512
    top_p: float = 0.9
    top_k: int = 50
    thinking: bool = False
    system_prompt: Optional[str] = None

class InferenceResponse(BaseModel):
    model: str
    response: str
    tokens_used: int
    latency_ms: float
    thinking: Optional[str] = None
    timestamp: str

class HealthResponse(BaseModel):
    status: str
    model: str
    device: str
    model_size_gb: int
    context_window: int

@lru_cache(maxsize=1)
def get_tokenizer():
    logger.info(f"Loading tokenizer for {HF_REPO_ID}...")
    return AutoTokenizer.from_pretrained(HF_REPO_ID)

@lru_cache(maxsize=1)
def get_model():
    logger.info(f"Loading model {HF_REPO_ID} on {DEVICE}...")
    kwargs = dict(low_cpu_mem_usage=True)
    if DEVICE == "cuda":
        kwargs["torch_dtype"] = torch.float16
        model = AutoModelForCausalLM.from_pretrained(HF_REPO_ID, device_map="auto", **kwargs)
    elif DEVICE == "cpu" and MODEL_SIZE > 10:
        kwargs["torch_dtype"] = torch.float32
        try:
            model = AutoModelForCausalLM.from_pretrained(HF_REPO_ID, device_map="auto", load_in_4bit=True, **kwargs)
        except Exception:
            model = AutoModelForCausalLM.from_pretrained(HF_REPO_ID, **kwargs)
    else:
        kwargs["torch_dtype"] = torch.float32
        model = AutoModelForCausalLM.from_pretrained(HF_REPO_ID, **kwargs)
    return model

def build_chat_prompt(messages: list[ChatMessage], system_prompt: Optional[str] = None) -> str:
    parts = []
    if system_prompt:
        parts.append(f"<|system|>\n{system_prompt}<|end_of_turn|>\n")
    for msg in messages:
        parts.append(f"<|{msg.role}|>\n{msg.content}<|end_of_turn|>\n")
    parts.append("<|assistant|>\n")
    return "".join(parts)

@app.get("/health", response_model=HealthResponse)
async def health_check():
    return HealthResponse(status="healthy", model=MODEL_NAME, device=DEVICE, model_size_gb=MODEL_SIZE, context_window=MODEL_CONTEXT)

@app.get("/info")
async def model_info():
    return {"model": MODEL_NAME, "repo_id": HF_REPO_ID, "device": DEVICE, "model_size_gb": MODEL_SIZE, "context_window": MODEL_CONTEXT, "quantization": QUANTIZATION}

@app.post("/infer", response_model=InferenceResponse)
async def infer(request: InferenceRequest):
    start_time = time.time()
    try:
        tokenizer = get_tokenizer()
        model = get_model()
        prompt = build_chat_prompt(request.messages, system_prompt=request.system_prompt)
        if request.thinking:
            prompt = "<|think|>\n" + prompt
        inputs = tokenizer(prompt, return_tensors="pt")
        model_device = getattr(model, 'device', None)
        if model_device is not None:
            inputs = {k: v.to(model_device) for k, v in inputs.items()}
        input_length = inputs['input_ids'].shape[1]
        max_new = max(1, min(request.max_tokens, MODEL_CONTEXT - input_length))
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new,
            temperature=request.temperature,
            top_p=request.top_p,
            top_k=request.top_k,
            do_sample=request.temperature > 0,
            pad_token_id=tokenizer.eos_token_id,
        )
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        response_text = generated_text.split("<|assistant|>")[-1].strip() if "<|assistant|>" in generated_text else generated_text
        latency_ms = (time.time() - start_time) * 1000
        tokens_generated = outputs.shape[1] - input_length
        return InferenceResponse(model=MODEL_NAME, response=response_text, tokens_used=int(tokens_generated), latency_ms=latency_ms, timestamp=datetime.utcnow().isoformat())
    except Exception as e:
        logger.exception("Inference error")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/chat")
async def chat(request: InferenceRequest):
    return await infer(request)

@app.post("/complete")
async def complete(request: InferenceRequest):
    if not request.messages:
        raise HTTPException(status_code=400, detail="No messages provided")
    simple_request = InferenceRequest(
        messages=[ChatMessage(role="user", content=request.messages[-1].content)],
        temperature=request.temperature,
        max_tokens=request.max_tokens,
        top_p=request.top_p,
        top_k=request.top_k,
        thinking=request.thinking,
        system_prompt=request.system_prompt,
    )
    return await infer(simple_request)

@app.get("/")
async def root():
    return {
        "name": "Gemma Inference API",
        "model": MODEL_NAME,
        "version": "1.0",
        "endpoints": {
            "/health": "Health check",
            "/info": "Model information",
            "/infer": "Run inference (POST)",
            "/chat": "Chat interface (POST)",
            "/complete": "Text completion (POST)"
        }
    }