import concurrent.futures
import json
import os
from pathlib import Path
from typing import Dict, Generator, List, Literal, Optional

import click
import rich
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
from tqdm import tqdm

from lightning_sdk.api.lit_container_api import DockerNotRunningError, LCRAuthFailedError, LitContainerApi
from lightning_sdk.api.utils import _get_cloud_url
from lightning_sdk.cli.legacy.exceptions import StudioCliError
from lightning_sdk.cli.utils.resource_resolution import resolve_studio, resolve_teamspace
from lightning_sdk.constants import _LIGHTNING_DEBUG
from lightning_sdk.exceptions import DeprecatedCommand, DeprecatedError
from lightning_sdk.models import upload_model as _upload_model
from lightning_sdk.studio import Studio

_STUDIO_UPLOAD_STATUS_PATH = "~/.lightning/studios/uploads"

UploadRecovery = Optional[Literal["resume", "restart"]]


def resolve_upload_recovery(resume: bool, restart: bool) -> UploadRecovery:
    if resume and restart:
        raise click.UsageError("--resume and --restart are mutually exclusive.")
    if resume:
        return "resume"
    if restart:
        return "restart"
    return None


@click.group("upload")
def upload() -> None:
    """Upload assets to Lightning AI."""


@upload.command("model")
@click.argument("name", metavar="ORG-NAME/TEAMSPACE-NAME/MODEL-NAME")
@click.argument("path")
@click.option(
    "--cloud-account",
    "--cloud_account",
    default=None,
    help="The name of the cloud account to store the Model in.",
)
def model(name: str, path: str, cloud_account: Optional[str] = None) -> None:
    """Upload a model a teamspace.

    Example:
        lightning upload model NAME

    NAME: the name of the model to upload (Should be of format <ORGANIZATION-NAME>/<TEAMSPACE-NAME>/<MODEL-NAME>).
    """
    _upload_model(name, path, cloud_account=cloud_account)


@upload.command(
    name="folder",
    cls=DeprecatedCommand,
    message="Studio uploads via 'lightning upload folder' are deprecated. Use 'lightning studio cp -r' instead.",
)
@click.argument("path", type=click.Path(), required=False, nargs=-1)
@click.option("--studio", default=None, hidden=True)
@click.option(
    "--remote-path",
    "--remote_path",
    default=None,
    hidden=True,
)
def folder(path: str, studio: Optional[str], remote_path: Optional[str]) -> None:
    """[DEPRECATED] Use 'lightning studio cp -r' instead."""
    raise DeprecatedError(
        "Studio uploads via 'lightning upload folder' are deprecated. Use 'lightning studio cp -r' instead."
    )


@upload.command(
    name="file",
    cls=DeprecatedCommand,
    message="Studio uploads via 'lightning upload file' are deprecated. Use 'lightning studio cp' instead.",
)
@click.argument("path", type=click.Path(), required=False, nargs=-1)
@click.option(
    "--studio",
    default=None,
    hidden=True,
)
@click.option(
    "--remote-path",
    "--remote_path",
    default=None,
    hidden=True,
)
def file(path: str, studio: Optional[str] = None, remote_path: Optional[str] = None) -> None:
    """[DEPRECATED] Use 'lightning studio cp' instead."""
    raise DeprecatedError(
        "Studio uploads via 'lightning upload file' are deprecated. Use 'lightning studio cp' instead."
    )


