"""
aaiclick.orchestration.models - Data models for orchestration backend.

This module defines SQLModel models for jobs, tasks, execution workers, groups, and dependencies.
All IDs are snowflake IDs (64-bit integers) generated using aaiclick.snowflake.
"""

from collections.abc import Sequence
from datetime import datetime
from typing import Any, ClassVar, Literal, Union, cast

from sqlalchemy import BigInteger, Boolean, ForeignKey, Index, String, UniqueConstraint
from sqlalchemy.orm import Mapped
from sqlmodel import JSON, Column, Field, Relationship, SQLModel

from ..datetime_utils import utc_now
from .runner_config import ENTRY_MODULE, EntryType
from .task_registry import register_task

# Enum columns are stored as plain ``String`` and validated by their ``Literal``
# type plus boundary validation (SQLModel/Pydantic on write, the CLI's
# ``choices=``, API request models). No DB CHECK constraints — widening a closed
# set is then a code change, not a hand-written constraint migration. Native
# ENUM types stay avoided (``ALTER TYPE ADD VALUE`` is non-transactional in
# Postgres). See CLAUDE.md, "Prefer Literal".

# Dependency type constants
DEPENDENCY_TASK = "task"
DEPENDENCY_GROUP = "group"
DEPENDENCY_TYPES = [DEPENDENCY_TASK, DEPENDENCY_GROUP]

# Type alias for dependency type annotations (Literal requires hardcoded values)
DependencyType = Literal["task", "group"]


RUN_SCHEDULED = "SCHEDULED"
RUN_MANUAL = "MANUAL"
RunType = Literal["SCHEDULED", "MANUAL"]
"""How a job run was triggered."""


JOB_PENDING = "PENDING"
JOB_RUNNING = "RUNNING"
JOB_COMPLETED = "COMPLETED"
JOB_FAILED = "FAILED"
JOB_CANCELLED = "CANCELLED"
JobStatus = Literal["PENDING", "RUNNING", "COMPLETED", "FAILED", "CANCELLED"]
"""Job execution status."""


TASK_PENDING = "PENDING"
TASK_CLAIMED = "CLAIMED"
TASK_RUNNING = "RUNNING"
TASK_COMPLETED = "COMPLETED"
TASK_FAILED = "FAILED"
TASK_CANCELLED = "CANCELLED"
TASK_PENDING_CLEANUP = "PENDING_CLEANUP"
TASK_UPSTREAM_FAILED = "UPSTREAM_FAILED"
TaskStatus = Literal[
    "PENDING", "CLAIMED", "RUNNING", "COMPLETED", "FAILED", "CANCELLED", "PENDING_CLEANUP", "UPSTREAM_FAILED"
]
"""Task execution status.

``UPSTREAM_FAILED`` is a terminal state assigned by the cascade sweep when a
task's transitive upstream is ``FAILED``, ``CANCELLED``, or ``UPSTREAM_FAILED``.
Job rollup treats it as a failure."""


EXECUTION_WORKER_ACTIVE = "ACTIVE"
EXECUTION_WORKER_IDLE = "IDLE"
EXECUTION_WORKER_STOPPING = "STOPPING"
EXECUTION_WORKER_STOPPED = "STOPPED"
ExecutionWorkerStatus = Literal["ACTIVE", "IDLE", "STOPPING", "STOPPED"]
"""ExecutionWorker status."""


BUILD_PENDING = "PENDING"
BUILD_BUILDING = "BUILDING"
BUILD_READY = "READY"
BUILD_FAILED = "FAILED"
BuildStatus = Literal["PENDING", "BUILDING", "READY", "FAILED"]
"""Lifecycle of an on-demand image build.

A ``BuildTask`` row is created at claim time (``BUILDING``), before the build
starts — so concurrent tasks see the in-flight build and attach instead of
starting a second one. ``READY`` once the image exists; ``FAILED`` after a
build error (retried while ``attempts <= max_retries``)."""


