import os
import os.path
from collections.abc import Callable
from itertools import chain
from pathlib import Path
from typing import Optional, Protocol, TypeVar

from radicli import Arg

from prodigy_teams_broker_sdk import Client as BrokerClient
from prodigy_teams_broker_sdk import models as broker_models

from ...cli import cli
from ...errors import BrokerError, CLIError
from ...messages import Messages
from ...ui import Progress, print_mutation_result
from ...util import _resolve_broker_ref, is_local, resolve_remote_path
from .._state import get_auth_state, get_saved_settings


class RsyncCallable(Protocol):
    def __call__(
        self,
        *,
        src: str,
        dest: str,
        broker_client: BrokerClient,
    ) -> str | None: ...


@cli.subcommand(
    "files",
    "rsync",
    src=Arg(help=Messages.remote_local_path.format(noun="source")),
    dest=Arg(help=Messages.remote_local_path.format(noun="destination")),
    cluster_host=Arg("--cluster-host", help=Messages.cluster_host),
    as_json=Arg("--json", help=Messages.as_json),
)
def rsync(
    src: str, dest: str, cluster_host: Optional[str] = None, as_json: bool = False
) -> None:
    """Rsync files to and from the cluster"""
    settings = get_saved_settings()
    auth = get_auth_state()
    broker_host = str(
        _resolve_broker_ref(auth.client, cluster_host or settings.broker_host)
    )
    is_src_local = is_local(src)
    is_dest_local = is_local(dest)

    def _rsync(f: RsyncCallable, *, src, dest) -> str | None:
        return f(
            src=src,
            dest=dest,
            broker_client=auth.broker_client,
        )

    if is_src_local and is_dest_local:
        raise CLIError(Messages.E015.format(verb="rsync"), f"{src} -> {dest}")
    elif is_src_local and not is_dest_local:
        _rsync(
            _rsync_upload,
            src=src,
            dest=resolve_remote_path(auth.client, dest, broker_host),
        )
    elif not is_src_local and is_dest_local:
        _rsync(
            _rsync_download,
            src=resolve_remote_path(auth.client, src, broker_host),
            dest=dest,
        )
    else:
        _rsync(
            _rsync_on_remote,
            src=resolve_remote_path(auth.client, src, broker_host),
            dest=resolve_remote_path(auth.client, dest, broker_host),
        )
    if as_json:
        print_mutation_result(
            {"status": "ok", "src": src, "dest": dest},
            "",
            as_json=True,
        )


T1 = TypeVar("T1")
T2 = TypeVar("T2")
R = TypeVar("R")


def diff(
    left: list[T1],
    right: list[T2],
    left_key: Callable[[T1], R],
    right_key: Callable[[T2], R],
) -> tuple[list[T1], list[T2]]:
    left_idx = {left_key(value): value for value in left}
    right_idx = {right_key(value): value for value in right}
    left_result = [left_idx[k] for k in left_idx.keys() - right_idx.keys()]
    right_result = [right_idx[k] for k in right_idx.keys() - left_idx.keys()]
    return left_result, right_result


def _rsync_upload(
    *,
    broker_client: BrokerClient,
    src: str,
    dest: str,
) -> None:
    src_path = Path(src).resolve().absolute()
    src = str(src_path)
    src_paths = sorted(
        [
            str(path.resolve().absolute())
            for path in chain((src_path,), src_path.glob("**/*"))
            if path.is_file()
        ],
        key=lambda path: path.removeprefix(src),
    )
    body = broker_models.Listing(path=dest, recurse=True, include_stats=False)
    try:
        files = broker_client.files.list_dir(body)
    except BrokerError as e:
        raise CLIError(Messages.E018, e)
    dest_paths = sorted(
        files.paths,
        key=lambda path: path.removeprefix(dest),
    )
    to_copy, to_delete = diff(
        src_paths,
        dest_paths,
        left_key=lambda path: path.removeprefix(src),
        right_key=lambda path: path.removeprefix(dest),
    )
    _rsync_upload_util(
        src_base=src,
        dest_base=dest,
        to_copy=to_copy,
        to_delete=to_delete,
        broker_client=broker_client,
    )