@upload.command("container")
@click.argument("container")
@click.option("--tag", default="latest", help="The tag of the container to upload.")
@click.option(
    "--teamspace",
    default=None,
    help=(
        "The teamspace the studio is part of. "
        "Should be of format <OWNER>/<TEAMSPACE_NAME>. "
        "If not specified, tries to infer from the environment (e.g. when run from within a Studio.)"
    ),
)
@click.option(
    "--cloud-account",
    "--cloud_account",  # The UI will present the above variant, using this as a secondary to be consistent w/ models
    default=None,
    help="The name of the cloud account to store the Container in.",
)
@click.option(
    "--platform",
    default="linux/amd64",
    help="This is the platform the container pulled and push to Lightning AI will run on.",
)
def upload_container(
    container: str,
    tag: str = "latest",
    teamspace: Optional[str] = None,
    cloud_account: Optional[str] = None,
    platform: Optional[str] = "linux/amd64",
) -> None:
    """Upload a container to Lightning AI's container registry."""
    resolved_teamspace = resolve_teamspace(teamspace)
    platform = platform or "linux/amd64"
    console = Console()
    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        TimeElapsedColumn(),
        console=console,
        transient=False,
    ) as progress:
        try:
            if platform != "linux/amd64":
                console.print(
                    "[yellow]Warning: The platform you selected (" + platform + ") may not deploy successfully[/yellow]"
                )
            api = LitContainerApi()
            push_task = progress.add_task("Pushing Docker image", total=None)
            console.print("Authenticating with Lightning Container Registry...")
            try:
                api.authenticate()
                console.print("Authenticated with Lightning Container Registry", style="green")
            except Exception:
                # let the push with retry take control of auth moving forward
                pass

            lines = api.upload_container(
                container, resolved_teamspace, tag, cloud_account, platform, return_final_dict=True
            )
            _print_docker_push(lines, console, progress, push_task)
        except DockerNotRunningError as e:
            e.print_help()
            return
        except LCRAuthFailedError:
            console.print("Re-authenticating with Lightning Container Registry...")
            if not api.authenticate(reauth=True):
                raise StudioCliError("Failed to authenticate with Lightning Container Registry") from None
            console.print("Authenticated with Lightning Container Registry", style="green")
            lines = api.upload_container(container, resolved_teamspace, tag, cloud_account, platform)
            _print_docker_push(lines, console, progress, push_task)
        except Exception as e:
            if _LIGHTNING_DEBUG:
                print(e)
                if e.__cause__:
                    print(e.__cause__)

            progress.update(push_task, description=f"[bold red]Error: {e!s}[/]")
            progress.stop()
            return
        progress.update(push_task, description="[green]Container pushed![/green]")


def _folder(path: str, studio: Optional[str] = None, remote_path: Optional[str] = None) -> None:
    """Upload a folder to a Studio."""
    console = Console()
    if remote_path is None:
        remote_path = os.path.basename(path)

    if not Path(path).exists():
        raise FileNotFoundError(f"The provided path does not exist: {path}.")
    if not Path(path).is_dir():
        raise StudioCliError(f"The provided path is not a folder: {path}. Use `lightning upload file` instead.")

    selected_studio = _resolve_studio(studio)

    console.print(f"Uploading to {selected_studio.teamspace.name}/{selected_studio.name}")

    _upload_folder(path, remote_path, selected_studio)

    studio_url = (
        _get_cloud_url().replace(":443", "")
        + "/"
        + selected_studio.owner.name
        + "/"
        + selected_studio.teamspace.name
        + "/studios/"
        + selected_studio.name
    )
    console.print(f"See your files at {studio_url}")


def _upload_folder(path: str, remote_path: str, studio: Studio, recovery: UploadRecovery = None) -> None:
    pairs = {}
    for root, _, files in os.walk(path):
        rel_root = os.path.relpath(root, path)
        for file_name in files:
            pairs[os.path.join(root, file_name)] = os.path.join(remote_path, rel_root, file_name)

    upload_state = _resolve_previous_upload_state(studio, remote_path, pairs, recovery)

    with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
        futures = _start_parallel_upload(executor, studio, upload_state)

        update_fn = tqdm(total=len(upload_state)).update if _global_upload_progress(upload_state) else lambda x: None

        for future in concurrent.futures.as_completed(futures):
            upload_state.pop(future.result())
            _dump_current_upload_state(studio, remote_path, upload_state)
            update_fn(1)


