main.py

import os
import json
import logging
import asyncio
from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from pydantic import BaseModel, Field

from agents import Runner
from doc_agents.implementation_agent import implementation_agent
from job_manager import job_manager, Job, PipelineTracker, PipelineStep

app = FastAPI(
    title="Autonomous Document Agent API",
    description="Asynchronous Batch/Job API for Autonomous AI Document Creation",
    servers=[{"url": "http://localhost:8000", "description": "Local Development Server"}],
    version="1.0.0",
    summary="This API provides endpoints for submitting document creation requests and monitoring their progress."
)

# Enable CORS Middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


class AgentRequest(BaseModel):
    request: str = Field(..., description="Natural language document creation request.")


async def run_agent_pipeline(job: Job) -> None:
    """
    Asynchronous background worker executing the multi-agent OpenAI Agents SDK pipeline.

    Args:
        job (Job): The queued Job instance containing request details and execution state.

    Workflow:
        1. Updates job status to 'in_progress' and sets current_step to 'task_planning'.
        2. Executes Runner.run(implementation_agent, input=job.request, hooks=tracker).
        3. Parses output items and logs to record document path and execution plan.
        4. Updates job status to 'completed' or 'failed'.
    """
    job_id = job.job_id
    logging.info(f"Starting background execution pipeline for job: {job_id}")
    
    job_manager.update_job(job_id, status="in_progress", current_step=PipelineStep.TASK_PLANNING)
    tracker = PipelineTracker(job_manager, job_id)

    try:
        result = await Runner.run(
            implementation_agent,
            input=job.request,
            hooks=tracker
        )
        
        # Extract execution plan and document path from tool calls / logs
        doc_path = None
        plan_list = []

        for item in result.new_items:
            output_str = str(getattr(item, "output", ""))
            if "[DOCX_OUTPUT]" in output_str:
                doc_path = output_str.split("[DOCX_OUTPUT]")[-1].strip()
            elif hasattr(item, "function_name") and item.function_name == "generate_task_plan":
                try:
                    plan_list = json.loads(item.output)
                except Exception:
                    pass

        # Fallback inspection of agent logs
        if not doc_path:
            for log in job.agent_logs:
                if "[DOCX_OUTPUT]" in log:
                    doc_path = log.split("[DOCX_OUTPUT]")[-1].strip()
                elif "Generated Word document at:" in log:
                    doc_path = log.split("Generated Word document at:")[-1].strip()

        if not plan_list:
            plan_list = [
                {"step": 1, "description": "Generate task execution plan", "status": "completed"},
                {"step": 2, "description": "Co-author document specification", "status": "completed"},
                {"step": 3, "description": "Fetch domain context", "status": "completed"},
                {"step": 4, "description": "Validate document quality", "status": "completed"},
                {"step": 5, "description": "Render Word (.docx) document", "status": "completed"}
            ]

        job_manager.update_job(
            job_id,
            status="completed",
            current_step=PipelineStep.RENDERING,
            document_path=doc_path,
            execution_plan=plan_list
        )
        job_manager.append_log(job_id, f"[SUCCESS] Pipeline completed successfully. Output doc: {doc_path}")

    except Exception as e:
        logging.error(f"Pipeline failure for job {job_id}: {e}", exc_info=True)
        job_manager.update_job(job_id, status="failed", error=str(e))
        job_manager.append_log(job_id, f"[ERROR] Job execution failed: {e}")


@app.post("/agent", status_code=202)
async def submit_agent_request(body: AgentRequest, background_tasks: BackgroundTasks) -> dict:
    """
    Submits a natural language document creation request.
    Spawns background worker task and returns queued job state immediately.

    Args:
        body (AgentRequest): Pydantic body containing the natural language document prompt.
        background_tasks (BackgroundTasks): FastAPI background task launcher.

    Returns:
        dict: Initial queued job payload containing job_id and status.

    Input Example:
        POST /agent
        {
            "request": "Create a Product Requirements Document (PRD) for an AI Chatbot."
        }

    Response Example:
        HTTP 202 Accepted
        {
            "job_id": "8da8e816-ce07-421b-a3f6-4d1f82ab05c8",
            "status": "queued"
        }
    """
    if not body.request.strip():
        raise HTTPException(status_code=400, detail="Request field cannot be empty.")

    job = job_manager.create_job(body.request)
    background_tasks.add_task(run_agent_pipeline, job)

    return {
        "job_id": job.job_id,
        "status": job.status
    }


@app.get("/agent/job/{job_id}")
async def get_job_status(job_id: str) -> dict:
    """
    Polling endpoint returning the current job status, step progress, logs, and output path.

    Args:
        job_id (str): Unique UUID string identifier of the submitted job.

    Returns:
        dict: Complete job status payload formatted with status first.

    Response Example:
        HTTP 200 OK
        {
            "status": "in_progress",
            "job_id": "8da8e816-ce07-421b-a3f6-4d1f82ab05c8",
            "current_step": "document_drafting",
            "request": "Create a Product Requirements Document...",
            "execution_plan": [...],
            "agent_logs": [...],
            "document_path": None,
            "error": None,
            "created_at": "2026-07-25T02:00:00.000000"
        }
    """
    job = job_manager.get_job(job_id)
    if not job:
        raise HTTPException(status_code=404, detail=f"Job ID '{job_id}' not found.")
    return job.to_dict()


@app.get("/agent/job/{job_id}/download")
async def download_document(job_id: str) -> FileResponse:
    """
    Serves the generated Microsoft Word (.docx) file as a raw binary download stream.

    Args:
        job_id (str): Unique UUID string identifier of the completed job.

    Returns:
        FileResponse: Binary download stream with 'application/vnd.openxmlformats-officedocument.wordprocessingml.document' media type.

    Response Example:
        HTTP 200 OK (Binary stream attachment download)
    """
    job = job_manager.get_job(job_id)
    if not job:
        raise HTTPException(status_code=404, detail=f"Job ID '{job_id}' not found.")

    if job.status != "completed" or not job.document_path:
        raise HTTPException(status_code=400, detail=f"Job '{job_id}' is not yet completed or has no output file.")

    file_path = os.path.abspath(job.document_path)
    if not os.path.exists(file_path):
        raise HTTPException(status_code=404, detail=f"Output file '{job.document_path}' not found on server.")

    filename = os.path.basename(file_path)
    return FileResponse(
        path=file_path,
        media_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
        filename=filename
    )



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