import torch
import time
import uuid
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Optional, Union, Dict, Any
import uvicorn

app = FastAPI()

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

# INT4量子化設定
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4"
)

# モデルとトークナイザーの初期化
model_name = "Qwen/Qwen3-235B-A22B"
model_id = "qwen3-235b-a22b-int4"  # OpenAI API互換のモデルID

print(f"Loading model: {model_name}")
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=quantization_config,
    device_map="auto",
    torch_dtype=torch.float16,
    trust_remote_code=True
)
print("Model loaded successfully")

# OpenAI API互換のデータモデル
class ChatMessage(BaseModel):
    role: str
    content: str

class ChatCompletionRequest(BaseModel):
    model: str
    messages: List[ChatMessage]
    temperature: Optional[float] = 0.7
    top_p: Optional[float] = 0.9
    max_tokens: Optional[int] = 512
    stream: Optional[bool] = False

class ChatCompletionChoice(BaseModel):
    index: int
    message: ChatMessage
    finish_reason: str

class ChatCompletionUsage(BaseModel):
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int

class ChatCompletionResponse(BaseModel):
    id: str
    object: str = "chat.completion"
    created: int
    model: str
    choices: List[ChatCompletionChoice]
    usage: ChatCompletionUsage

class ModelInfo(BaseModel):
    id: str
    object: str = "model"
    created: int
    owned_by: str = "huggingface"

class ModelsResponse(BaseModel):
    object: str = "list"
    data: List[ModelInfo]

# チャット形式のプロンプト構築
def build_chat_prompt(messages: List[ChatMessage]) -> str:
    """Qwen形式のチャットプロンプトを構築"""
    prompt = ""
    for msg in messages:
        if msg.role == "system":
            prompt += f"System: {msg.content}\n\n"
        elif msg.role == "user":
            prompt += f"Human: {msg.content}\n\n"
        elif msg.role == "assistant":
            prompt += f"Assistant: {msg.content}\n\n"
    
    # 最後にAssistantのプロンプトを追加
    if messages[-1].role != "assistant":
        prompt += "Assistant: "
    
    return prompt

@app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
async def chat_completions(request: ChatCompletionRequest):
    """OpenAI互換のチャット補完エンドポイント"""
    try:
        # プロンプト構築
        prompt = build_chat_prompt(request.messages)
        
        # トークナイズ
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        input_length = inputs["input_ids"].shape[1]
        
        # 生成
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=request.max_tokens,
                temperature=request.temperature,
                top_p=request.top_p,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id
            )
        
        # デコード
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        # プロンプトを除去
        response_text = generated_text[len(prompt):].strip()
        
        # トークン数計算
        output_length = outputs[0].shape[0] - input_length
        
        # レスポンス構築
        response = ChatCompletionResponse(
            id=f"chatcmpl-{uuid.uuid4().hex[:8]}",
            created=int(time.time()),
            model=request.model,
            choices=[
                ChatCompletionChoice(
                    index=0,
                    message=ChatMessage(role="assistant", content=response_text),
                    finish_reason="stop"
                )
            ],
            usage=ChatCompletionUsage(
                prompt_tokens=input_length,
                completion_tokens=output_length,
                total_tokens=input_length + output_length
            )
        )
        
        return response
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/v1/models", response_model=ModelsResponse)
async def list_models():
    """利用可能なモデルのリストを返す"""
    return ModelsResponse(
        data=[
            ModelInfo(
                id=model_id,
                created=int(time.time())
            )
        ]
    )

@app.get("/health")
async def health():
    return {
        "status": "healthy",
        "model": model_name,
        "model_id": model_id,
        "quantization": "INT4"
    }

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