PRESERVATION_NONE = "NONE"
PRESERVATION_FULL = "FULL"
PreservationMode = Literal["NONE", "FULL"]
"""Which tables survive after a job completes.

- ``NONE``: persistent tables only (default) — intermediate tables are
  dropped as soon as their refs fall to zero.
- ``FULL``: every table the job produced stays until the job TTL expires,
  useful for development and debugging.
"""


RUNNER_SUBPROCESS = "subprocess"
RUNNER_DOCKER = "docker"
RUNNER_KUBERNETES = "kubernetes"
RunnerMode = Literal["subprocess", "docker", "kubernetes"]
"""Which task-execution runner the orchestrator uses for a job.

- ``subprocess`` (default): each task runs in a multiprocessing child
  spawned by the host worker process.
- ``docker``: each task runs in a fresh container built on demand from
  the user's repo at a specific git SHA.
- ``kubernetes``: each task runs in a fresh Pod built on demand from the
  user's repo at a specific git SHA, scheduled on a cluster. The result
  is handed back via the ``remote_task_results`` table rather than a
  bind-mounted file.
"""
RUNNER_MODES: list[RunnerMode] = [RUNNER_SUBPROCESS, RUNNER_DOCKER, RUNNER_KUBERNETES]


class RegisteredJob(SQLModel, table=True):
    """
    RegisteredJob model - catalog of known jobs.

    Stores the job definition (entrypoint, schedule, defaults) separately
    from individual job runs. The background worker uses cron schedules
    to create Job rows automatically.
    """

    __tablename__: ClassVar[str] = "registered_jobs"
    __table_args__ = (UniqueConstraint("name"),)

    id: int = Field(sa_column=Column(BigInteger, primary_key=True))
    name: str = Field(index=True)
    entrypoint: str = Field()
    enabled: bool = Field(sa_column=Column(Boolean, nullable=False, server_default="1"), default=True)
    schedule: str | None = Field(default=None)
    default_kwargs: dict[str, Any] | None = Field(default=None, sa_column=Column(JSON, nullable=True))
    preservation_mode: PreservationMode | None = Field(
        default=None,
        sa_column=Column(String, nullable=True),
    )
    runner_mode: RunnerMode = Field(
        default=RUNNER_SUBPROCESS,
        sa_column=Column(String, nullable=False, server_default=RUNNER_SUBPROCESS),
    )
    dockerfile: str | None = Field(default=None)
    git_remote: str | None = Field(default=None)
    kubernetes_config: dict[str, Any] | None = Field(default=None, sa_column=Column(JSON, nullable=True))
    runner: dict[str, Any] | None = Field(default=None, sa_column=Column(JSON, nullable=True))
    next_run_at: datetime | None = Field(default=None, index=True)
    created_at: datetime = Field(default_factory=utc_now)
    updated_at: datetime = Field(default_factory=utc_now)


class Job(SQLModel, table=True):
    """
    Job model - represents a workflow execution.

    Each job has a unique snowflake ID and tracks overall workflow status.
    """

    __tablename__: ClassVar[str] = "jobs"

    id: int = Field(sa_column=Column(BigInteger, primary_key=True))
    name: str = Field(index=True)
    status: JobStatus = Field(
        default=JOB_PENDING,
        sa_column=Column(String, nullable=False, index=True),
    )
    run_type: RunType = Field(
        sa_column=Column(String, nullable=False),
    )
    registered_job_id: int | None = Field(
        default=None,
        sa_column=Column(BigInteger, ForeignKey("registered_jobs.id"), nullable=True, index=True),
    )
    preservation_mode: PreservationMode = Field(
        default=PRESERVATION_NONE,
        sa_column=Column(String, nullable=False, server_default=PRESERVATION_NONE),
    )
    runner_mode: RunnerMode = Field(
        default=RUNNER_SUBPROCESS,
        sa_column=Column(String, nullable=False, server_default=RUNNER_SUBPROCESS),
    )
    runner: dict[str, Any] | None = Field(default=None, sa_column=Column(JSON, nullable=True))
    created_at: datetime = Field(default_factory=utc_now, index=True)
    started_at: datetime | None = Field(default=None)
    completed_at: datetime | None = Field(default=None)
    error: str | None = Field(default=None)


