"""Pydantic models for GA4GH TES (Task Execution Service) v1.1."""

from __future__ import annotations

from enum import Enum
from typing import Any

from ai4gh.core.types import GA4GHBase


class TESState(str, Enum):
    """TES task state."""

    UNKNOWN = "UNKNOWN"
    QUEUED = "QUEUED"
    INITIALIZING = "INITIALIZING"
    RUNNING = "RUNNING"
    PAUSED = "PAUSED"
    COMPLETE = "COMPLETE"
    EXECUTOR_ERROR = "EXECUTOR_ERROR"
    SYSTEM_ERROR = "SYSTEM_ERROR"
    CANCELED = "CANCELED"
    CANCELING = "CANCELING"
    PREEMPTED = "PREEMPTED"


class TESView(str, Enum):
    """TES task view level."""

    MINIMAL = "MINIMAL"
    BASIC = "BASIC"
    FULL = "FULL"


class TESInput(GA4GHBase):
    """Input file for a TES task."""

    name: str | None = None
    description: str | None = None
    url: str | None = None
    path: str
    type: str | None = None
    content: str | None = None
    streamable: bool | None = None


class TESOutput(GA4GHBase):
    """Output file for a TES task."""

    name: str | None = None
    description: str | None = None
    url: str
    path: str
    type: str | None = None


class TESExecutor(GA4GHBase):
    """An executor (command) in a TES task."""

    image: str
    command: list[str]
    workdir: str | None = None
    stdin: str | None = None
    stdout: str | None = None
    stderr: str | None = None
    env: dict[str, str] | None = None


class TESResources(GA4GHBase):
    """Resource requirements for a TES task."""

    cpu_cores: int | None = None
    ram_gb: float | None = None
    disk_gb: float | None = None
    preemptible: bool | None = None
    zones: list[str] | None = None
    backend_parameters: dict[str, str] | None = None
    backend_parameters_strict: bool | None = None


class TESTaskLog(GA4GHBase):
    """Log for a TES task execution."""

    logs: list[dict[str, Any]] | None = None
    metadata: dict[str, str] | None = None
    start_time: str | None = None
    end_time: str | None = None
    outputs: list[dict[str, Any]] | None = None
    system_logs: list[str] | None = None


class TESTask(GA4GHBase):
    """A TES task."""

    id: str | None = None
    state: TESState | None = None
    name: str | None = None
    description: str | None = None
    inputs: list[TESInput] | None = None
    outputs: list[TESOutput] | None = None
    resources: TESResources | None = None
    executors: list[TESExecutor] | None = None
    volumes: list[str] | None = None
    tags: dict[str, str] | None = None
    logs: list[TESTaskLog] | None = None
    creation_time: str | None = None


class TESCreateTaskResponse(GA4GHBase):
    """Response when a task is created."""

    id: str


class TESListTasksResponse(GA4GHBase):
    """TES task list response."""

    tasks: list[TESTask] | None = None
    next_page_token: str | None = None