def _rsync_upload_util(
    *,
    src_base: str,
    dest_base: str,
    to_copy: list[str],
    to_delete: list[str],
    broker_client,
) -> None:
    def get_dest(path: str) -> str:
        return os.path.join(dest_base, path.removeprefix(src_base).removeprefix("/"))

    with Progress(to_copy) as paths:
        for path in paths:
            body = Path(path).open("rb").read()
            try:
                broker_client.files.upload(
                    body, dest=get_dest(path), overwrite=True, make_dirs=True
                )
            except BrokerError as e:
                raise CLIError(Messages.E017.format(verb="upload"), e)
    with Progress(to_delete) as paths:
        for path in paths:
            body = broker_models.Deleting(path=path, missing_ok=False, recurse=False)
            try:
                broker_client.files.delete(body)
            except BrokerError as e:
                raise CLIError(Messages.E017.format(verb="delete"), e)


def _rsync_download(
    *,
    broker_client: BrokerClient,
    src: str,
    dest: str,
) -> None:
    body = broker_models.Listing(path=src, recurse=True, include_stats=False)
    try:
        paths = broker_client.files.list_dir(body)
    except BrokerError as e:
        raise CLIError(Messages.E018, e)
    src_paths = sorted(paths.paths)
    if src_paths == src:
        raise CLIError(Messages.E024, src)
    dest_path = Path(dest).resolve().absolute()
    dest = str(dest_path)
    dest_paths = sorted(
        [
            str(path.resolve().absolute())
            for path in chain((dest_path,), dest_path.glob("**/*"))
            if path.is_file()
        ],
    )
    to_copy, to_delete = diff(
        src_paths,
        dest_paths,
        left_key=lambda path: path.removeprefix(src),
        right_key=lambda path: path.removeprefix(dest),
    )
    _rsync_download_util(
        src_base=src,
        dest_base=dest,
        to_copy=to_copy,
        to_delete=to_delete,
        broker_client=broker_client,
    )


def _rsync_download_util(
    *,
    src_base: str,
    dest_base: str,
    to_copy: list[str],
    to_delete: list[str],
    broker_client: BrokerClient,
) -> None:
    def get_dest(path: str) -> str:
        return os.path.join(dest_base, path.removeprefix(src_base).removeprefix("/"))

    with Progress(to_copy) as paths:
        for path in paths:
            body = broker_models.Downloading(target=path)
            try:
                content = broker_client.files.download(body)
            except BrokerError as e:
                raise CLIError(Messages.E017.format(verb="download"), e)
            dest = Path(get_dest(path))
            dest.parent.mkdir(parents=True, exist_ok=True)
            with dest.open("wb") as file_:
                file_.write(content.read())
    with Progress(to_delete) as paths:
        for path in paths:
            path = Path(path)
            if path.exists():
                path.unlink()


def _rsync_on_remote(
    *,
    broker_client: BrokerClient,
    src: str,
    dest: str,
) -> None:
    body = broker_models.Copying(src=src, dest=dest, make_dirs=True, overwrite=True)
    try:
        plan = broker_client.files.plan_directory_rsync(body)
    except BrokerError as e:
        raise CLIError(Messages.E022, e)
    with Progress(plan.copy_) as copies:
        for file_copy_plan in copies:
            body = broker_models.Copying(
                src=file_copy_plan.src,
                dest=file_copy_plan.dest,
                make_dirs=True,
                overwrite=True,
            )
            try:
                broker_client.files.copy(body)
            except BrokerError as e:
                raise CLIError(Messages.E017.format(verb="copy"), e)
    with Progress(plan.delete) as paths:
        for path in paths:
            body = broker_models.Deleting(path=path, missing_ok=False, recurse=False)
            try:
                broker_client.files.delete(body)
            except BrokerError as e:
                raise CLIError(Messages.E017.format(verb="delete"), e)