class ExecutionWorker(SQLModel, table=True):
    """
    ExecutionWorker model - represents a worker process that executes tasks.

    Workers claim tasks from the queue and execute them.
    """

    __tablename__: ClassVar[str] = "execution_workers"

    id: int = Field(sa_column=Column(BigInteger, primary_key=True))
    hostname: str = Field(index=True)
    pid: int = Field()
    status: ExecutionWorkerStatus = Field(
        default=EXECUTION_WORKER_ACTIVE,
        sa_column=Column(String, nullable=False, index=True),
    )
    created_at: datetime = Field(default_factory=utc_now)
    started_at: datetime = Field(default_factory=utc_now)
    last_heartbeat: datetime = Field(default_factory=utc_now, index=True)
    tasks_completed: int = Field(default=0)
    tasks_failed: int = Field(default=0)


class Dependency(SQLModel, table=True):
    """
    Dependency model - represents dependencies between tasks and groups.

    Supports all combinations:
    - Task → Task
    - Task → Group
    - Group → Task
    - Group → Group
    """

    __tablename__: ClassVar[str] = "dependencies"

    # Entity that must complete first
    previous_id: int = Field(sa_column=Column(BigInteger, primary_key=True, index=True))
    previous_type: DependencyType = Field(sa_column=Column(String, primary_key=True))

    # Entity that waits (executes after previous completes)
    next_id: int = Field(sa_column=Column(BigInteger, primary_key=True, index=True))
    next_type: DependencyType = Field(sa_column=Column(String, primary_key=True))

    created_at: datetime = Field(default_factory=utc_now)


class _DependencyOps:
    """Shared dependency wiring for Task and Group.

    Provides ``depends_on`` plus the ``>>`` / ``<<`` operators. Subclasses set
    ``_dep_next_type`` and define their own ``previous_dependencies``
    relationship (the primaryjoin differs per table). ``_node()`` narrows
    ``self`` to ``Task | Group`` — a bare attribute annotation here would be
    collected as a pydantic field by the SQLModel metaclass.
    """

    _dep_next_type: ClassVar[DependencyType]

    def _node(self) -> Union["Task", "Group"]:
        return cast(Union["Task", "Group"], self)

    def depends_on(self, other: Union["Task", "Group"]) -> Union["Task", "Group"]:
        """
        Declare that this item depends on a task or group.

        Creates a Dependency record that will be committed when commit_tasks() is called.

        Args:
            other: Task or Group that must complete before this one

        Returns:
            self (for chaining)
        """
        node = self._node()
        dependency = Dependency(
            previous_id=other.id,
            previous_type=DEPENDENCY_TASK if isinstance(other, Task) else DEPENDENCY_GROUP,
            next_id=node.id,
            next_type=node._dep_next_type,
        )
        node.previous_dependencies.append(dependency)
        return node

    def __rshift__(
        self, other: Union["Task", "Group", Sequence[Union["Task", "Group"]]]
    ) -> Union["Task", "Group", Sequence[Union["Task", "Group"]]]:
        """A >> B: B depends on A (A executes before B)."""
        if isinstance(other, (Task, Group)):
            other.depends_on(self._node())
            return other
        for item in other:
            item.depends_on(self._node())
        return other

    def __lshift__(self, other: Union["Task", "Group", list[Union["Task", "Group"]]]) -> Union["Task", "Group"]:
        """A << B: A depends on B (B executes before A)."""
        if isinstance(other, list):
            for item in other:
                self.depends_on(item)
        else:
            self.depends_on(other)
        return self._node()

    def __rrshift__(self, other: Union["Task", "Group", Sequence[Union["Task", "Group"]]]) -> Union["Task", "Group"]:
        """Reverse: [A, B] >> C means C depends on A and B (fan-in)."""
        if isinstance(other, list):
            for item in other:
                self.depends_on(item)
        else:
            self.depends_on(other)
        return self._node()

    def __rlshift__(
        self, other: Union["Task", "Group", list[Union["Task", "Group"]]]
    ) -> Union["Task", "Group", list[Union["Task", "Group"]]]:
        """Reverse: [A, B] << C means A and B depend on C (fan-out)."""
        if isinstance(other, list):
            for item in other:
                item.depends_on(self._node())
            return other
        else:
            other.depends_on(self._node())
            return other


