from fastapi import FastAPI, BackgroundTasks
from pydantic import BaseModel
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread
import uuid
import torch

app = FastAPI()

model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)

tasks = {}

class GenerateRequest(BaseModel):
    prompt: str
    system_prompt: str = "Ты — полезный ассистент Сайга. Отвечай кратко на русском языке."
    custom_id: str = None

def run_ai(task_id, prompt, system_prompt):
    try:
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": prompt}
        ]
        input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True)
        
        generation_kwargs = dict(
            input_ids=input_ids,
            streamer=streamer,
            max_new_tokens=256,
            do_sample=True,
            temperature=0.7,
            top_p=0.9
        )
        
        thread = Thread(target=model.generate, kwargs=generation_kwargs)
        thread.start()

        generated_text = ""
        tasks[task_id]["status"] = "generating"

        for new_text in streamer:
            clean_text = new_text.replace("<|eot_id|>", "").replace("<|end_of_text|>", "")
            generated_text += clean_text
            tasks[task_id]["result"] = generated_text.strip()
        
        tasks[task_id]["status"] = "ready"
    except Exception as e:
        tasks[task_id] = {"status": "error", "result": str(e)}

@app.post("/generate")
async def generate(req: GenerateRequest, background_tasks: BackgroundTasks):
    task_id = req.custom_id if req.custom_id else str(uuid.uuid4())
    tasks[task_id] = {"status": "pending", "result": ""}
    background_tasks.add_task(run_ai, task_id, req.prompt, req.system_prompt)
    return {"task_id": task_id}

@app.get("/check/{task_id}")
async def check(task_id: str):
    data = tasks.get(task_id)
    if not data:
        return {"task_id": task_id, "status": "not_found", "result": None}
    return {"task_id": task_id, "status": data["status"], "result": data["result"]}

@app.get("/")
async def root():
    return {"status": "running", "model": "Saiga Llama 3 Tiny"}