import os
import tempfile
import typing as t
from pathlib import Path

from filelock import FileLock
from packaging.version import Version
from pydantic import AnyHttpUrl, BaseModel, Field, create_model, validator

from tungstenkit._internal.constants import LOCK_DIR, PKG_DIR
from tungstenkit._internal.io_schema import IOClasses
from tungstenkit.exceptions import NotLoggedIn

CONFIG_PATH = PKG_DIR / "config.json"
CONFIG_LOCK_PATH = LOCK_DIR / "config.json.lock"


class _Version(Version):
    validate_always = True

    @classmethod
    def __get_validators__(cls):
        yield cls.validate

    @classmethod
    def validate(cls, v: t.Any):
        if isinstance(v, Version):
            return v
        if not isinstance(v, str):
            raise ValueError(f"Should be 'str', not '{type(v)}'.")
        return cls(v)


def default_include_files():
    return ["*"]


def default_exclude_files():
    return [".*/", "__pycache__/", "*.pyc", "*.pyo", "*.pyd"]


# TODO print warnings if there are ignored fields
class BuildConfig(BaseModel):
    gpu: bool = False
    mem_gb: int = 8
    gpu_mem_gb: int = 16

    description: str = Field(default_factory=str)
    python_packages: t.List[str] = Field(default_factory=list)
    include_files: t.List[str] = Field(default_factory=default_include_files)
    copy_files: t.List[t.Tuple[str, str]] = Field(default_factory=list)
    exclude_files: t.List[str] = Field(default_factory=default_exclude_files)
    pip_wheels: t.List[Path] = Field(default_factory=list)
    conda_environment_yml: t.Optional[Path] = None
    environment_variables: t.Dict[str, str] = Field(default_factory=dict)

    system_packages: t.List[str] = Field(default_factory=list)
    dockerfile_commands: t.List[str] = Field(default_factory=list)
    python_version: t.Optional[_Version] = None
    cuda_version: t.Optional[_Version] = None

    base_image: t.Optional[str] = None

    class Config:
        arbitrary_types_allowed = True

    @validator("cuda_version")
    def validate_cuda_version(cls, v, values, **kwargs):
        if not values["gpu"] and v is not None:
            raise ValueError("'gpu' is not set")
        return v


class ModelConfig(BuildConfig):
    batch_size: int = 1
    readme_md: t.Optional[Path] = None
    input_schema: t.Dict = Field(default_factory=dict)
    output_schema: t.Dict = Field(default_factory=dict)

    @classmethod
    def with_types(cls, io_classes: IOClasses):
        return create_model(
            cls.__name__,
            __base__=cls,
            input_schema=(t.Dict, Field(default=io_classes.input.schema())),
            output_schema=(t.Dict, Field(default=io_classes.output.schema())),
        )

    @validator("readme_md")
    def validate_readme(cls, v):
        if v is not None:
            readme_path = Path(v).resolve()
            if not readme_path.exists():
                raise ValueError(f"Not found: {v}")
            if not readme_path.is_file():
                raise ValueError(f"Not a file: {v}")
            return readme_path


class TungstenClientConfig(BaseModel):
    url: AnyHttpUrl
    access_token: str

    def save(self, path: t.Optional[Path] = None, lock_path: t.Optional[Path] = None):
        path = path if path else CONFIG_PATH
        lock_path = lock_path if lock_path else CONFIG_LOCK_PATH

        dumped = self.json()
        tmp_config_fd, tmp_config_path_str = tempfile.mkstemp(suffix=path.name)
        try:
            tmp_config_path = Path(tmp_config_path_str)
            with FileLock(lock_path, timeout=30.0):
                tmp_config_path.write_text(dumped)
                os.replace(tmp_config_path, path)
        finally:
            os.close(tmp_config_fd)

    @staticmethod
    def from_env(path: t.Optional[Path] = None) -> "TungstenClientConfig":
        path = path if path else CONFIG_PATH
        if not path.exists():
            raise NotLoggedIn("Please log in first")
        # TODO catch a pydantic exception and convert to a tungsten exception.
        return TungstenClientConfig.parse_raw(path.read_text())