class Group(_DependencyOps, SQLModel, table=True):
    """
    Group model - represents a logical grouping of tasks.

    Groups support nesting via parent_group_id.
    Non-DB attribute ``_tasks`` carries associated Task objects so that
    returning a Group from a @task/@job function also registers its tasks.
    """

    __tablename__: ClassVar[str] = "groups"
    _dep_next_type: ClassVar[DependencyType] = DEPENDENCY_GROUP

    id: int = Field(sa_column=Column(BigInteger, primary_key=True))
    job_id: int = Field(default=0, sa_column=Column(BigInteger, ForeignKey("jobs.id"), index=True))
    parent_group_id: int | None = Field(
        default=None, sa_column=Column(BigInteger, ForeignKey("groups.id"), index=True, nullable=True)
    )
    name: str = Field()
    created_at: datetime = Field(default_factory=utc_now)

    _tasks: list = []
    _result_task: Any = None

    def model_post_init(self, __context: Any) -> None:
        self._tasks = []
        self._result_task = None
        register_task(self.id, self)

    def add_task(self, task: "Task") -> None:
        """Attach a Task to this group for co-registration."""
        self._tasks.append(task)

    def get_tasks(self) -> list["Task"]:
        """Return tasks attached to this group (non-DB)."""
        return self._tasks

    # Dependencies where this group is the "next" (i.e., this group depends on previous)
    # Note: overlaps="previous_dependencies" tells SQLAlchemy that both Task and Group
    # intentionally write to Dependency.next_id (polymorphic design via next_type)
    previous_dependencies: Mapped[list[Dependency]] = Relationship(
        sa_relationship_kwargs={
            "primaryjoin": "and_(Group.id == foreign(Dependency.next_id), Dependency.next_type == 'group')",
            "cascade": "all, delete-orphan",
            "overlaps": "previous_dependencies",
        }
    )


class Task(_DependencyOps, SQLModel, table=True):
    """
    Task model - represents a unit of work to be executed.

    Tasks can be part of a group and have dependencies on other tasks/groups.
    """

    __tablename__: ClassVar[str] = "tasks"
    _dep_next_type: ClassVar[DependencyType] = DEPENDENCY_TASK
    # Composite index: the rollup aggregate and the cascade UPDATE in
    # background.handler both filter on (job_id, status) — hot path on every
    # task transition.
    __table_args__ = (Index("ix_tasks_job_id_status", "job_id", "status"),)

    id: int = Field(sa_column=Column(BigInteger, primary_key=True))
    job_id: int = Field(default=0, sa_column=Column(BigInteger, ForeignKey("jobs.id"), index=True))
    group_id: int | None = Field(
        default=None, sa_column=Column(BigInteger, ForeignKey("groups.id"), index=True, nullable=True)
    )
    entrypoint: str = Field()
    name: str = Field()
    kwargs: dict[str, Any] = Field(default_factory=dict, sa_column=Column(JSON))
    entry_type: EntryType = Field(default=ENTRY_MODULE, sa_column=Column(String, nullable=True))
    command: list[str] | None = Field(default=None, sa_column=Column(JSON, nullable=True))
    command_env: dict[str, str] | None = Field(default=None, sa_column=Column(JSON, nullable=True))
    status: TaskStatus = Field(
        default=TASK_PENDING,
        sa_column=Column(String, nullable=False, index=True),
    )
    created_at: datetime = Field(default_factory=utc_now, index=True)
    claimed_at: datetime | None = Field(default=None)
    started_at: datetime | None = Field(default=None)
    completed_at: datetime | None = Field(default=None)
    execution_worker_id: int | None = Field(
        default=None, sa_column=Column(BigInteger, ForeignKey("execution_workers.id"), index=True, nullable=True)
    )
    build_task_id: int | None = Field(
        default=None, sa_column=Column(BigInteger, ForeignKey("build_tasks.id"), index=True, nullable=True)
    )
    result: dict[str, Any] | None = Field(default=None, sa_column=Column(JSON, nullable=True))
    log_path: str | None = Field(default=None)
    error: str | None = Field(default=None)
    max_retries: int = Field(default=0)
    attempt: int = Field(default=0)
    retry_after: datetime | None = Field(default=None)
    run_ids: list[int] = Field(default_factory=list, sa_column=Column(JSON, nullable=False, server_default="[]"))
    run_statuses: list[str] = Field(default_factory=list, sa_column=Column(JSON, nullable=False, server_default="[]"))
    # Fencing token bumped by clear_task. A worker captures this at claim time
    # and gates every status write on it, so a cleared run's late writes are
    # rejected — see clear_task / check_run_aborted in execution/claiming.py.
    run_epoch: int = Field(default=0, sa_column=Column(BigInteger, nullable=False, server_default="0"))

    def model_post_init(self, __context: Any) -> None:
        register_task(self.id, self)

    # Dependencies where this task is the "next" (i.e., this task depends on previous)
    # Note: overlaps="previous_dependencies" tells SQLAlchemy that both Task and Group
    # intentionally write to Dependency.next_id (polymorphic design via next_type)
    previous_dependencies: Mapped[list[Dependency]] = Relationship(
        sa_relationship_kwargs={
            "primaryjoin": "and_(Task.id == foreign(Dependency.next_id), Dependency.next_type == 'task')",
            "cascade": "all, delete-orphan",
            "overlaps": "previous_dependencies",
        }
    )


