import os

import uvicorn
import torch
from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse, Response
from pydantic import BaseModel
from transformers import pipeline


# Utwórz instancję FastAPI
app = FastAPI(
    title="Bielik Text Generation API",
    description="API do generowania tekstu za pomocą modelu Bielik-1.5B-v3.0-Instruct",
    version="1.0.0"
)

# Ścieżka do modelu - Hugging Face automatycznie pobierze model
MODEL_NAME = "speakleash/Bielik-1.5B-v3.0-Instruct"
generator = None  # Zostanie załadowany później


# Model wejściowy dla POST request
class GenerationRequest(BaseModel):
    prompt: str
    max_new_tokens: int = 50
    temperature: float = 0.7
    top_p: float = 0.9


@app.on_event("startup")
async def startup_event():
    """
    Ładowanie modelu podczas uruchamiania aplikacji.
    To zajmie trochę czasu, ale dzieje się tylko raz.
    """
    global generator

    print(f"Ładowanie modelu: {MODEL_NAME}...")
    try:
        # Możesz dostosować device=0 (GPU) lub device=-1 (CPU) w zależności od wybranej maszyny Space
        # Free tier spaces usually run on CPU, unless you explicitly select a GPU.
        # It's safer to not specify device if you want it to auto-detect or default to CPU.
        generator = pipeline(
            "text-generation",
            model=MODEL_NAME,
            # device=0 if torch.cuda.is_available() else -1 # Odkomentuj dla detekcji GPU
        )
        print("Model załadowany pomyślnie!")
    except Exception as e:
        print(f"Błąd ładowania modelu: {e}")
        # Możesz zdecydować, czy aplikacja ma zakończyć działanie, czy kontynuować bez modelu
        # W przypadku błędu ładowania modelu, endpoint generacji tekstu będzie zwracał błąd
        generator = None  # Ustaw na None, aby sygnalizować problem


@app.get("/")
async def root():
    """
    Główny endpoint (health check).
    """
    return {"message": "Bielik Text Generation API is running!"}


@app.post("/generate")
async def generate_text(request: GenerationRequest):
    """
    Endpoint do generowania tekstu na podstawie promptu.
    """
    if generator is None:
        raise HTTPException(status_code=503, detail="Model nie został załadowany lub wystąpił błąd.")

    try:
        generated_text = generator(
            request.prompt,
            max_new_tokens=request.max_new_tokens,
            temperature=request.temperature,
            top_p=request.top_p,
            do_sample=True,  # Ważne dla generowania z temperaturą
            return_full_text=False # Dodaj ten parametr, aby model nie zwracał od razu promptu
        )
        # Pipeline zwraca listę słowników, bierzemy pierwszy wynik
        #gen_text = {"generated_text": generated_text[0]["generated_text"]}

        response_data = {"generated_text": generated_text[0]["generated_text"]}
        return JSONResponse(
            content=response_data,
            media_type="application/json; charset=utf-8"
        )
        #return Response(content=generated_text[0]["generated_text"], media_type="text/plain; charset=utf-8")
        # return {"generated_text": generated_text[0]["generated_text"]}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Błąd podczas generowania tekstu: {e}")


# Uruchamianie serwera Uvicorn bezpośrednio (dla Dockera)
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", 7860)))