from __future__ import annotations

import os
import typing as t
from pathlib import Path

from rich.pretty import Pretty
from rich.syntax import Syntax

from starwhale.utils import console, load_yaml, pretty_bytes, in_production
from starwhale.consts import (
    PythonRunEnv,
    DEFAULT_PAGE_IDX,
    DEFAULT_PAGE_SIZE,
    STANDALONE_INSTANCE,
)
from starwhale.base.uri import URI
from starwhale.base.type import URIType, InstanceType
from starwhale.base.view import BaseTermView
from starwhale.utils.venv import get_venv_env, get_conda_env, get_python_run_env
from starwhale.utils.error import NotFoundError, NoSupportError, ExclusiveArgsError
from starwhale.utils.config import SWCliConfigMixed

from .model import Runtime, RuntimeInfoFilter, StandaloneRuntime


class RuntimeTermView(BaseTermView):
    def __init__(self, runtime_uri: str | URI) -> None:
        super().__init__()

        if isinstance(runtime_uri, URI):
            self.uri = runtime_uri
        else:
            self.uri = URI(runtime_uri, expected_type=URIType.RUNTIME)

        self.runtime = Runtime.get_runtime(self.uri)

    @BaseTermView._simple_action_print
    def remove(self, force: bool = False) -> t.Tuple[bool, str]:
        return self.runtime.remove(force)

    @BaseTermView._simple_action_print
    def recover(self, force: bool = False) -> t.Tuple[bool, str]:
        return self.runtime.recover(force)

    @BaseTermView._pager
    @BaseTermView._header
    def history(self, fullname: bool = False) -> t.List[t.Dict[str, t.Any]]:
        fullname = fullname or self.uri.instance_type == InstanceType.CLOUD
        return self._print_history(
            title="Runtime History", history=self.runtime.history(), fullname=fullname
        )

    def info(
        self,
        output_filter: RuntimeInfoFilter = RuntimeInfoFilter.basic,
    ) -> None:
        info = self.runtime.info()
        if not info:
            console.print(
                f":anguished_face: No runtime info found: {self.uri}", style="red"
            )
            return

        basic_content = Pretty(info["basic"], expand_all=True)
        runtime_content = Syntax(
            info.get("runtime_yaml", ""), "yaml", theme="ansi_dark"
        )
        manifest_content = Pretty(info.get("manifest", {}), expand_all=True)
        _locks = []
        for fname, content in info.get("lock", {}).items():
            _locks.append(f"#lock file: {fname}")
            _locks.append(content)
        lock_content = "\n".join(_locks)

        if output_filter == RuntimeInfoFilter.basic:
            console.print(basic_content)
        elif output_filter == RuntimeInfoFilter.runtime_yaml:
            console.print(runtime_content)
        elif output_filter == RuntimeInfoFilter.lock:
            console.print(lock_content)
        elif output_filter == RuntimeInfoFilter.manifest:
            console.print(manifest_content)
        else:
            console.rule("[green bold] Runtime Basic Info")
            console.print(basic_content)
            console.rule("[green bold] runtime.yaml")
            console.print(runtime_content)
            console.rule("[green bold] _manifest.yaml")
            console.print(manifest_content)
            console.rule("[green bold] Runtime Lock Files")
            console.print(lock_content)

    @classmethod
    @BaseTermView._only_standalone
    def activate(cls, uri: URI, force_restore: bool = False) -> None:
        Runtime.activate(uri, force_restore)

    @BaseTermView._only_standalone
    def dockerize(
        self,
        tags: t.List[str],
        push: bool,
        platforms: t.List[str],
        dry_run: bool,
        use_starwhale_builder: bool,
        reset_qemu_static: bool,
    ) -> None:
        self.runtime.dockerize(
            tags=tags,
            push=push,
            platforms=platforms,
            dry_run=dry_run,
            use_starwhale_builder=use_starwhale_builder,
            reset_qemu_static=reset_qemu_static,
        )

    @classmethod
    @BaseTermView._only_standalone
    def lock(
        cls,
        target_dir: str | Path,
        yaml_path: str | Path,
        env_name: str = "",
        env_prefix_path: str = "",
        no_cache: bool = False,
        stdout: bool = False,
        include_editable: bool = False,
        include_local_wheel: bool = False,
        emit_pip_options: bool = False,
        env_use_shell: bool = False,
    ) -> None:
        Runtime.lock(
            target_dir=target_dir,
            yaml_path=yaml_path,
            env_name=env_name,
            env_prefix_path=env_prefix_path,
            stdout=stdout,
            no_cache=no_cache,
            include_editable=include_editable,
            include_local_wheel=include_local_wheel,
            emit_pip_options=emit_pip_options,
            env_use_shell=env_use_shell,
        )

    @classmethod
    @BaseTermView._only_standalone
    def build_from_docker_image(
        cls, image: str, runtime_name: str = "", project: str = ""
    ) -> None:
        runtime_name = runtime_name or image.split(":")[0].split("/")[-1]
        runtime_uri = cls.prepare_build_bundle(
            project=project, bundle_name=runtime_name, typ=URIType.RUNTIME
        )

        rt = Runtime.get_runtime(runtime_uri)
        rt.build_from_docker_image(image=image, runtime_name=runtime_name)

    @classmethod
    @BaseTermView._only_standalone
    def build_from_python_env(
        cls,
        runtime_name: str = "",
        conda_name: str = "",
        conda_prefix: str = "",
        venv_prefix: str = "",
        project: str = "",
        cuda: str = "",
        cudnn: str = "",
        arch: str = "",
        download_all_deps: bool = False,
        include_editable: bool = False,
        include_local_wheel: bool = False,
    ) -> URI:
        set_args = list(filter(bool, (conda_name, conda_prefix, venv_prefix)))
        if len(set_args) >= 2:
            raise ExclusiveArgsError(
                f"conda_prefix({conda_prefix}), conda_name({conda_name}), venv_prefix({venv_prefix}) are the mutex args."
            )

        if conda_name:
            mode = PythonRunEnv.CONDA
            candidate_runtime_name = conda_name
        elif conda_prefix:
            mode = PythonRunEnv.CONDA
            candidate_runtime_name = conda_prefix
        elif venv_prefix:
            mode = PythonRunEnv.VENV
            candidate_runtime_name = venv_prefix
        else:
            mode = get_python_run_env()
            candidate_runtime_name = (
                get_conda_env() if mode == PythonRunEnv.CONDA else get_venv_env()
            ) or "default"

        if runtime_name == "":
            runtime_name = candidate_runtime_name.strip("/").split("/")[-1]

        runtime_uri = cls.prepare_build_bundle(
            project=project, bundle_name=runtime_name, typ=URIType.RUNTIME
        )
        rt = Runtime.get_runtime(runtime_uri)
        rt.build_from_python_env(
            runtime_name=runtime_name,
            mode=mode,
            conda_name=conda_name,
            conda_prefix=conda_prefix,
            venv_prefix=venv_prefix,
            cuda=cuda,
            cudnn=cudnn,
            arch=arch,
            download_all_deps=download_all_deps,
            include_editable=include_editable,
            include_local_wheel=include_local_wheel,
        )
        return runtime_uri

    @classmethod
    @BaseTermView._only_standalone
    def build_from_runtime_yaml(
        cls,
        workdir: str | Path,
        yaml_path: str | Path,
        runtime_name: str = "",
        project: str = "",
        download_all_deps: bool = False,
        include_editable: bool = False,
        include_local_wheel: bool = False,
        no_cache: bool = False,
        disable_env_lock: bool = False,
    ) -> URI:
        workdir = Path(workdir)
        yaml_path = Path(yaml_path)

        if not yaml_path.exists():
            raise NotFoundError(f"not found runtime yaml:{yaml_path}")
        _config = load_yaml(yaml_path)
        runtime_name = runtime_name or _config["name"]

        _runtime_uri = cls.prepare_build_bundle(
            project=project, bundle_name=runtime_name, typ=URIType.RUNTIME
        )
        if include_editable:
            console.print(
                ":bell: [red bold]runtime will include pypi editable package[/] :bell:"
            )
        else:
            console.print(
                ":bird: [red bold]runtime will ignore pypi editable package[/]"
            )

        _rt = Runtime.get_runtime(_runtime_uri)
        _rt.build_from_runtime_yaml(
            workdir=workdir,
            yaml_path=yaml_path,
            download_all_deps=download_all_deps,
            include_editable=include_editable,
            include_local_wheel=include_local_wheel,
            no_cache=no_cache,
            disable_env_lock=disable_env_lock,
        )
        return _runtime_uri

    @BaseTermView._only_standalone
    def extract(self, force: bool = False, target: t.Union[str, Path] = "") -> None:
        console.print(":oncoming_police_car: try to extract ...")
        path = self.runtime.extract(force, target)
        console.print(f":clap: extracted @ {path.resolve()} :tada:")

    @classmethod
    def list(
        cls,
        project_uri: str = "",
        fullname: bool = False,
        show_removed: bool = False,
        page: int = DEFAULT_PAGE_IDX,
        size: int = DEFAULT_PAGE_SIZE,
        filters: t.Optional[t.List[str]] = None,
    ) -> t.Tuple[t.List[t.Dict[str, t.Any]], t.Dict[str, t.Any]]:
        filters = filters or []
        _uri = URI(project_uri, expected_type=URIType.PROJECT)
        cls.must_have_project(_uri)
        fullname = fullname or (_uri.instance_type == InstanceType.CLOUD)
        _runtimes, _pager = Runtime.list(_uri, page, size, filters)
        _data = BaseTermView.list_data(_runtimes, show_removed, fullname)
        return _data, _pager

    @classmethod
    @BaseTermView._only_standalone
    def quickstart_from_uri(
        cls,
        workdir: t.Union[Path, str],
        name: str,
        uri: URI,
        force: bool = False,
        disable_restore: bool = False,
    ) -> None:
        console.print(
            f":construction: quickstart Starwhale Runtime[{name}] environment from runtime URI({uri})..."
        )
        _sw_config = SWCliConfigMixed()
        if _sw_config.current_instance != STANDALONE_INSTANCE:
            raise NoSupportError(f"{_sw_config.current_instance} quickstart")
        StandaloneRuntime.quickstart_from_uri(
            workdir, name, uri, force, disable_restore
        )
        console.print(":clap: Starwhale Runtime environment is ready to use :tada:")

    @classmethod
    @BaseTermView._only_standalone
    def quickstart_from_ishell(
        cls,
        workdir: t.Union[Path, str],
        name: str,
        mode: str,
        disable_create_env: bool = False,
        force: bool = False,
        interactive: bool = False,
    ) -> None:
        console.print(
            f":construction: quickstart Starwhale Runtime[{name}] environment..."
        )
        StandaloneRuntime.quickstart_from_ishell(
            workdir, name, mode, disable_create_env, force, interactive
        )
        console.print(":clap: Starwhale Runtime environment is ready to use :tada:")

    @classmethod
    @BaseTermView._only_standalone
    def restore(cls, target: str) -> None:
        if in_production() or (os.path.exists(target) and os.path.isdir(target)):
            workdir = Path(target)
        else:
            _uri = URI(target, URIType.RUNTIME)
            _runtime = StandaloneRuntime(_uri)
            workdir = _runtime.store.snapshot_workdir
            if not workdir.exists():
                _runtime.extract(force=True, target=workdir)

        console.print(
            f":golfer: try to restore python runtime environment: {workdir} ..."
        )
        Runtime.restore(workdir)

    @classmethod
    def copy(
        cls,
        src_uri: str,
        dest_uri: str,
        force: bool = False,
        dest_local_project_uri: str = "",
    ) -> None:
        Runtime.copy(src_uri, dest_uri, force, dest_local_project_uri)
        console.print(":clap: copy done.")

    @BaseTermView._header
    def tag(
        self, tags: t.List[str], remove: bool = False, ignore_errors: bool = False
    ) -> None:
        # TODO: refactor model/runtime/dataset tag view-model
        if remove:
            console.print(f":golfer: remove tags [red]{tags}[/] @ {self.uri}...")
            self.runtime.remove_tags(tags, ignore_errors)
        else:
            console.print(f":surfer: add tags [red]{tags}[/] @ {self.uri}...")
            self.runtime.add_tags(tags, ignore_errors)


