from __future__ import annotations

import enum
import uuid
from collections.abc import Mapping
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final

import sqlalchemy as sa
from sqlalchemy.dialects import postgresql as pgsql
from sqlalchemy.exc import NoResultFound
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import Mapped, foreign, mapped_column, relationship, selectinload

from ai.backend.manager.models.base import (
    GUID,
    Base,
    SlugType,
)

if TYPE_CHECKING:
    from ai.backend.manager.models.domain import DomainRow
    from ai.backend.manager.models.group import GroupRow

__all__: Final[tuple[str, ...]] = (
    "NetworkRow",
    "NetworkType",
)


class NetworkType(enum.StrEnum):
    VOLATILE = "volatile"
    PERSISTENT = "persistent"
    HOST = "host"


def _get_project_join_condition() -> sa.ColumnElement[bool]:
    from ai.backend.manager.models.group import GroupRow

    return GroupRow.id == foreign(NetworkRow.project)


def _get_domain_join_condition() -> sa.ColumnElement[bool]:
    from ai.backend.manager.models.domain import DomainRow

    return DomainRow.name == foreign(NetworkRow.domain_name)


class NetworkRow(Base):  # type: ignore[misc]
    __tablename__ = "networks"

    id: Mapped[uuid.UUID] = mapped_column(
        "id", GUID, primary_key=True, server_default=sa.text("uuid_generate_v4()")
    )

    name: Mapped[str] = mapped_column("name", sa.String(length=128), nullable=False)
    """Name of the network"""
    ref_name: Mapped[str] = mapped_column("ref_name", sa.String(length=128), nullable=False)
    """Network ID generated by network plugin"""
    driver: Mapped[str] = mapped_column("driver", sa.String(length=64), nullable=False)
    """Network plugin in charge of the network created"""
    options: Mapped[Mapping[str, Any]] = mapped_column(
        "options", pgsql.JSONB, nullable=False, default="{}", server_default="{}"
    )

    project: Mapped[uuid.UUID] = mapped_column(
        "project",
        GUID,
        nullable=False,
    )
    domain_name: Mapped[str] = mapped_column(
        "domain_name",
        SlugType(length=64, allow_unicode=True, allow_dot=True),
        nullable=False,
    )

    created_at: Mapped[datetime | None] = mapped_column(
        "created_at",
        sa.DateTime(timezone=True),
        server_default=sa.text("now()"),
        nullable=True,
    )
    updated_at: Mapped[datetime | None] = mapped_column(
        "updated_at",
        sa.DateTime(timezone=True),
        server_default=sa.text("now()"),
        nullable=True,
    )

    project_row: Mapped[GroupRow] = relationship(
        "GroupRow",
        back_populates="networks",
        primaryjoin=_get_project_join_condition,
    )
    domain_row: Mapped[DomainRow] = relationship(
        "DomainRow",
        back_populates="networks",
        primaryjoin=_get_domain_join_condition,
    )

    def __init__(
        self,
        name: str,
        ref_name: str,
        driver: str,
        domain: str,
        project: uuid.UUID,
        *,
        options: Mapping[str, Any] | None = None,
    ) -> None:
        self.id = uuid.uuid4()
        self.name = name
        self.ref_name = ref_name
        self.driver = driver
        self.domain_name = domain
        self.project = project
        self.options = options or {}

    @classmethod
    async def get(
        cls,
        session: AsyncSession,
        network_id: uuid.UUID,
        load_project: bool = False,
        load_domain: bool = False,
    ) -> NetworkRow:
        query = sa.select(NetworkRow).filter(NetworkRow.id == network_id)
        if load_project:
            query = query.options(selectinload(cls.project_row))
        if load_domain:
            query = query.options(selectinload(cls.domain_row))
        row = await session.scalar(query)
        if not row:
            raise NoResultFound

        return row
