from fastapi import FastAPI, UploadFile, File, Form
from typing import List, Optional
import uvicorn
from agent import generate_task_prompt

import io
import pypdf
import docx2txt

app = FastAPI(title="Task Prompting API")

@app.post("/generate_prompt")
async def generate_prompt(
    description: str = Form(...),
    field: str = Form(...),
    files: Optional[List[UploadFile]] = File(None)
):
    files_data = []
    
    if files:
        for file in files:
            content = await file.read()
            filename_lower = file.filename.lower()
            
            try:
                if filename_lower.endswith(".pdf"):
                    reader = pypdf.PdfReader(io.BytesIO(content))
                    text = "\n".join([page.extract_text() or "" for page in reader.pages])
                    files_data.append({"filename": file.filename, "content": text})
                elif filename_lower.endswith(".docx"):
                    text = docx2txt.process(io.BytesIO(content))
                    files_data.append({"filename": file.filename, "content": text})
                else:
                    decoded_content = content.decode('utf-8')
                    files_data.append({"filename": file.filename, "content": decoded_content})
            except Exception as e:
                print(f"Error processing {file.filename}: {e}")
                pass # Skip files that fail encoding or extraction
                
    try:
        generated_prompt = generate_task_prompt(description, field, files_data)
        return {"status": "success", "prompt": generated_prompt}
    except Exception as e:
        return {"status": "error", "message": str(e)}

if __name__ == "__main__":
    uvicorn.run("api:app", host="0.0.0.0", port=8000, reload=True)