from fastapi import FastAPI
from pydantic import BaseModel
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

app = FastAPI(
    title="Rapnss DevOps-Ultra-125M API",
    description="A pure REST API for generating Python code via HTTPS using the custom Rapnss architecture",
    version="1.0.0"
)

# Configuration
REPO_ID = "Rapnss/DevOps-Ultra-125M"
DEVICE = "cpu"

model = None
tokenizer = None

@app.on_event("startup")
def load_model():
    global model, tokenizer
    print("Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(REPO_ID, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    
    print("Loading custom Rapnss model weights directly from the Hub...")
    # trust_remote_code=True automatically downloads your custom PyTorch architecture files!
    model = AutoModelForCausalLM.from_pretrained(REPO_ID, trust_remote_code=True)
    model.to(DEVICE)
    model.eval()
    print("API Ready for Requests!")

class CodeRequest(BaseModel):
    instruction: str
    input_code: str = ""
    max_tokens: int = 150
    temperature: float = 0.7

@app.post("/v1/generate")
def generate_code(req: CodeRequest):
    if model is None or tokenizer is None:
        return {"error": "Model not loaded yet."}
        
    prompt = f"Instruction: {req.instruction}\n"
    if req.input_code:
        prompt += f"Input: {req.input_code}\n"
    prompt += "Output:\n"
    
    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(DEVICE)
    
    output_ids = model.generate(
        input_ids,
        max_new_tokens=req.max_tokens,
        do_sample=False,
        pad_token_id=tokenizer.eos_token_id
    )
    
    response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    
    # Clean up the output to only return the generated code
    if "Output:\n" in response:
        response = response.split("Output:\n")[-1]
        
    return {"generated_code": response.strip()}

@app.get("/")
def health_check():
    return {"status": "active", "model": "Rapnss DevOps-Ultra-125M API is running!"}