# generated by datamodel-codegen:
#   filename:  application.json
#   timestamp: 2022-10-27T05:19:51+00:00

from __future__ import annotations

from enum import Enum
from typing import Dict, List, Optional, Union

from pydantic import BaseModel, Field, conint, constr


class DockerFileBuild(BaseModel):
    """
        +label=Docker File (I already have Docker File)
    +icon=fa-brands fa-docker:#0db7ed
    """

    type: constr(regex=r"dockerfile")
    dockerfile_path: str = Field(
        "./Dockerfile",
        description="+label=Path to Dockerfile\n+usage=The file path of the Dockerfile relative to project root path.",
    )
    build_context_path: str = Field(
        "./",
        description="+label=Path to build context\n+usage=Build context path for the Dockerfile relative to project root path.",
    )


class GitSource(BaseModel):
    """
        +label=Git Source
    +icon=fa-solid fa-code-branch:black
    """

    type: constr(regex=r"git")
    repo_url: constr(
        regex=r"(http(s?)://)(github.com|(.+@)*bitbucket.org|gitlab.com).*$"
    ) = Field(..., description="+label=Repo URL\n+usage=The repository URL.\n+sort=1")
    ref: str = Field(
        ...,
        description="+label=Branch/Tag/Commit SHA\n+usage=The commit sha, branch name or the tag.\n+sort=2",
    )
    branch_name: Optional[str] = Field(
        None,
        description="+label=Branch Name\n+usage=branch name of given commit sha in the `ref` field.\n+sort=3\n+ignore",
    )


class HttpProbe(BaseModel):
    """
    +label=Instructions for assessing container health by executing an HTTP GET request.
    """

    type: constr(regex=r"http")
    path: str = Field(
        ...,
        description="+usage=The endpoint, relative to the port, to which the HTTP GET request should be directed.",
    )
    port: conint(ge=0, le=65535) = Field(
        ...,
        description="+usage=The TCP socket within the container to which the HTTP GET request should be directed.",
    )
    host: Optional[str] = None
    scheme: str = "HTTP"


class Image(BaseModel):
    """
        +label=Docker Image (Deploy an existing image)
    +icon=fa-brands fa-docker:#0db7ed
    """

    type: constr(regex=r"image")
    docker_registry: Optional[str] = Field(
        None,
        description="+label=Docker Registry\n+usage=FQN of the container registry. If you can't find your registry here,\nadd it through the [Settings](/settings?tab=registry) page",
    )
    image_uri: str = Field(
        ...,
        description="+label=Image URI\n+usage=The image URI. Specify the name of the image and the tag.\nIf the image is in Dockerhub, you can skip registry-url (for e.g. `tensorflow/tensorflow`).\nYou can use an image from a private registry using Advanced fields\n+placeholder=registry-url/account/image:version (e.g. docker.io/tensorflow/tensorflow)",
    )


class Restart(str, Enum):
    """
    +usage=
    """

    Never = "Never"
    OnFailure = "OnFailure"
    Always = "Always"


class LocalSource(BaseModel):
    """
        +label=Local
    +icon=fa-folder:black
    """

    type: constr(regex=r"local")
    project_root_path: str = Field("./", description="+usage=Local project root path.")


class Manual(BaseModel):
    """
    +label= Manual - Trigger the job manually
    """

    type: constr(regex=r"manual")


class ModelFormat(str, Enum):
    """
        +usage=Format of the model package.
    +sort=3
    +label=Model package format
    """

    pytorch = "pytorch"
    tensorflow = "tensorflow"
    mlflow = "mlflow"


class Protocol(str, Enum):
    """
    +usage=Protocol for the port.
    """

    TCP = "TCP"
    UDP = "UDP"


class Port(BaseModel):
    protocol: Protocol = Field("TCP", description="+usage=Protocol for the port.")
    port: conint(ge=0, le=65535) = Field(
        80, description="+usage=Port number to expose."
    )
    expose: bool = Field(True, description="+usage=Expose the port")


class PythonBuild(BaseModel):
    """
        +label=Python Buildpack (I don't have Dockerfile)
    +icon=fa-brands fa-python:#306998
    """

    type: constr(regex=r"tfy-python-buildpack")
    python_version: constr(regex=r"^\d+(\.\d+){1,2}$") = Field(
        "3.9",
        description="+label=Python version\n+usage=Python version to run your application.",
    )
    build_context_path: str = Field(
        "./",
        description="+label=Path to build context\n+usage=Build path relative to project root path.",
    )
    requirements_path: Optional[str] = Field(
        None,
        description="+label=Path to requirements\n+usage=Path to `requirements.txt` relative to\n`Path to build context`",
    )
    pip_packages: Optional[List[str]] = Field(
        None,
        description="+label=Pip packages to install\n+usage=Define pip package requirements.",
    )
    command: Union[str, List[str]] = Field(
        ...,
        description="+label=Entrypoint override\n+usage=Override the command to run when the container starts.\nCommand will be set as the Entrypoint of the generated\nimage.",
    )