def _file(path: str, studio: Optional[str] = None, remote_path: Optional[str] = None) -> None:
    """Upload a file to a Studio."""
    console = Console()
    if remote_path is None:
        remote_path = os.path.basename(path)

    if Path(path).is_dir():
        raise StudioCliError(f"The provided path is a folder: {path}. Use `lightning upload folder` instead.")
    if not Path(path).exists():
        raise FileNotFoundError(f"The provided path does not exist: {path}.")

    selected_studio = _resolve_studio(studio)

    console.print(f"Uploading to {selected_studio.teamspace.name}/{selected_studio.name}")

    _single_file_upload(selected_studio, path, remote_path, True)

    studio_url = (
        _get_cloud_url().replace(":443", "")
        + "/"
        + selected_studio.owner.name
        + "/"
        + selected_studio.teamspace.name
        + "/studios/"
        + selected_studio.name
    )
    console.print(f"See your file at {studio_url}")


def _resolve_studio(studio: Optional[str]) -> Studio:
    try:
        resolved_teamspace = resolve_teamspace()
        return resolve_studio(studio, resolved_teamspace)

    except Exception as e:
        raise StudioCliError(
            f"Could not find the given Studio {studio} to upload files to. "
            "Please contact Lightning AI directly to resolve this issue."
        ) from e


def _print_docker_push(lines: Generator, console: Console, progress: Progress, push_task: rich.progress.TaskID) -> None:
    for line in lines:
        if "status" in line:
            console.print(line["status"], style="bright_black")
            progress.update(push_task, description="Pushing Docker image")
        elif "aux" in line:
            console.print(line["aux"], style="bright_black")
        elif "error" in line:
            progress.stop()
            console.print(f"\n[red]{line}[/red]")
            return
        elif "finish" in line:
            if "url" in line:
                console.print(f"Container available at [i]{line['url']}[/i]")
            return
        else:
            console.print(line, style="bright_black")


def _start_parallel_upload(
    executor: concurrent.futures.ThreadPoolExecutor, studio: Studio, upload_state: Dict[str, str]
) -> List[concurrent.futures.Future]:
    # only add progress bar on individual uploads with less than 10 files
    progress_bar = not _global_upload_progress(upload_state)

    futures = []
    for k, v in upload_state.items():
        futures.append(
            executor.submit(_single_file_upload, studio=studio, local_path=k, remote_path=v, progress_bar=progress_bar)
        )

    return futures


def _single_file_upload(studio: Studio, local_path: str, remote_path: str, progress_bar: bool) -> str:
    studio.upload_file(local_path, remote_path, progress_bar)
    return local_path


def _dump_current_upload_state(studio: Studio, remote_path: str, state_dict: Dict[str, str]) -> None:
    """Dumps the current upload state so that we can safely resume later."""
    curr_path = _upload_state_path(studio, remote_path)

    dirpath = curr_path.parent
    if state_dict:
        dirpath.mkdir(parents=True, exist_ok=True)
        curr_path.write_text(json.dumps(state_dict, indent=4))
        return

    if curr_path.exists():
        curr_path.unlink()
    if dirpath.exists():
        os.removedirs(dirpath)


def _upload_state_path(studio: Studio, remote_path: str) -> Path:
    return Path(
        os.path.abspath(
            os.path.expandvars(
                os.path.expanduser(
                    os.path.join(
                        _STUDIO_UPLOAD_STATUS_PATH,
                        studio._studio.id,
                        remote_path + ".json",
                    )
                )
            )
        )
    )


def _resolve_previous_upload_state(
    studio: Studio,
    remote_path: str,
    state_dict: Dict[str, str],
    recovery: UploadRecovery,
) -> Dict[str, str]:
    """Resolves potential previous uploads to continue if possible."""
    curr_path = _upload_state_path(studio, remote_path)
    exists = curr_path.is_file()
    if recovery == "resume":
        if not exists:
            raise click.UsageError("There is nothing to resume.")
        return json.loads(curr_path.read_text())
    if recovery == "restart":
        return state_dict
    if exists:
        raise click.UsageError("An incomplete upload exists. Pass --resume or --restart.")
    return state_dict


def _global_upload_progress(upload_state: Dict[str, str]) -> bool:
    return len(upload_state) > 10