"""AsyncJob domain model for CleverAgents.

Represents an asynchronous job that executes plan phases in background
workers.  Jobs follow a strict state machine:

```
queued -> running -> succeeded
                  -> failed
                  -> cancelled
```

Only the transitions listed above are valid.  Attempting an invalid
transition raises ``InvalidJobTransitionError``.

## Fields

| Field            | Type              | Description                           |
|------------------|-------------------|---------------------------------------|
| ``job_id``       | ``str`` (ULID)    | Unique job identifier                 |
| ``plan_id``      | ``str`` (ULID)    | Plan this job belongs to              |
| ``phase``        | ``str``           | Plan phase (execute / apply)          |
| ``status``       | ``AsyncJobStatus``| Current job status                    |
| ``payload_json`` | ``str``           | Serialized job payload                |
| ``created_at``   | ``datetime``      | When the job was enqueued             |
| ``started_at``   | ``datetime|None`` | When a worker picked it up            |
| ``finished_at``  | ``datetime|None`` | When execution completed              |
| ``worker_id``    | ``str|None``      | ID of the worker processing this job  |
| ``last_heartbeat``| ``datetime|None``| Last heartbeat from the worker        |
| ``schema_version``| ``int``          | Payload schema version                |

Based on issue #312 (feat(async): add async command execution and workers).
"""

from __future__ import annotations

import json
import re
from datetime import UTC, datetime
from enum import StrEnum
from typing import Any

from pydantic import BaseModel, ConfigDict, Field, field_validator

# ULID is 26 characters, Crockford's base32
ULID_PATTERN = r"^[0-9A-HJKMNP-TV-Z]{26}$"
_ULID_RE = re.compile(ULID_PATTERN)

# Valid phases for async jobs
VALID_PHASES = frozenset({"execute", "apply"})


class AsyncJobStatus(StrEnum):
    """Status of an asynchronous job.

    Jobs follow a strict state machine:
    ``queued`` -> ``running`` -> ``succeeded`` / ``failed`` / ``cancelled``
    """

    QUEUED = "queued"
    RUNNING = "running"
    SUCCEEDED = "succeeded"
    FAILED = "failed"
    CANCELLED = "cancelled"


# Valid status transitions: from -> set of allowed targets
VALID_JOB_TRANSITIONS: dict[AsyncJobStatus, frozenset[AsyncJobStatus]] = {
    AsyncJobStatus.QUEUED: frozenset(
        {AsyncJobStatus.RUNNING, AsyncJobStatus.CANCELLED}
    ),
    AsyncJobStatus.RUNNING: frozenset(
        {AsyncJobStatus.SUCCEEDED, AsyncJobStatus.FAILED, AsyncJobStatus.CANCELLED}
    ),
    AsyncJobStatus.SUCCEEDED: frozenset(),
    AsyncJobStatus.FAILED: frozenset(),
    AsyncJobStatus.CANCELLED: frozenset(),
}

# Terminal states
TERMINAL_STATUSES: frozenset[AsyncJobStatus] = frozenset(
    {AsyncJobStatus.SUCCEEDED, AsyncJobStatus.FAILED, AsyncJobStatus.CANCELLED}
)


class InvalidJobTransitionError(Exception):
    """Raised when an invalid job status transition is attempted."""

    def __init__(self, from_status: AsyncJobStatus, to_status: AsyncJobStatus) -> None:
        self.from_status = from_status
        self.to_status = to_status
        super().__init__(
            f"Invalid job transition: {from_status.value} -> {to_status.value}"
        )


def can_transition_job(from_status: AsyncJobStatus, to_status: AsyncJobStatus) -> bool:
    """Check if a job status transition is valid.

    Args:
        from_status: Current job status.
        to_status: Target job status.

    Returns:
        True if the transition is valid.
    """
    if not isinstance(from_status, AsyncJobStatus):
        raise TypeError(f"from_status must be AsyncJobStatus, got {type(from_status)}")
    if not isinstance(to_status, AsyncJobStatus):
        raise TypeError(f"to_status must be AsyncJobStatus, got {type(to_status)}")
    return to_status in VALID_JOB_TRANSITIONS.get(from_status, frozenset())