class RemoteSource(BaseModel):
    """
        +label=S3
    +icon=fa-brands fa-aws:black
    """

    type: constr(regex=r"remote")
    remote_uri: str = Field(..., description="+label=Remote URI")


class Resources(BaseModel):
    """
        +icon=fa-microchip
    +label=Resources
    """

    cpu_request: str = Field(
        "0.2",
        description="+label=CPU Request\n+sort=1\n+usage=Requested CPU which determines the minimum cost incurred. The CPU usage can exceed the requested\namount, but not the value specified in the limit. 1 CPU means 1 CPU core. Fractional CPU can be requested\nlike `0.5` or `0.05`",
    )
    cpu_limit: str = Field(
        "0.5",
        description="+label=CPU Limit\n+usage=CPU limit beyond which the usage cannot be exceeded. 1 CPU means 1 CPU core. Fractional CPU can be requested\nlike `0.5`. CPU limit should be >= cpu request.\n+sort=2",
    )
    memory_request: constr(regex=r"^[1-9]\d*$") = Field(
        "200",
        description="+label=Memory Request\n+usage=Requested memory which determines the minimum cost incurred. The unit of memory is in megabytes(MB).\nSo 1 means 1 MB and 2000 means 2GB.\n+sort=3",
    )
    memory_limit: constr(regex=r"^[1-9]\d*$") = Field(
        "500",
        description="+label=Memory Limit\n+usage=Memory limit after which the application will be killed with an OOM error. The unit of memory is\nin megabytes(MB). So 1 means 1 MB and 2000 means 2GB. MemoryLimit should be greater than memory request.\n+sort=4",
    )
    instance_family: Optional[List[str]] = Field(
        None,
        description="+label=Instance family\n+usage=Instance family of the underlying machine to use. Multiple instance families can be supplied.\nThe workload is guaranteed to be scheduled on one of them.",
    )


class ConcurrencyPolicy(str, Enum):
    """
        +usage=Choose whether to allow this job to run while another instance of the job is running, or to replace the currently running instance. Allow
    will enable multiple instances of this job to run. Forbid will keep the current instance of the job running and stop a new instance from being run.
    Replace will terminate any currently running instance of the job and start a new one.
    """

    Forbid = "Forbid"
    Allow = "Allow"
    Replace = "Replace"


class Schedule(BaseModel):
    """
    +label=Schedule - Run the job on a schedule
    """

    type: constr(regex=r"scheduled")
    schedule: str = Field(
        ...,
        description="+usage=Specify the schedule for this job to be run periodically in cron format.\n```\n* * * * *\n| | | | |\n| | | | |___ day of week (0-6) (Sunday is 0)\n| | | |_____ month (1-12)\n| | |_______ day of month (1-31)\n| |_________ hour (0-23)\n|___________ minute (0-59)\n```",
    )
    concurrency_policy: ConcurrencyPolicy = Field(
        "Forbid",
        description="+usage=Choose whether to allow this job to run while another instance of the job is running, or to replace the currently running instance. Allow\nwill enable multiple instances of this job to run. Forbid will keep the current instance of the job running and stop a new instance from being run.\nReplace will terminate any currently running instance of the job and start a new one.",
    )


class Build(BaseModel):
    """
        +label=Source Code (Build and deploy source code)
    +icon=fa-code
    """

    type: constr(regex=r"build")
    docker_registry: Optional[str] = Field(
        None,
        description="+label=Docker Registry\n+usage=FQN of the container registry. If you can't find your registry here,\nadd it through the [Settings](/settings?tab=registry) page",
    )
    build_source: Union[RemoteSource, GitSource, LocalSource] = Field(
        ...,
        description="+label=Fetch source code to build and deploy\n+icon=fa-code\n+sort=1",
    )
    build_spec: Union[DockerFileBuild, PythonBuild] = Field(
        ...,
        description="+label=Build using DockerFile or using Buildpack\n+icon=fa-wrench\n+sort=2",
    )


class HealthProbe(BaseModel):
    """
    +icon=fa-heart
    """

    config: HttpProbe
    initial_delay_seconds: conint(ge=0, le=36000) = Field(
        0,
        description="+usage=Number of seconds after the container is started before the first probe is initiated.",
    )
    period_seconds: conint(ge=0, le=36000) = Field(
        10, description="+usage=How often, in seconds, to execute the probe."
    )
    timeout_seconds: conint(ge=0, le=36000) = Field(
        1, description="+usage=Number of seconds after which the probe times out."
    )
    success_threshold: conint(ge=0, le=100) = Field(
        1,
        description="+usage=Minimum consecutive successes for the probe to be considered successful after having failed.",
    )
    failure_threshold: conint(ge=0, le=100) = Field(
        3,
        description="+usage=Number of consecutive failures required to determine the container is not alive (liveness probe) or not ready (readiness probe).",
    )


