from fastapi import FastAPI
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import os

app = FastAPI()

# Load HF_TOKEN from environment variable
HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
    raise ValueError("HF_TOKEN environment variable not set. Please set it in Hugging Face Spaces settings.")

# Load LLaMA model with updated token parameter
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B", token=HF_TOKEN)
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B", token=HF_TOKEN, device_map="auto")

def generate_code(query):
    prompt = f"Generate Python code for: {query}\n```python\n"
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
    outputs = model.generate(**inputs, max_new_tokens=200, temperature=0.7, do_sample=True)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    code_start = response.find("```python")
    code_end = response.find("```", code_start + 1)
    if code_start != -1 and code_end != -1:
        return response[code_start + 9:code_end].strip()
    return response.split("\n```python\n")[-1].strip()

@app.get("/generate")
async def generate(query: str):
    code = generate_code(query)
    return {"code": code}

@app.get("/")
async def root():
    return {"message": "CodeCraft Server is running!"}