import asyncio
from bson import ObjectId
from app.database import db
from app.models.assignment import AssignmentInput
from app.services.ai_service import generate_paper
from app.services.ws_service import manager

async def generate_paper_in_background(assignment_id: str, input_data: AssignmentInput):
    try:
        # 1. Update status to processing
        await db.assignments.update_one({"_id": ObjectId(assignment_id)}, {"$set": {"status": "processing"}})
        await manager.broadcast(assignment_id, {
            "type": "JOB_STARTED",
            "assignmentId": assignment_id,
            "message": "Starting paper generation..."
        })

        # 2. Progress update
        await manager.broadcast(assignment_id, {
            "type": "JOB_PROGRESS",
            "assignmentId": assignment_id,
            "progress": 20,
            "message": "Analyzing assessment criteria..."
        })

        # 3. Invoke LLM generation service
        paper = await generate_paper(input_data, assignment_id)

        # 4. Progress update
        await manager.broadcast(assignment_id, {
            "type": "JOB_PROGRESS",
            "assignmentId": assignment_id,
            "progress": 70,
            "message": "Structuring exam sections..."
        })

        # 5. Save generated paper representation to database
        await db.generated_papers.insert_one(paper.copy()) # copy since _id is injected sometimes

        # 6. Update main assignment document status
        await db.assignments.update_one({"_id": ObjectId(assignment_id)}, {"$set": {"status": "completed", "result": paper}})

        # 7. Progress update
        await manager.broadcast(assignment_id, {
            "type": "JOB_PROGRESS",
            "assignmentId": assignment_id,
            "progress": 90,
            "message": "Finalizing exam document..."
        })

        # 8. Stream completed message
        await manager.broadcast(assignment_id, {
            "type": "JOB_COMPLETED",
            "assignmentId": assignment_id,
            "paper": paper
        })
    except Exception as e:
        print(f"[AI Engine] Generation failed for assignment {assignment_id}: {e}")
        error_message = str(e)

        # Update assignment status to failed
        await db.assignments.update_one({"_id": ObjectId(assignment_id)}, {"$set": {"status": "failed", "error": error_message}})

        # Broadcast failure
        await manager.broadcast(assignment_id, {
            "type": "JOB_FAILED",
            "assignmentId": assignment_id,
            "error": error_message
        })