class RemoteTaskResult(SQLModel, table=True):
    """Per-attempt result handoff from a remote task vehicle (Kubernetes Pod)
    back to the host worker.

    The Pod writes one row keyed by ``(task_id, run_epoch)``; the host reads
    it at the epoch it claimed and then writes the terminal ``Task`` status
    itself. The Pod never writes ``Task.status``/``run_statuses`` — preserving
    the reaper invariant (see ``docker_worker`` module docstring). ``run_epoch``
    fences stale attempts: a cleared/re-run task bumps the live epoch, so an
    old Pod's row lands under a stale key and is ignored.
    """

    __tablename__: ClassVar[str] = "remote_task_results"

    task_id: int = Field(
        sa_column=Column(BigInteger, ForeignKey("tasks.id"), primary_key=True),
    )
    run_epoch: int = Field(sa_column=Column(BigInteger, primary_key=True))
    success: bool = Field(sa_column=Column(Boolean, nullable=False))
    result_ref: dict[str, Any] | None = Field(default=None, sa_column=Column(JSON, nullable=True))
    log_path: str | None = Field(default=None)
    error: str | None = Field(default=None)
    created_at: datetime = Field(default_factory=utc_now)


class BuildTask(SQLModel, table=True):
    """On-demand image build, keyed by image identity (not by job).

    ``image_key`` is a sha256 of ``(git_remote, git_sha, dockerfile)``; its
    UNIQUE constraint makes a duplicate build record structurally impossible, so
    every job/dynamic task on the same image shares this one row. ``ImagePrebuilt``
    images never create a ``BuildTask``."""

    __tablename__: ClassVar[str] = "build_tasks"
    __table_args__ = (UniqueConstraint("image_key"),)

    id: int = Field(sa_column=Column(BigInteger, primary_key=True))
    image_key: str = Field(sa_column=Column(String, nullable=False, index=True))
    image_tag: str = Field()
    git_remote: str = Field()
    git_sha: str = Field()
    dockerfile: str | None = Field(default=None)
    status: BuildStatus = Field(
        default=BUILD_PENDING,
        sa_column=Column(String, nullable=False, index=True),
    )
    holder_execution_worker_id: int | None = Field(default=None, sa_column=Column(BigInteger, nullable=True))
    lease_expires_at: datetime | None = Field(default=None)
    log_path: str | None = Field(default=None)
    error: str | None = Field(default=None)
    attempts: int = Field(default=0)
    max_retries: int = Field(default=2)
    created_at: datetime = Field(default_factory=utc_now)
    started_at: datetime | None = Field(default=None)
    finished_at: datetime | None = Field(default=None)


# Type alias for tasks/groups that can be applied
TasksType = Task | Group | list[Task | Group]
