from pydantic import BaseModel, Field, validator
import datetime
from typing import List, Optional, Union, Literal, Dict
from enum import Enum
import semver
import uuid
from arkitekt_next.model import Requirement
from rekuest.api.schema import DefinitionInput
from string import Formatter
import os

from rekuest_next.api.schema import TemplateInput

ALLOWED_BUILDER_KEYS = [
    "tag",
    "dockerfile",
]


class Framework(str, Enum):
    """Do we support other frameworks?"""

    VANILLA = "vanilla"
    TENSORFLOW = "tensorflow"
    PYTORCH = "pytorch"


class Packager(str, Enum):
    CONDA = "conda"
    POETRY = "poetry"
    PIP = "pip"


class Manifest(BaseModel):
    identifier: str
    version: str
    author: str
    logo: Optional[str]
    entrypoint: str
    scopes: List[str]
    created_at: datetime.datetime = Field(default_factory=datetime.datetime.now)

    @validator("version", pre=True)
    def version_must_be_semver(cls, v) -> str:
        """Checks that the version is a valid semver version"""
        if isinstance(v, str):
            try:
                semver.VersionInfo.parse(v)
            except ValueError:
                raise ValueError("Version must be a valid semver version")
        return str(v)

    def to_console_string(self):
        return f"📦 {self.identifier} ({self.version}) by {self.author}"

    def to_builder_dict(self):
        return {
            "identifier": self.identifier,
            "version": self.version,
            "logo": self.logo,
            "scopes": self.scopes,
        }

    class Config:
        validate_assignment = True


class SelectorType(str, Enum):
    RAM = "ram"
    CPU = "cpu"
    GPU = "gpu"
    LABEL = "label"


class BaseSelector(BaseModel):
    """A selector is a way to describe a flavours preference for a
    compute node. It contains the node_id, the selector and the flavour_id.
    """

    required: bool = True

    class Config:
        extra = "forbid"

    def build_docker_params(self) -> List[str]:
        """Builds the docker params for this selector

        Should return a list of strings that can be used as docker params
        If the selector is not required, it should return an empty list

        Returns
        -------
        List[str]
            The docker params for this selector
        """
        return []

    def build_arkitekt_next_params(self) -> List[str]:
        """Builds the arkitekt_next params for this selector

        Returns
        -------
        List[str]
            The docker params for this selector
        """
        return []


class RAMSelector(BaseSelector):
    """A selector is a way to describe a flavours preference for a
    compute node. It contains the node_id, the selector and the flavour_id.
    """

    type: Literal["ram"]
    min: int


class CPUSelector(BaseSelector):
    """A selector is a way to describe a flavours preference for a
    compute node. It contains the node_id, the selector and the flavour_id.
    """

    type: Literal["cpu"]
    min: int
    frequency: Optional[int] = None


class CudaSelector(BaseSelector):
    """A selector is a way to describe a flavours preference for a
    compute node. It contains the node_id, the selector and the flavour_id.
    """

    type: Literal["cuda"]
    frequency: Optional[int] = Field(default=None, description="The frequency in MHz")
    memory: Optional[int] = Field(default=None, description="The memory in MB")
    architecture: Optional[str] = Field(
        default=None, description="The architecture of the GPU"
    )
    compute_capability: str = Field(
        default="3.5", description="The minimum compute capability"
    )
    cuda_cores: Optional[int] = None
    cuda_version: str = Field(default="10.2", description="The minimum cuda version")

    def build_docker_params(self) -> List[str]:
        """Builds the docker params for this selector

        Should return a list of strings that can be used as docker params
        If the selector is not required, it should return an empty list

        Returns
        -------
        List[str]
            The docker params for this selector
        """
        return [
            "--gpus",
            "all",
        ]


class RocmSelector(BaseSelector):
    """A selector is a way to describe a flavours preference for a
    compute node. It contains the node_id, the selector and the flavour_id.
    """

    type: Literal["rocm"]
    min: int
    frequency: Optional[int] = None
    memory: Optional[int] = None
    architecture: Optional[str] = None
    compute_capability: Optional[str] = None
    cuda_cores: Optional[int] = None
    cuda_version: Optional[str] = None

    def build_docker_params(self) -> List[str]:
        """Builds the docker params for this selector"""

        return [
            "--device=/dev/kfd",
            "--device=/dev/dri",
            "--group-add",
            "video",
            "--cap-add=SYS_PTRACE",
            "--security-opt",
            "seccomp=unconfined",
        ]


class LabelSelector(BaseSelector):
    """A selector is a way to describe a flavours preference for a
    compute node. It contains the node_id, the selector and the flavour_id.
    """

    type: Literal["label"]
    key: str
    value: str


