#!/usr/bin/env python3
# ============================================================
# ZabaanAI-v2: FastAPI Server for Production Deployment
# Run: uvicorn scripts.06_deploy_api:app --host 0.0.0.0 --port 8000
# ============================================================
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Optional, List
import torch
import os, warnings
warnings.filterwarnings('ignore')

app = FastAPI(title='ZabaanAI v2 API', version='2.0.0',
              description='Pakistan Multilingual AI - Supports Urdu, Punjabi, Sindhi, Pashto, Balochi, Saraiki, English, Roman Urdu')

USE_GPU    = os.getenv('USE_GPU', '0') == '1'
MODEL_PATH = os.getenv('MODEL_PATH', 'shaikhsalman/zabaanai-v2-sft')
DEVICE     = 'cuda' if USE_GPU else 'cpu'

# Load on startup
print(f'Loading model from {MODEL_PATH} on {DEVICE}...')
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel

tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
base = AutoModelForCausalLM.from_pretrained(
    'Qwen/Qwen2.5-7B-Instruct',
    device_map=DEVICE,
    load_in_4bit=(not USE_GPU),
    torch_dtype=torch.bfloat16 if USE_GPU else torch.float16,
    trust_remote_code=True
)
model = PeftModel.from_pretrained(base, MODEL_PATH)
model.eval()
print('Model ready!')

SYSTEM_MSG = (
    'You are ZabaanAI, a helpful multilingual AI assistant specializing in '
    'Pakistan languages. Respond in the same language as the user.'
)

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

class ChatRequest(BaseModel):
    messages: List[Message]
    max_tokens: int = 512
    temperature: float = 0.7
    top_p: float = 0.9

class GenerateRequest(BaseModel):
    prompt: str
    max_tokens: int = 512
    temperature: float = 0.7
    system: Optional[str] = None

def build_prompt(messages: List[Message], system: str = None) -> str:
    sys_msg = system or SYSTEM_MSG
    prompt = f'<|im_start|>system\n{sys_msg}<|im_end|>\n'
    for m in messages:
        role = 'user' if m.role in ('user', 'human') else 'assistant'
        prompt += f'<|im_start|>{role}\n{m.content}<|im_end|>\n'
    prompt += '<|im_start|>assistant\n'
    return prompt

@app.post('/v1/chat/completions')
async def chat_completions(req: ChatRequest):
    prompt = build_prompt(req.messages)
    inputs = tokenizer(prompt, return_tensors='pt').to(model.device)
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=req.max_tokens,
            temperature=req.temperature,
            do_sample=True,
            top_p=req.top_p,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
    return {'model': 'zabaanai-v2', 'choices': [{'message': {'role': 'assistant', 'content': response.strip()}}]}

@app.post('/v1/generate')
async def generate(req: GenerateRequest):
    sys_msg = req.system or SYSTEM_MSG
    prompt = f'<|im_start|>system\n{sys_msg}<|im_end|>\n<|im_start|>user\n{req.prompt}<|im_end|>\n<|im_start|>assistant\n'
    inputs = tokenizer(prompt, return_tensors='pt').to(model.device)
    with torch.no_grad():
        outputs = model.generate(
            **inputs, max_new_tokens=req.max_tokens,
            temperature=req.temperature, do_sample=True,
            pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id,
        )
    return {'text': tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True).strip()}

@app.get('/health')
async def health():
    return {'status': 'ok', 'model': 'zabaanai-v2'}

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