class Job(BaseModel):
    type: constr(regex=r"job")
    name: constr(regex=r"^[a-z][a-z0-9-]{1,23}[a-z0-9]$") = Field(
        ..., description="+usage=Name of the job\n+sort=1"
    )
    image: Union[Build, Image] = Field(
        ...,
        description="+label=Deploy a Docker image or build and deploy from source code\n+icon=fa-solid fa-cloud-arrow-up:#21B6A8\n+sort=2",
    )
    resources: Optional[Resources] = None
    trigger: Union[Manual, Schedule] = Field({"type": "manual"}, description="+sort=3")
    retries: conint(ge=0) = Field(
        1,
        description="+label=Retries\n+usage=Specify the maximum number of attempts to retry a job before it is marked as failed.\n+icon=fa-clone",
    )
    timeout: conint(ge=0, le=3600) = Field(
        1000,
        description="+label=Timeout\n+usage=Job timeout in seconds.\n+icon=fa-clone",
    )
    successful_jobs_history_limit: conint(ge=0, le=100) = Field(
        20,
        description="+usage=Number of successful invocation of this job to keep in history",
    )
    failed_jobs_history_limit: conint(ge=0, le=100) = Field(
        20,
        description="+usage=Number of failed invocation of this job to keep in history",
    )
    restart: Restart = Field("Never", description="+usage=")
    command: Optional[Union[str, List[str]]] = Field(
        None,
        description="+label=Commands\n+usage=Commands to execute to start the application. Overrides the `Entrypoint` command.\n+icon=fa-terminal\n+sort=6",
    )
    env: Optional[Dict[str, str]] = Field(
        None,
        description="+label=Environment Variables\n+usage=Configure environment variables to be injected in the service.\n+icon=fa-globe\n+sort=7",
    )
    service_account: Optional[str] = Field(
        None, description="+usage=Service account that this workload should use"
    )


class ModelDeployment(BaseModel):
    type: constr(regex=r"model-deployment")
    name: constr(regex=r"^[a-z][a-z0-9-]{1,23}[a-z0-9]$") = Field(
        ...,
        description="+usage=Name of the inference service. This uniquely identifies this service in the workspace.\n> Name can only contain alphanumeric characters and '-' and can be atmost 20 characters long\n+sort=1\n+label=Inference service name",
    )
    model_uri: constr(regex=r"\w+:(/?/?)[^\s]+") = Field(
        ...,
        description="+usage=URI of the model to be deployed.\n+sort=2\n+label=Model URI",
    )
    model_format: Optional[ModelFormat] = Field(
        None,
        description="+usage=Format of the model package.\n+sort=3\n+label=Model package format",
    )
    resources: Optional[Resources] = None


class Notebook(BaseModel):
    type: constr(regex=r"notebook")
    name: constr(regex=r"^[a-z][a-z0-9-]{1,23}[a-z0-9]$") = Field(
        ..., description="+usage=Name of the notebook"
    )
    resources: Optional[Resources] = None


class Service(BaseModel):
    type: constr(regex=r"service")
    name: constr(regex=r"^[a-z][a-z0-9-]{1,23}[a-z0-9]$") = Field(
        ...,
        description="+usage=Name of the service. This uniquely identifies this service in the workspace.\n> Name can only contain alphanumeric characters and '-' and can be atmost 25 characters long\n+sort=1",
    )
    image: Union[Build, Image] = Field(
        ...,
        description="+label=Deploy a Docker image or build and deploy from source code\n+icon=fa-solid fa-cloud-arrow-up:#21B6A8\n+sort=2",
    )
    resources: Optional[Resources] = None
    replicas: conint(ge=1, le=5) = Field(
        1,
        description="+label=Replicas\n+usage=Number of service instances/replicas you want to run.\n+icon=fa-clone\n+sort=3",
    )
    command: Optional[Union[str, List[str]]] = Field(
        None,
        description="+label=Commands\n+usage=Commands to execute to start the application. Overrides the `Entrypoint` command.\n+icon=fa-terminal",
    )
    env: Optional[Dict[str, str]] = Field(
        None,
        description="+label=Environment Variables\n+usage=Configure environment variables to be injected in the service.\n+icon=fa-globe\n+sort=6",
    )
    ports: List[Port] = Field(
        ...,
        description="+label=Ports (Ports to route customer traffic to)\n+icon=fa-plug\n+sort=4",
    )
    liveness_probe: Optional[HealthProbe] = None
    readiness_probe: Optional[HealthProbe] = None
    service_account: Optional[str] = Field(
        None, description="+usage=Service account that this workload should use"
    )


class Application(BaseModel):
    name: constr(regex=r"^[a-z][a-z0-9-]{1,23}[a-z0-9]$") = Field(
        ..., description="+label=Name of the application"
    )
    components: List[Union[Service, Job, Notebook, ModelDeployment]] = Field(
        ..., description="+label=A list of components"
    )
