from fastapi import FastAPI, Depends, HTTPException
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from transformers import pipeline, AutoTokenizer
import torch
import jwt
import os
import re
from dotenv import load_dotenv
from typing import List, Dict, Optional
from pydantic import BaseModel

load_dotenv()
SECRET_KEY = os.getenv('SECRET_KEY')
security = HTTPBearer()

app = FastAPI()

model_name = 'Qwen/Qwen3-0.6B'
device = 'cuda' if torch.cuda.is_available() else 'cpu'
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
pipe = pipeline('text-generation', model=model_name, device=device, trust_remote_code=True)

class GenerateRequest(BaseModel):
    messages: List[Dict[str, str]]
    enable_thinking: Optional[bool] = False

def verify_token(credentials: HTTPAuthorizationCredentials = Depends(security)):
    try:
        payload = jwt.decode(credentials.credentials, SECRET_KEY, algorithms=['HS256'])
        return payload
    except Exception as e:
        raise HTTPException(status_code=401, detail=str(e))

@app.post('/generate')
def generate(req: GenerateRequest, user=Depends(verify_token)):
    try:
        prompt = tokenizer.apply_chat_template(req.messages, tokenize=False, add_generation_prompt=True)
        result = pipe(prompt, max_new_tokens=200)
        full_text = result[0]['generated_text']

        # Extract assistant response
        response_split = full_text.split('<|im_start|>assistant')
        content = response_split[-1] if len(response_split) > 1 else full_text

        # Handle thinking block based on flag
        if not req.enable_thinking:
            content = re.sub(r'<think>.*?</think>', '', content, flags=re.DOTALL)

        clean_response = content.replace('<|im_end|>', '').strip()

        return {'generated_text': clean_response}
    except Exception as e:
        return {'error': str(e)}

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