backend/app/api/workflows.py

from __future__ import annotations

from uuid import UUID

from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy.orm import Session, joinedload, selectinload

from app.core.dependencies import get_current_user
from app.database.session import get_db
from app.models.document_version import DocumentVersion
from app.models.user import User
from app.models.workflow_run import WorkflowRun, WorkflowStatus
from app.models.workspace import Workspace

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


def _workflow_response(workflow: WorkflowRun) -> dict:
    dv = workflow.document_version

    # Determine completed stages from checkpoints
    checkpoints = workflow.checkpoints or []
    completed_stages = [cp.agent_name for cp in checkpoints]

    # Infer current node from status + checkpoints
    if workflow.status == WorkflowStatus.COMPLETED:
        current_node = "COMPLETED"
    elif workflow.status == WorkflowStatus.WAITING_FOR_REVIEW:
        current_node = "HUMAN_REVIEW"
    elif workflow.status == WorkflowStatus.FAILED:
        current_node = "FAILED"
    elif workflow.status == WorkflowStatus.CANCELLED:
        current_node = "CANCELLED"
    elif workflow.status == WorkflowStatus.PENDING:
        current_node = "PENDING"
    else:
        # RUNNING — read live progress from PROGRESS checkpoint
        progress_cp = None
        for cp in reversed(checkpoints):
            if cp.agent_name == "PROGRESS":
                progress_cp = cp
                break
        if progress_cp and progress_cp.state:
            current_node = progress_cp.state.get("current_node", "EXTRACTION")
        else:
            current_node = "EXTRACTION"

    # Get validation results from checkpoint state if available
    validation_results = None
    for cp in reversed(checkpoints):
        if cp.state and "validation_results" in cp.state:
            validation_results = cp.state["validation_results"]
            break

    return {
        "id": str(workflow.id),
        "workspace_id": str(workflow.workspace_id),
        "document_version_id": str(workflow.document_version_id),
        "document_id": str(dv.document_id) if dv else None,
        "filename": dv.filename if dv else None,
        "status": workflow.status.value,
        "current_node": current_node,
        "completed_stages": completed_stages,
        "validation_results": validation_results,
        "started_at": workflow.started_at,
        "completed_at": workflow.completed_at,
    }


def _sync_workflow_review_state(workflow: WorkflowRun, db: Session) -> None:
    """
    If a workflow is WAITING_FOR_REVIEW but has no PENDING proposals left
    (all proposals were approved, rejected, or archived), resume/complete the workflow.
    """
    if workflow.status == WorkflowStatus.WAITING_FOR_REVIEW and workflow.document_version_id:
        from app.services.proposal_review_service import ProposalReviewService
        service = ProposalReviewService(db)
        pending = service.get_pending_for_document_version(
            workflow.workspace_id,
            workflow.document_version_id,
        )
        if not pending:
            from app.workflow import WorkflowExecutor
            from datetime import datetime, timezone
            try:
                executor = WorkflowExecutor(db)
                try:
                    executor.resume(workflow)
                finally:
                    executor.close()
            except Exception:
                workflow.status = WorkflowStatus.FAILED
                if not workflow.completed_at:
                    workflow.completed_at = datetime.now(timezone.utc)
                db.commit()


@router.get("")
def list_workflows(
    workspace_id: UUID,
    status: str | None = None,
    current_user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
):
    workspace = (
        db.query(Workspace)
        .filter(
            Workspace.id == workspace_id,
            Workspace.created_by == current_user.id,
        )
        .first()
    )
    if workspace is None:
        raise HTTPException(
            status_code=403,
            detail="You do not have access to this workspace.",
        )

    query = (
        db.query(WorkflowRun)
        .options(
            joinedload(WorkflowRun.document_version),
            selectinload(WorkflowRun.checkpoints),
        )
        .filter(WorkflowRun.workspace_id == workspace_id)
    )

    if status:
        try:
            ws = WorkflowStatus(status.upper())
            query = query.filter(WorkflowRun.status == ws)
        except ValueError:
            pass

    workflows = query.order_by(WorkflowRun.started_at.desc()).all()
    for w in workflows:
        _sync_workflow_review_state(w, db)
    return [_workflow_response(w) for w in workflows]