class ServiceSelector(BaseSelector):
    """A service selector is a way to describe a flavours preference for a
    service. It contains the service_id,
    """

    type: Literal["service"]


Selector = Union[
    RAMSelector, CPUSelector, CudaSelector, RocmSelector, LabelSelector, ServiceSelector
]


class Inspection(BaseModel):
    templates: List[TemplateInput]
    requirements: Dict[str, Requirement]
    size: int


class Flavour(BaseModel):
    selectors: List[Selector]
    description: str = Field(default="")
    dockerfile: str = Field(default="Dockerfile")
    build_command: List[str] = Field(
        default_factory=lambda: [
            "docker",
            "build",
            "-t",
            "{tag}",
            "-f",
            "{dockerfile}",
            ".",
        ]
    )

    @validator("build_command", each_item=True, always=True)
    def check_valid_template_name(cls, v):
        """Checks that the template name is valid"""
        for literal_text, field_name, format_spec, conversion in Formatter().parse(v):
            if field_name is not None:
                assert (
                    field_name in ALLOWED_BUILDER_KEYS
                ), f"Invalid template key {field_name}. Allowed keys are {ALLOWED_BUILDER_KEYS}"

        return v

    def generate_build_command(self, tag: str, relative_dir: str):
        """Generates the build command for this flavour"""

        dockerfile = os.path.join(relative_dir, self.dockerfile)

        return [v.format(tag=tag, dockerfile=dockerfile) for v in self.build_command]

    def check_relative_paths(self, flavour_folder: str):
        """Checks that the paths are relative to the flavour folder"""

        dockerfile_path = os.path.join(flavour_folder, self.dockerfile)

        if not os.path.exists(dockerfile_path):
            raise Exception(
                f"Could not find Dockerfile {self.dockerfile} in flavour {flavour_folder}"
            )


class Deployment(BaseModel):
    """A deployment is a Release of a Build.
    It contains the build_id, the manifest, the builder, the definitions, the image and the deployed_at timestamp.



    """

    deployment_run: str = Field(
        default_factory=lambda: str(uuid.uuid4()),
        description="The unique identifier of the deployment run",
    )

    deployment_id: str = Field(
        default_factory=lambda: str(uuid.uuid4()),
        description="The unique identifier of the deployment",
    )
    manifest: Manifest = Field(description="The manifest of the app that was deployed")
    selectors: List[Selector] = Field(
        description="The selectors are used to place this image on the nodes",
        default_factory=list,
    )
    flavour: str = Field(
        description="The flavour that was used to build this deployment",
        default="vanilla",
    )
    build_id: str = Field(
        description="The build_id of the build that was deployed. Is referenced in the build.yaml file."
    )
    inspection: Optional[Inspection] = Field(
        description="The inspection of the build that was deployed"
    )
    image: str = Field(
        description="The docker image that was built for this deployment"
    )
    deployed_at: datetime.datetime = Field(
        default_factory=datetime.datetime.now,
        description="The timestamp of the deployment",
    )


class DeploymentsConfigFile(BaseModel):
    """The ConfigFile is a pydantic model that represents the deployments.yaml file


    Parameters
    ----------
    BaseModel : _type_
        _description_
    """

    deployments: List[Deployment] = []
    latest_deployment_run: Optional[str] = None


class Build(BaseModel):
    build_run: str
    build_id: str
    inspection: Optional[Inspection] = None
    description: str = Field(default="")
    selectors: List[Selector] = Field(default_factory=list)
    flavour: str = Field(default="vanilla")
    manifest: Manifest
    build_at: datetime.datetime = Field(default_factory=datetime.datetime.now)
    base_docker_command: List[str] = Field(
        default_factory=lambda: ["docker", "run", "-it", "--net", "host"]
    )
    base_arkitekt_next_command: List[str] = Field(
        default_factory=lambda: ["arkitekt-next", "run", "prod", "--headless"]
    )

    def build_docker_command(self) -> List[str]:
        """Builds the docker command for this build"""

        base_command = self.base_docker_command

        for selector in self.selectors:
            base_command = base_command + selector.build_docker_params()

        base_command = base_command + [self.build_id]

        return base_command

    def build_arkitekt_next_command(self, fakts_url: str):
        """Builds the arkitekt_next command for this build"""

        base_command = self.base_arkitekt_next_command

        for selector in self.selectors:
            base_command = base_command + selector.build_arkitekt_next_params()

        base_command = base_command + ["--url", fakts_url]

        return base_command


class BuildsConfigFile(BaseModel):
    builds: List[Build] = Field(default_factory=list)
    latest_build_run: Optional[str] = None
