# main.py

import logging
from contextlib import asynccontextmanager
import torch
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForCausalLM

# --- Configuration ---
# The repository ID for your model on the Hugging Face Hub
HF_REPO_ID = "rxmha125/Rx_Codex_V1_Tiny_test"
# Use GPU if available (CUDA), otherwise fallback to CPU
MODEL_LOAD_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# --- Logging Setup ---
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# --- Global variables to hold the model and tokenizer ---
model = None
tokenizer = None

# --- Application Lifespan (Model Loading) ---
@asynccontextmanager
async def lifespan(app: FastAPI):
    global model, tokenizer
    logger.info(f"API Startup: Loading model '{HF_REPO_ID}' to device '{MODEL_LOAD_DEVICE}'...")
    
    # Load the tokenizer from the Hub
    try:
        tokenizer = AutoTokenizer.from_pretrained(HF_REPO_ID)
        logger.info("✅ Tokenizer loaded successfully.")
    except Exception as e:
        logger.error(f"❌ FATAL: Tokenizer loading failed: {e}")
        # In a real app, you might want to handle this more gracefully
        # For Spaces, it will just fail to start, which is okay.

    # Load the model from the Hub
    try:
        model = AutoModelForCausalLM.from_pretrained(HF_REPO_ID)
        model.to(MODEL_LOAD_DEVICE)
        model.eval() # Set to evaluation mode for inference
        logger.info("✅ Model loaded successfully.")
    except Exception as e:
        logger.error(f"❌ FATAL: Model loading failed: {e}")

    yield # The API is now running

    # --- Code below this line runs on shutdown ---
    logger.info("API Shutting down.")
    model = None
    tokenizer = None


# --- Initialize FastAPI ---
app = FastAPI(
    title="Rx Codex V1-Tiny API",
    description="An API for generating text with the Rx_Codex_V1_Tiny model.",
    lifespan=lifespan
)

# --- Pydantic Models for API Data Validation ---
class GenerationRequest(BaseModel):
    prompt: str
    max_new_tokens: int = 150
    temperature: float = 0.7
    top_k: int = 50

class GenerationResponse(BaseModel):
    generated_text: str

# --- API Endpoints ---
@app.get("/")
def root():
    """A simple endpoint to check if the API is running."""
    status = "loaded" if model and tokenizer else "not loaded"
    return {"message": "Rx Codex V1-Tiny API is running", "model_status": status}

@app.post("/generate", response_model=GenerationResponse)
async def generate_text(request: GenerationRequest):
    """The main endpoint to generate text from a prompt."""
    if not model or not tokenizer:
        raise HTTPException(status_code=503, detail="Model is not ready. Please try again later.")
    
    logger.info(f"Received generation request for prompt: '{request.prompt}'")
    
    # --- CRITICAL: Format the prompt correctly for the model ---
    formatted_prompt = f"### Human:\n{request.prompt}\n\n### Assistant:"
    
    # Prepare the input text for the model
    inputs = tokenizer(formatted_prompt, return_tensors="pt").to(MODEL_LOAD_DEVICE)
    
    # Generate text using the model
    with torch.no_grad():
        output_sequences = model.generate(
            input_ids=inputs["input_ids"],
            attention_mask=inputs["attention_mask"],
            max_new_tokens=request.max_new_tokens,
            temperature=request.temperature,
            top_k=request.top_k,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id
        )
    
    # Decode the generated tokens back into text
    full_text = tokenizer.decode(output_sequences[0], skip_special_tokens=True)
    
    # Remove the original formatted prompt from the output to return only the new text
    generated_text = full_text[len(formatted_prompt):].strip()

    logger.info("Generation complete.")
    return GenerationResponse(generated_text=generated_text)


# --- Uvicorn Runner (for local testing) ---
if __name__ == "__main__":
    import uvicorn
    logger.info("Starting API locally via Uvicorn...")
    uvicorn.run(app, host="0.0.0.0", port=8000)