"""Async workflow submission — fire-and-forget execution with polling.

Usage::

    from core.workflow import submit, wait

    job = await submit(agent, {"description": "kitchen sink leaking"})
    result = await wait(job, timeout=120)

Internally uses ``asyncio.create_task()`` with an in-memory job registry.
For production use cases requiring persistence, plug in a real job queue
(Redis, Celery, etc.) behind the same interface.
"""

from __future__ import annotations

import asyncio
import logging
import time
from typing import TYPE_CHECKING, Literal
from uuid import uuid4

from pydantic import BaseModel, Field

if TYPE_CHECKING:
    from core.agent import Agent
    from core.types.agents import AgentDefinition
    from core.types.execution import ExecutionBudget

from core.types.execution import ExecutionResult

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Job model
# ---------------------------------------------------------------------------

JobStatus = Literal["pending", "running", "completed", "failed", "awaiting_approval"]


class WorkflowJob(BaseModel):
    """Handle for a background workflow execution.

    Attributes
    ----------
    job_id
        Unique job identifier.
    status
        Current status.
    result
        Populated when the job completes (or fails).
    error
        Error message if the job failed.
    submitted_at
        Epoch timestamp when the job was submitted.
    completed_at
        Epoch timestamp when the job finished (or None).
    """

    job_id: str = Field(default_factory=lambda: f"job_{uuid4().hex[:12]}")
    status: JobStatus = "pending"
    result: ExecutionResult | None = None
    error: str | None = None
    submitted_at: float = Field(default_factory=time.time)
    completed_at: float | None = None


# ---------------------------------------------------------------------------
# In-memory job registry
# ---------------------------------------------------------------------------

_jobs: dict[str, WorkflowJob] = {}
_tasks: dict[str, asyncio.Task[None]] = {}


def get_job(job_id: str) -> WorkflowJob | None:
    """Lookup a job by ID."""
    return _jobs.get(job_id)


def list_jobs(*, status: JobStatus | None = None) -> list[WorkflowJob]:
    """List all jobs, optionally filtered by status."""
    if status is None:
        return list(_jobs.values())
    return [j for j in _jobs.values() if j.status == status]


# ---------------------------------------------------------------------------
# submit / wait
# ---------------------------------------------------------------------------


async def submit(
    agent_or_definition: Agent | AgentDefinition,
    inputs: dict[str, object] | str,
    *,
    pathway: object | None = None,
    budget: ExecutionBudget | None = None,
    state_store: object | None = None,
) -> WorkflowJob:
    """Submit a workflow for background execution.

    Parameters
    ----------
    agent_or_definition
        Agent or AgentDefinition to execute.
    inputs
        Payload — either a message string or a dict of inputs.
    pathway
        Optional pathway to deploy.
    budget
        Execution budget.
    state_store
        Optional state store.

    Returns
    -------
    WorkflowJob
        A handle that can be polled or awaited.
    """
    from core.runner import Session
    from core.types.execution import ExecutionMode

    job = WorkflowJob()
    _jobs[job.job_id] = job

    message = inputs if isinstance(inputs, str) else str(inputs.get("message", inputs))

    session = Session(
        agent_or_definition,
        pathway=pathway,
        state_store=state_store,
        mode=ExecutionMode.WORKFLOW,
    )

    async def _run() -> None:
        job.status = "running"
        try:
            result = await session.run(message, budget=budget)
            job.result = result
            job.status = "completed" if result.success else "failed"
            if not result.success:
                job.error = result.error or result.response
        except Exception as e:
            job.status = "failed"
            job.error = str(e)
            logger.exception("Workflow job %s failed", job.job_id)
        finally:
            job.completed_at = time.time()

    task = asyncio.create_task(_run(), name=f"workflow:{job.job_id}")
    _tasks[job.job_id] = task
    return job


async def wait(job: WorkflowJob, *, timeout: float = 300) -> ExecutionResult:
    """Block until the job completes or timeout.

    Uses ``asyncio.wait_for`` on the underlying task — no polling needed.

    Parameters
    ----------
    job
        The WorkflowJob returned by ``submit()``.
    timeout
        Maximum seconds to wait.

    Returns
    -------
    ExecutionResult

    Raises
    ------
    TimeoutError
        If the job doesn't complete within *timeout*.
    RuntimeError
        If the job failed without producing an ExecutionResult.
    """
    task = _tasks.get(job.job_id)
    if task is not None:
        try:
            await asyncio.wait_for(asyncio.shield(task), timeout=timeout)
        except TimeoutError:
            raise TimeoutError(
                f"Workflow job {job.job_id} did not complete within {timeout}s"
            ) from None
        except Exception:
            pass  # Error is captured in job.error

    # Refresh from registry
    job = _jobs.get(job.job_id, job)

    if job.result is not None:
        return job.result

    if job.error:
        raise RuntimeError(f"Workflow job {job.job_id} failed: {job.error}")

    raise RuntimeError(f"Workflow job {job.job_id} ended without result")


__all__ = [
    "WorkflowJob",
    "get_job",
    "list_jobs",
    "submit",
    "wait",
]
