from __future__ import annotations

import os
import platform
from typing import Optional, Tuple

from ramalama.common import MNT_DIR, RAG_DIR, ContainerEntryPoint, check_nvidia, get_accel_env_vars, get_gpu_devices
from ramalama.file import PlainFile
from ramalama.path_utils import normalize_host_path_for_container
from ramalama.version import version


class Kube:
    def __init__(
        self,
        model_name: str,
        model_paths: Tuple[str, str],
        chat_template_paths: Optional[Tuple[str, str]],
        mmproj_paths: Optional[Tuple[str, str]],
        args,
        exec_args,
        draft_model_paths: Optional[Tuple[str, str]],
        artifact: bool,
    ):
        self.src_chat_template_path, self.dest_chat_template_path = (
            chat_template_paths if chat_template_paths is not None else ("", "")
        )
        self.src_mmproj_path, self.dest_mmproj_path = mmproj_paths if mmproj_paths is not None else ("", "")
        self.model_paths = {'model': model_paths}
        if draft_model_paths is not None:
            self.model_paths['model-draft'] = draft_model_paths
        self.ai_image = model_name
        if getattr(args, "name", None):
            self.name = args.name
        else:
            self.name = "ramalama"

        self.args = args
        self.exec_args = exec_args
        self.image = args.image
        self.artifact = artifact

    def _gen_volumes(self) -> Tuple[str, str]:
        mounts = """\
        volumeMounts:"""

        volumes = """\
      volumes:"""
        for volume_name, (src_model_path, dest_model_path) in self.model_paths.items():
            src_model_path = src_model_path.removeprefix("oci://")
            if os.path.exists(src_model_path):
                m, v = self._gen_path_volume(volume_name, src_model_path, dest_model_path)
                mounts += m
                volumes += v
            else:
                subPath = ""
                if not self.artifact:
                    subPath = """
          subPath: /models"""
                mounts += f"""
        - mountPath: {MNT_DIR}{subPath}
          name: {volume_name}"""
                volumes += self._gen_oci_volume(volume_name, src_model_path)

        if getattr(self.args, 'rag', None):
            m, v = self._gen_rag_volume()
            mounts += m
            volumes += v

        if self.src_chat_template_path and os.path.exists(self.src_chat_template_path):
            m, v = self._gen_chat_template_volume()
            mounts += m
            volumes += v

        if self.src_mmproj_path and os.path.exists(self.src_mmproj_path):
            m, v = self._gen_mmproj_volume()
            mounts += m
            volumes += v

        m, v = self._gen_devices()
        mounts += m
        volumes += v
        return mounts, volumes

    def _gen_devices(self) -> Tuple[str, str]:
        mounts = ""
        volumes = ""
        for name, path in get_gpu_devices().items():
            mounts += f"""
        - mountPath: {path}
          name: {name}"""
            volumes += f"""
      - hostPath:
          path: {path}
        name: {name}"""
        return mounts, volumes

    def _gen_path_volume(self, volume_name, src_model_path, dest_model_path) -> Tuple[str, str]:
        host_model_path = normalize_host_path_for_container(src_model_path)
        if platform.system() == "Windows":
            #  Workaround https://github.com/containers/podman/issues/16704
            host_model_path = '/mnt' + host_model_path
        mount = f"""
        - mountPath: {dest_model_path}
          name: {volume_name}"""
        volume = f"""
      - hostPath:
          path: {host_model_path}
        name: {volume_name}"""
        return mount, volume

    def _gen_oci_volume(self, volume_name, src_model_path) -> str:
        return f"""
      - image:
          reference: {src_model_path}
          pullPolicy: IfNotPresent
        name: {volume_name}"""

    def _gen_rag_volume(self) -> Tuple[str, str]:
        mounts = f"""
        - mountPath: {RAG_DIR}
          name: rag"""

        volumes = f"""
      - image:
          reference: {self.args.rag}
          pullPolicy: IfNotPresent
        name: rag"""

        return mounts, volumes

    def _gen_chat_template_volume(self) -> Tuple[str, str]:
        host_chat_template_path = normalize_host_path_for_container(self.src_chat_template_path)
        if platform.system() == "Windows":
            #  Workaround https://github.com/containers/podman/issues/16704
            host_chat_template_path = '/mnt' + host_chat_template_path
        mount = f"""
        - mountPath: {self.dest_chat_template_path}
          name: chat-template"""
        volume = f"""
      - hostPath:
          path: {host_chat_template_path}
        name: chat-template"""
        return mount, volume

    def _gen_mmproj_volume(self) -> Tuple[str, str]:
        host_mmproj_path = normalize_host_path_for_container(self.src_mmproj_path)
        if platform.system() == "Windows":
            #  Workaround https://github.com/containers/podman/issues/16704
            host_mmproj_path = '/mnt' + host_mmproj_path
        mount = f"""
        - mountPath: {self.dest_mmproj_path}
          name: mmproj"""
        volume = f"""
      - hostPath:
          path: {host_mmproj_path}
        name: mmproj"""
        return mount, volume

    def __gen_ports(self) -> str:
        if not hasattr(self.args, "port"):
            return ""

        p = self.args.port.split(":", 2)
        ports = f"""\
        ports:
        - containerPort: {p[0]}"""
        if len(p) > 1:
            ports += f"""
          hostPort: {p[1]}"""

        return ports

    def __gen_env_vars(self) -> str:
        env_vars = get_accel_env_vars()

        if hasattr(self.args, "env"):
            extra_env = dict(i.split("=", 1) for i in self.args.env)
            env_vars |= extra_env

        if not env_vars:
            return ""

        env_spec = """\
        env:"""

        for k, v in sorted(env_vars.items()):
            env_spec += f"""
        - name: {k}
          value: \"{v}\""""

        return env_spec

    def __gen_security_context(self) -> str:
        return """\
        securityContext:
          allowPrivilegeEscalation: false
          capabilities:
            drop:
            - CAP_CHOWN
            - CAP_FOWNER
            - CAP_FSETID
            - CAP_KILL
            - CAP_NET_BIND_SERVICE
            - CAP_SETFCAP
            - CAP_SETGID
            - CAP_SETPCAP
            - CAP_SETUID
            - CAP_SYS_CHROOT
            add:
            - CAP_DAC_OVERRIDE
          seLinuxOptions:
            type: spc_t"""

    def __gen_resources(self) -> str:
        if check_nvidia() == "cuda":
            return """
        resources:
          limits:
             'nvidia.com/gpu=all': 1"""

        devices = get_gpu_devices()
        if devices:
            limits = "".join(f"\n             'podman.io/device={path}': 1" for path in devices.values())
            return f"""
        resources:
          limits:{limits}"""

        return ""

    def __gen_container(self, container_args) -> str:
        content = f"""\
      - name: {container_args["name"]}
        image: {container_args["image"]}
"""
        if "command" in container_args:
            content += f"""\
        command: ["{container_args["command"]}"]
"""
        content += f"""\
        args: {container_args["args"]}
{container_args["env_string"]}"""
        content += f"""
{self.__gen_security_context()}
{container_args["port_string"]}"""
        if "mounts_string" in container_args:
            content += f"""
{container_args["mounts_string"]}"""
        if "resources_string" in container_args:
            content += f"""\
{container_args["resources_string"]}"""
        return content

    def generate_content(self, name=None, labels="", container=None) -> str:
        env_string = self.__gen_env_vars()
        port_string = self.__gen_ports()
        mounts_string, volume_string = self._gen_volumes()
        _version = version()
        model_container = {
            "name": self.name if name is None else name,
            "image": self.image,
            "args": self.exec_args[1:],
            "env_string": env_string,
            "port_string": port_string,
            "mounts_string": mounts_string,
            "resources_string": self.__gen_resources(),
        }
        if not isinstance(self.exec_args[0], ContainerEntryPoint):
            model_container["command"] = self.exec_args[0]

        content = f"""\
# Save the output of this file and use kubectl create -f to import
# it into Kubernetes.
#
# Created with ramalama-{_version}
apiVersion: apps/v1
kind: Deployment
metadata:
  name: {self.name}
  labels:
    app: {self.name}
spec:
  replicas: 1
  selector:
    matchLabels:
      app: {self.name}
  template:
    metadata:
      labels:
        app: {self.name}{labels}
    spec:
      containers:
{self.__gen_container(model_container)}
"""
        if container is not None:
            content += f"""{self.__gen_container(container)}
"""
        content += f"""{volume_string}
"""

        return content

    def generate(self) -> PlainFile:
        return genfile(self.name, self.generate_content())


def genfile(name, content) -> PlainFile:
    file_name = f"{name}.yaml"
    print(f"Generating Kubernetes YAML file: {file_name}")

    file = PlainFile(file_name)
    file.content = content
    return file
