src/api/v1/drafts.py

import json
from collections.abc import AsyncIterator
from typing import Any

from fastapi import APIRouter, Response
from fastapi.responses import StreamingResponse

from src.api.dependencies import DbSessionDep, GuestSessionHeaderDep, OptionalCurrentUserDep
from src.api.errors import error_response
from src.models.schemas import (
    DraftCheckRequest,
    DraftCheckResponse,
    DraftExportRequest,
    DraftPatchRequest,
    DraftPatchResponse,
    DraftPlanningResult,
    DraftRequest,
    DraftResponse,
)
from src.services.auth import AuthServiceError
from src.services.drafting import (
    check_draft_grounding,
    export_docx,
    generate_draft,
    generate_draft_stream,
    patch_draft_selection,
    plan_draft,
)
from src.services.guest_quota import ensure_guest_session_id
from src.services.retrieval_access import RetrievalAccessCoordinator
from src.services.usage_tracking import UsageTrackingService

router = APIRouter(prefix="/drafts", tags=["drafts"])


@router.post("/plan", response_model=DraftPlanningResult)
async def plan_draft_request(request: DraftRequest) -> DraftPlanningResult | dict:
    try:
        return plan_draft(request)
    except Exception as e:
        return error_response(
            status_code=500,
            code="draft_planning_failed",
            message=str(e),
        )


@router.post("", response_model=DraftResponse)
async def create_draft(
    request: DraftRequest,
    response: Response,
    db: DbSessionDep,
    current_user: OptionalCurrentUserDep,
    guest_session_id: GuestSessionHeaderDep = None,
) -> DraftResponse | dict:
    resolved_guest_session_id: str | None = None
    try:
        if current_user is None:
            resolved_guest_session_id = ensure_guest_session_id(guest_session_id)
            response.headers["X-Guest-Session-Id"] = resolved_guest_session_id

        coordinator = RetrievalAccessCoordinator(db)
        draft_response = generate_draft(
            request=request,
            coordinator=coordinator,
            current_user=current_user,
            guest_session_id=resolved_guest_session_id or guest_session_id,
        )
        if current_user is not None and draft_response.html:
            UsageTrackingService.record_event_safe(
                event_type="draft_created",
                user_id=current_user.id,
                source="draft",
                metadata={"template_key": request.template_key.value},
            )
        return draft_response

    except AuthServiceError as exc:
        return error_response(
            status_code=exc.status_code,
            code=exc.code,
            message=exc.message,
        )
    except Exception as e:
        return error_response(
            status_code=500,
            code="draft_generation_failed",
            message=str(e),
        )


@router.post("/stream")
async def create_draft_stream(
    request: DraftRequest,
    response: Response,
    db: DbSessionDep,
    current_user: OptionalCurrentUserDep,
    guest_session_id: GuestSessionHeaderDep = None,
) -> StreamingResponse:
    resolved_guest_session_id: str | None = None
    try:
        if current_user is None:
            resolved_guest_session_id = ensure_guest_session_id(guest_session_id)

        coordinator = RetrievalAccessCoordinator(db)

        async def event_generator() -> AsyncIterator[str]:
            tracked = False
            try:
                async for frame in generate_draft_stream(
                    request=request,
                    coordinator=coordinator,
                    current_user=current_user,
                    guest_session_id=resolved_guest_session_id or guest_session_id,
                ):
                    if (
                        not tracked
                        and isinstance(frame, dict)
                        and frame.get("event") == "done"
                        and isinstance(frame.get("data"), dict)
                        and frame["data"].get("html")
                        and current_user is not None
                    ):
                        UsageTrackingService.record_event_safe(
                            event_type="draft_created",
                            user_id=current_user.id,
                            source="draft_stream",
                            metadata={"template_key": request.template_key.value},
                        )
                        tracked = True
                    yield f"event: {frame['event']}\ndata: {json.dumps(frame['data'], ensure_ascii=False)}\n\n"
            except Exception as exc:
                yield f"event: error\ndata: {json.dumps({'message': str(exc)})}\n\n"

        headers = {"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}
        if resolved_guest_session_id:
            headers["X-Guest-Session-Id"] = resolved_guest_session_id
        return StreamingResponse(
            event_generator(),
            media_type="text/event-stream",
            headers=headers,
        )

    except AuthServiceError as exc:
        return error_response(status_code=exc.status_code, code=exc.code, message=exc.message)
    except Exception as e:
        return error_response(status_code=500, code="draft_stream_failed", message=str(e))


@router.post("/check", response_model=DraftCheckResponse)
async def check_draft(request: DraftCheckRequest) -> DraftCheckResponse | dict:
    try:
        checked_html, uncertain_count, issues = check_draft_grounding(
            request.html, request.requirement, request.citations, request.template_key.value
        )
        return DraftCheckResponse(checked_html=checked_html, uncertain_count=uncertain_count, issues=issues)
    except Exception as e:
        return error_response(status_code=500, code="draft_check_failed", message=str(e))


@router.post("/patch", response_model=DraftPatchResponse)
async def patch_draft(request: DraftPatchRequest) -> DraftPatchResponse | dict:
    try:
        replacement = patch_draft_selection(request.selected_html, request.instruction)
        return DraftPatchResponse(replacement_html=replacement)
    except Exception as e:
        return error_response(status_code=500, code="draft_patch_failed", message=str(e))


@router.post("/export.docx")
async def export_draft_docx(
    request: DraftExportRequest,
    current_user: OptionalCurrentUserDep,
) -> Any:
    try:
        buffer = export_docx(request)
        headers = {
            "Content-Disposition": "attachment; filename=draft_export.docx"
        }
        return StreamingResponse(
            buffer,
            media_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
            headers=headers
        )
    except Exception as e:
        return error_response(
            status_code=500,
            code="docx_export_failed",
            message=str(e),
        )