from fastapi import FastAPI, HTTPException, UploadFile, File
from transformers import pipeline
import os
from typing import List
from pydantic import BaseModel

app = FastAPI()

# Initialize the AI model
model_name = "gpt-3.5-turbo"
generator = pipeline('text-generation', model=model_name)

class FileUpdate(BaseModel):
    filename: str
    content: str

@app.post("/generate")
def generate_text(prompt: str):
    try:
        result = generator(prompt, max_length=100)
        return {"response": result[0]['generated_text']}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/upload")
async def upload_file(file: UploadFile = File(...)):
    try:
        with open(f"./files/{file.filename}", "wb") as f:
            content = await file.read()
            f.write(content)
        return {"filename": file.filename}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/update")
def update_file(file_update: FileUpdate):
    try:
        with open(f"./files/{file_update.filename}", "w") as f:
            f.write(file_update.content)
        return {"status": "success"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/files")
def list_files():
    try:
        files = os.listdir("./files")
        return {"files": files}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))