class AsyncJob(BaseModel):
    """Domain model for an asynchronous job.

    An async job represents a unit of work (a plan phase) that is executed
    by a background worker.  Jobs are created when ``async.enabled`` is True
    and a plan transitions to the execute or apply phase.

    The ``payload_json`` field contains a JSON-serialized payload with the
    data needed to execute the job.  The ``schema_version`` field tracks
    the payload format version for forward compatibility.
    """

    job_id: str = Field(
        ...,
        description="Unique ULID identifier for this job",
        pattern=ULID_PATTERN,
    )
    plan_id: str = Field(
        ...,
        description="ULID of the plan this job belongs to",
        pattern=ULID_PATTERN,
    )
    phase: str = Field(
        ...,
        description="Plan phase this job executes (execute or apply)",
    )
    status: AsyncJobStatus = Field(
        default=AsyncJobStatus.QUEUED,
        description="Current job status",
    )
    payload_json: str = Field(
        default="{}",
        description="JSON-serialized job payload",
    )
    created_at: datetime = Field(
        default_factory=lambda: datetime.now(UTC),
        description="When the job was enqueued",
    )
    started_at: datetime | None = Field(
        default=None,
        description="When a worker picked up this job",
    )
    finished_at: datetime | None = Field(
        default=None,
        description="When execution completed",
    )
    worker_id: str | None = Field(
        default=None,
        description="ID of the worker processing this job",
    )
    last_heartbeat: datetime | None = Field(
        default=None,
        description="Last heartbeat timestamp from the worker",
    )
    schema_version: int = Field(
        default=1,
        ge=1,
        description="Payload schema version for forward compatibility",
    )
    error_message: str | None = Field(
        default=None,
        description="Error details when the job fails (for audit trail)",
    )

    @field_validator("phase")
    @classmethod
    def validate_phase(cls: type[AsyncJob], v: str) -> str:
        """Validate that phase is a valid async job phase."""
        if v not in VALID_PHASES:
            raise ValueError(f"phase must be one of {sorted(VALID_PHASES)}, got {v!r}")
        return v

    @field_validator("payload_json")
    @classmethod
    def validate_payload_json(cls: type[AsyncJob], v: str) -> str:
        """Validate that payload_json is valid JSON."""
        try:
            json.loads(v)
        except (json.JSONDecodeError, TypeError) as exc:
            raise ValueError(f"payload_json must be valid JSON: {exc}") from exc
        return v

    # -- Properties ---------------------------------------------------------

    @property
    def is_terminal(self) -> bool:
        """Check if the job is in a terminal state."""
        return self.status in TERMINAL_STATUSES

    @property
    def is_running(self) -> bool:
        """Check if the job is currently running."""
        return self.status == AsyncJobStatus.RUNNING

    @property
    def payload(self) -> dict[str, Any]:
        """Deserialize and return the job payload."""
        return dict(json.loads(self.payload_json))

    # -- Methods ------------------------------------------------------------

    def transition_to(self, target: AsyncJobStatus) -> None:
        """Transition the job to a new status.

        Validates that the transition is allowed by the state machine
        before applying the change.

        Args:
            target: The target status.

        Raises:
            TypeError: If target is not an AsyncJobStatus.
            InvalidJobTransitionError: If the transition is not allowed.
        """
        if not isinstance(target, AsyncJobStatus):
            raise TypeError(f"target must be AsyncJobStatus, got {type(target)}")
        if not can_transition_job(self.status, target):
            raise InvalidJobTransitionError(self.status, target)
        self.status = target

    def mark_running(self, worker_id: str) -> None:
        """Mark the job as running by a specific worker.

        Args:
            worker_id: ID of the worker picking up this job.

        Raises:
            ValueError: If worker_id is empty.
            InvalidJobTransitionError: If the job is not in QUEUED state.
        """
        if not worker_id or not worker_id.strip():
            raise ValueError("worker_id must be a non-empty string")
        self.transition_to(AsyncJobStatus.RUNNING)
        self.worker_id = worker_id
        self.started_at = datetime.now(UTC)
        self.last_heartbeat = self.started_at

    def mark_succeeded(self) -> None:
        """Mark the job as successfully completed.

        Raises:
            InvalidJobTransitionError: If the job is not in RUNNING state.
        """
        self.transition_to(AsyncJobStatus.SUCCEEDED)
        self.finished_at = datetime.now(UTC)

    def mark_failed(self, error: str | None = None) -> None:
        """Mark the job as failed.

        Args:
            error: Optional error message for audit trail.

        Raises:
            InvalidJobTransitionError: If the job is not in RUNNING state.
        """
        self.transition_to(AsyncJobStatus.FAILED)
        self.finished_at = datetime.now(UTC)
        if error:
            self.error_message = error

    def mark_cancelled(self) -> None:
        """Mark the job as cancelled.

        Raises:
            InvalidJobTransitionError: If the job is already terminal.
        """
        self.transition_to(AsyncJobStatus.CANCELLED)
        self.finished_at = datetime.now(UTC)

    def record_heartbeat(self) -> None:
        """Record a heartbeat from the worker.

        Raises:
            ValueError: If the job is not currently running.
        """
        if self.status != AsyncJobStatus.RUNNING:
            raise ValueError(
                f"Cannot record heartbeat for job in {self.status.value} state"
            )
        self.last_heartbeat = datetime.now(UTC)

    def as_cli_dict(self) -> dict[str, Any]:
        """Return a stable dictionary representation for CLI output."""
        result: dict[str, Any] = {
            "job_id": self.job_id,
            "plan_id": self.plan_id,
            "phase": self.phase,
            "status": self.status.value,
            "schema_version": self.schema_version,
            "created_at": self.created_at.isoformat(),
        }
        if self.started_at:
            result["started_at"] = self.started_at.isoformat()
        if self.finished_at:
            result["finished_at"] = self.finished_at.isoformat()
        if self.worker_id:
            result["worker_id"] = self.worker_id
        if self.last_heartbeat:
            result["last_heartbeat"] = self.last_heartbeat.isoformat()
        if self.error_message:
            result["error_message"] = self.error_message
        return result

    model_config = ConfigDict(
        str_strip_whitespace=True,
        validate_assignment=True,
        use_enum_values=False,
    )