class RuntimeTermViewRich(RuntimeTermView):
    @classmethod
    @BaseTermView._pager
    def list(
        cls,
        project_uri: str = "",
        fullname: bool = False,
        show_removed: bool = False,
        page: int = DEFAULT_PAGE_IDX,
        size: int = DEFAULT_PAGE_SIZE,
        filters: t.Optional[t.List[str]] = None,
    ) -> t.Tuple[t.List[t.Dict[str, t.Any]], t.Dict[str, t.Any]]:
        filters = filters or []
        _data, _pager = super().list(
            project_uri, fullname, show_removed, page, size, filters
        )

        custom_column: t.Dict[str, t.Callable[[t.Any], str]] = {
            "tags": lambda x: ",".join(x),
            "size": lambda x: pretty_bytes(x),
            "runtime": cls.place_holder_for_empty(),
        }

        cls.print_header()
        cls.print_table("Runtime List", _data, custom_column=custom_column)
        return _data, _pager


class RuntimeTermViewJson(RuntimeTermView):
    @classmethod
    def list(  # type: ignore
        cls,
        project_uri: str = "",
        fullname: bool = False,
        show_removed: bool = False,
        page: int = DEFAULT_PAGE_IDX,
        size: int = DEFAULT_PAGE_SIZE,
        filters: t.Optional[t.List[str]] = None,
    ) -> None:
        filters = filters or []
        _data, _pager = super().list(
            project_uri, fullname, show_removed, page, size, filters
        )
        cls.pretty_json(_data)

    def info(
        self,
        output_filter: RuntimeInfoFilter = RuntimeInfoFilter.basic,
    ) -> None:
        info = self.runtime.info()

        if output_filter == RuntimeInfoFilter.basic:
            info = {"basic": info.get("basic", {})}
        elif output_filter == RuntimeInfoFilter.lock:
            info = {"lock": info.get("lock", {})}
        elif output_filter == RuntimeInfoFilter.manifest:
            info = {"manifest": info.get("manifest", {})}
        elif output_filter == RuntimeInfoFilter.runtime_yaml:
            info = {"runtime_yaml": info.get("runtime_yaml", "")}

        self.pretty_json(info)

    def history(self, fullname: bool = False) -> None:
        fullname = fullname or self.uri.instance_type == InstanceType.CLOUD
        _data = BaseTermView.get_history_data(
            history=self.runtime.history(), fullname=fullname
        )
        self.pretty_json(_data)


def get_term_view(ctx_obj: t.Dict) -> t.Type[RuntimeTermView]:
    return (
        RuntimeTermViewJson if ctx_obj.get("output") == "json" else RuntimeTermViewRich
    )
