from datetime import datetime
from typing import Any
from typing import Self

from sqlalchemy import Column
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.types import DateTime
from sqlmodel import CheckConstraint
from sqlmodel import Field
from sqlmodel import SQLModel

from fractal_server.utils import get_timestamp


class Resource(SQLModel, table=True):
    """
    Resource table.
    """

    id: int | None = Field(default=None, primary_key=True)

    type: str
    """
    One of `local`, `slurm_sudo` or `slurm_ssh` - matching with
    `settings.FRACTAL_RUNNER_BACKEND`.
    """

    name: str = Field(unique=True)
    """
    Resource name.
    """

    timestamp_created: datetime = Field(
        default_factory=get_timestamp,
        sa_column=Column(DateTime(timezone=True), nullable=False),
    )
    """
    Creation timestamp (autogenerated).
    """

    host: str | None = None
    """
    Address for ssh connections, when `type="slurm_ssh"`.
    """

    jobs_local_dir: str
    """
    Base local folder for job subfolders (containing artifacts and logs).
    """

    jobs_runner_config: dict[str, Any] = Field(
        sa_column=Column(JSONB, nullable=False, server_default="{}")
    )
    """
    Runner configuration, matching one of `JobRunnerConfigLocal` or
    `JobRunnerConfigSLURM` schemas.
    """

    jobs_slurm_python_worker: str | None = None
    """
    On SLURM deloyments, this is the Python interpreter that runs the
    `fractal-server` worker from within the SLURM jobs.
    """

    jobs_poll_interval: int
    """
    On SLURM resources: the interval to wait before new `squeue` calls.
    On local resources: ignored.
    """

    # task_settings
    tasks_local_dir: str
    """
    Base local folder for task-package subfolders.
    """

    tasks_python_config: dict[str, Any] = Field(
        sa_column=Column(JSONB, nullable=False, server_default="{}")
    )
    """
    Python configuration for task collection. Example:
    ```json
    {
      "default_version": "3.10",
      "versions:{
        "3.10": "/xxx/venv-3.10/bin/python",
        "3.11": "/xxx/venv-3.11/bin/python",
        "3.12": "/xxx/venv-3.12/bin/python"
       }
    }
    ```
    """

    tasks_pixi_config: dict[str, Any] = Field(
        sa_column=Column(JSONB, nullable=False, server_default="{}")
    )
    """
    Pixi configuration for task collection. Basic example:
    ```json
    {
        "default_version": "0.41.0",
        "versions": {
            "0.40.0": "/xxx/pixi/0.40.0/",
            "0.41.0": "/xxx/pixi/0.41.0/"
        },
    }
    ```
    """

    @property
    def pip_cache_dir_arg(self: Self) -> str:
        """
        If `pip_cache_dir` is set (in `self.tasks_python_config`), then
        return `--cache_dir /something`; else return `--no-cache-dir`.
        """
        _pip_cache_dir = self.tasks_python_config.get("pip_cache_dir", None)
        if _pip_cache_dir is not None:
            return f"--cache-dir {_pip_cache_dir}"
        else:
            return "--no-cache-dir"

    # Check constraints
    __table_args__ = (
        # `type` column must be one of "local", "slurm_sudo" or "slurm_ssh"
        CheckConstraint(
            "type IN ('local', 'slurm_sudo', 'slurm_ssh')",
            name="correct_type",
        ),
        # If `type` is not "local", `jobs_slurm_python_worker` must be set
        CheckConstraint(
            "(type = 'local') OR (jobs_slurm_python_worker IS NOT NULL)",
            name="jobs_slurm_python_worker_set",
        ),
    )