@router.get("/by-document-version/{document_version_id}")
def get_workflow_by_document_version(
    document_version_id: UUID,
    current_user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
):
    workflow = (
        db.query(WorkflowRun)
        .options(
            joinedload(WorkflowRun.document_version),
            selectinload(WorkflowRun.checkpoints),
        )
        .filter(WorkflowRun.document_version_id == document_version_id)
        .order_by(WorkflowRun.started_at.desc())
        .first()
    )

    if workflow is None:
        raise HTTPException(status_code=404, detail="No workflow found for this document version.")

    workspace = (
        db.query(Workspace)
        .filter(
            Workspace.id == workflow.workspace_id,
            Workspace.created_by == current_user.id,
        )
        .first()
    )
    if workspace is None:
        raise HTTPException(status_code=403, detail="You do not have access to this workflow.")

    # Auto-detect stuck workflows: if RUNNING for >5 min with no progress, mark as FAILED
    if workflow.status == WorkflowStatus.RUNNING:
        from datetime import datetime, timezone, timedelta
        now = datetime.now(timezone.utc)
        started = workflow.started_at.replace(tzinfo=timezone.utc) if workflow.started_at.tzinfo is None else workflow.started_at
        if now - started > timedelta(minutes=5):
            # Check if there's been recent progress
            latest_checkpoint = None
            for cp in reversed(workflow.checkpoints or []):
                if cp.agent_name == "PROGRESS":
                    latest_checkpoint = cp
                    break
            stale = True
            if latest_checkpoint and latest_checkpoint.created_at:
                cp_time = latest_checkpoint.created_at.replace(tzinfo=timezone.utc) if latest_checkpoint.created_at.tzinfo is None else latest_checkpoint.created_at
                if now - cp_time < timedelta(minutes=5):
                    stale = False
            if stale:
                workflow.status = WorkflowStatus.FAILED
                db.commit()

    _sync_workflow_review_state(workflow, db)
    return _workflow_response(workflow)


@router.get("/{workflow_id}")
def get_workflow(
    workflow_id: UUID,
    current_user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
):
    workflow = (
        db.query(WorkflowRun)
        .options(
            joinedload(WorkflowRun.document_version),
            selectinload(WorkflowRun.checkpoints),
        )
        .filter(WorkflowRun.id == workflow_id)
        .first()
    )

    if workflow is None:
        raise HTTPException(
            status_code=404,
            detail="Workflow not found.",
        )

    _sync_workflow_review_state(workflow, db)

    workspace = (
        db.query(Workspace)
        .filter(
            Workspace.id == workflow.workspace_id,
            Workspace.created_by == current_user.id,
        )
        .first()
    )

    if workspace is None:
        raise HTTPException(
            status_code=403,
            detail="You do not have access to this workflow.",
        )

    # Auto-detect stuck workflows: if RUNNING for >5 min with no progress, mark as FAILED
    if workflow.status == WorkflowStatus.RUNNING:
        from datetime import datetime, timezone, timedelta
        now = datetime.now(timezone.utc)
        started = workflow.started_at.replace(tzinfo=timezone.utc) if workflow.started_at.tzinfo is None else workflow.started_at
        if now - started > timedelta(minutes=5):
            # Check if there's been recent progress
            latest_checkpoint = None
            for cp in reversed(workflow.checkpoints or []):
                if cp.agent_name == "PROGRESS":
                    latest_checkpoint = cp
                    break
            stale = True
            if latest_checkpoint and latest_checkpoint.created_at:
                cp_time = latest_checkpoint.created_at.replace(tzinfo=timezone.utc) if latest_checkpoint.created_at.tzinfo is None else latest_checkpoint.created_at
                if now - cp_time < timedelta(minutes=5):
                    stale = False
            if stale:
                workflow.status = WorkflowStatus.FAILED
                db.commit()

    return _workflow_response(workflow)


@router.post("/{workflow_id}/cancel", status_code=200)
def cancel_workflow(
    workflow_id: UUID,
    current_user: User = Depends(get_current_user),
    db: Session = Depends(get_db),
):
    """Mark a stuck/running workflow as FAILED so it stops polling."""
    workflow = (
        db.query(WorkflowRun)
        .filter(WorkflowRun.id == workflow_id)
        .first()
    )

    if workflow is None:
        raise HTTPException(status_code=404, detail="Workflow not found.")

    workspace = (
        db.query(Workspace)
        .filter(Workspace.id == workflow.workspace_id, Workspace.created_by == current_user.id)
        .first()
    )
    if workspace is None:
        raise HTTPException(status_code=403, detail="You do not have access to this workflow.")

    if workflow.status in (WorkflowStatus.COMPLETED, WorkflowStatus.CANCELLED):
        return {"id": str(workflow.id), "status": workflow.status.value, "message": "Already terminal."}

    workflow.status = WorkflowStatus.FAILED
    db.commit()

    return {"id": str(workflow.id), "status": "FAILED", "message": "Workflow marked as failed."}