def serialize_job_payload(
    plan_id: str,
    phase: str,
    extra: dict[str, Any] | None = None,
    schema_version: int = 1,
) -> str:
    """Serialize a job payload to JSON with schema versioning.

    Args:
        plan_id: The plan ULID.
        phase: The plan phase.
        extra: Additional payload data.
        schema_version: Schema version for forward compatibility.

    Returns:
        JSON string of the serialized payload.

    Raises:
        ValueError: If plan_id or phase is empty.
    """
    if not plan_id or not plan_id.strip():
        raise ValueError("plan_id must be a non-empty string")
    if not phase or not phase.strip():
        raise ValueError("phase must be a non-empty string")
    if schema_version < 1:
        raise ValueError("schema_version must be >= 1")
    payload: dict[str, Any] = {
        "schema_version": schema_version,
        "plan_id": plan_id,
        "phase": phase,
    }
    if extra:
        payload["extra"] = extra
    return json.dumps(payload, default=str)


def deserialize_job_payload(payload_json: str) -> dict[str, Any]:
    """Deserialize a job payload from JSON.

    Args:
        payload_json: The JSON string to deserialize.

    Returns:
        The deserialized payload dictionary.

    Raises:
        ValueError: If the JSON is invalid or missing required fields.
    """
    if not payload_json or not payload_json.strip():
        raise ValueError("payload_json must be a non-empty string")
    try:
        data: dict[str, Any] = json.loads(payload_json)
    except (json.JSONDecodeError, TypeError) as exc:
        raise ValueError(f"Invalid JSON payload: {exc}") from exc
    if not isinstance(data, dict):
        raise ValueError("Payload must be a JSON object")
    return data
