import warnings

with warnings.catch_warnings():
    warnings.simplefilter("ignore")

    import rich_click as click
    import rich
    from rich.console import Console
    from rich.panel import Panel
    from rich.table import Table
    from rich.text import Text
    from rich import box
    from thunder import auth, auth_helper
    import os
    from os.path import join
    from scp import SCPClient
    import paramiko
    import subprocess
    import time
    import platform

    from thunder import thunder_helper
    from thunder import utils

    try:
        from importlib.metadata import version
    except Exception as e:
        from importlib_metadata import version

    import requests
    from packaging import version as version_parser
    from thunder import api
    from yaspin import yaspin


PACKAGE_NAME = "tnr"
ENABLE_RUN_COMMANDS = True if platform.system() == "Linux" else False
INSIDE_INSTANCE = False
CURRENT_IP = None


# Remove the DefaultCommandGroup class
DISPLAYED_WARNING = False
def get_token():
    global DISPLAYED_WARNING
    
    if "TNR_API_TOKEN" in os.environ:
        return os.environ["TNR_API_TOKEN"]

    token_file = auth_helper.get_credentials_file_path()
    if not os.path.exists(token_file):
        auth.login()

    with open(auth_helper.get_credentials_file_path(), "r") as f:
        lines = f.readlines()
        if len(lines) == 1:
            token = lines[0].strip()
            return token

    if not DISPLAYED_WARNING:
        click.echo(
            click.style(
                f"We switched to an API-token based auth system. Please go to console.thundercompute.com and generate a new token. Then save it to the TNR_API_TOKEN environment variable",
                fg="yellow",
                bold=True,
            )
        )
        DISPLAYED_WARNING = True
        
    id_token, refresh_token, uid = auth.load_tokens()
    if not id_token or not refresh_token:
        click.echo("Please log in to begin using Thunder Compute.")
        token = auth.login()
        return token

    if not api.is_token_valid(id_token):
        id_token, refresh_token, uid = auth.handle_token_refresh(refresh_token)

    return id_token


def init():
    global CURRENT_IP, INSIDE_INSTANCE

    if ENABLE_RUN_COMMANDS:
        utils.setup_instance()

    # Determine if we are currently in a TC instance
    try:
        CURRENT_IP = requests.get("https://ifconfig.co/json").json()["ip"]
    except Exception as e:
        CURRENT_IP = None

    success = False
    num_attempts = 0
    while not success and num_attempts < 5:
        success, error, instances = utils.get_instances(get_token())
        if not success:
            if error == 'Unauthorized':
                click.echo(
                    click.style(
                        f"Invalid token, please update the TNR_API_TOKEN environment variable or login again",
                        fg="red",
                        bold=True,
                    )
                )
                auth.login()
            else:
                click.echo(
                    click.style(
                        f"Failed to list Thunder Compute instances: {error}",
                        fg="red",
                        bold=True,
                    )
                )
                exit(1)
            num_attempts += 1
        
    if not success and num_attempts == 5:
        click.echo(
                click.style(
                    f"Failed to login to thunder compute 5 times. Please sign in with a valid token",
                    fg="red",
                    bold=True,
                )
            )
        exit(1)

    for instance_id, metadata in instances.items():
        if "ip" in metadata and metadata["ip"] == CURRENT_IP:
            INSIDE_INSTANCE = True
            break
init()

click.rich_click.USE_RICH_MARKUP = True
click.rich_click.GROUP_ARGUMENTS_OPTIONS = True
click.rich_click.SHOW_ARGUMENTS = True
click.rich_click.COMMAND_GROUPS = {
    "cli": [
        {
            "name": "Process management",
            "commands": (["device", "status"] if INSIDE_INSTANCE else ["status"]),
        },
        {
            "name": "Instance management",
            "commands": (
                ["create", "connect", "start", "stop", "delete", "scp"]
                if not INSIDE_INSTANCE
                else ["scp"]
            ),
        },
        (
            {
                "name": "Account management",
                "commands": ["login", "logout"],
            }
            if not INSIDE_INSTANCE
            else {}
        ),
    ]
}

COLOR = 'green' if INSIDE_INSTANCE else 'cyan'
click.rich_click.STYLE_OPTION = COLOR
click.rich_click.STYLE_COMMAND = COLOR

main_message = (
    f":link: [{COLOR}]You're connected to a Thunder Compute instance. Any process you run will automatically use a GPU on-demand.[/]"
    if INSIDE_INSTANCE
    else f":laptop_computer: [{COLOR}]You're in a local environment, use these commands to manage your Thunder Compute instances.[/]"
)


@click.group(
    cls=click.RichGroup,
    help=main_message,
    context_settings={"ignore_unknown_options": True, "allow_extra_args": True},
)
@click.version_option(version=version(PACKAGE_NAME))
def cli():
    meets_version, versions = does_meet_min_required_version()
    if not meets_version:
        raise click.ClickException(
            f'Failed to meet minimum required tnr version to proceed (current=={versions[0]}, required=={versions[1]}), please run "pip install --upgrade tnr" to update'
        )


if ENABLE_RUN_COMMANDS:

    @cli.command(
        help="Runs process on a remote Thunder Compute GPU. The GPU type is specified in the ~/.thunder/dev file. For more details, please go to thundercompute.com",
        context_settings={"ignore_unknown_options": True, "allow_extra_args": True},
        hidden=True,
    )
    @click.argument("args", nargs=-1, type=click.UNPROCESSED)
    def run(args):
        if not args:  # Check if args is empty
            raise click.ClickException("No arguments provided. Exiting...")

        token = get_token()
        endpoint = f"https://api.thundercompute.com:8443/uid"
        response = requests.get(endpoint, headers={"Authorization": f"Bearer {token}"})

        if response.status_code != 200:
            raise click.ClickException(
                "Failed to get info about user, is the API token correct?"
            )
        uid = response.text

        # Identify current computer
        basedir = join(os.path.expanduser("~"), ".thunder")
        device_id_file = join(basedir, "device_id")
        if not os.path.exists(device_id_file):
            try:
                endpoint = f"https://api.thundercompute.com:8443/next_id"
                response = requests.get(
                    endpoint, headers={"Authorization": f"Bearer {token}"}
                )
                identity = response.json()["id"]
                with open(device_id_file, "w+", encoding="utf-8") as f:
                    f.write(str(identity))
            except Exception as e:
                raise click.ClickException(
                    "Unable to create device identity file. Please report this error to the developers"
                )

            os.chmod(device_id_file, 0o400)

        # Create an instance of Task with required arguments
        task = thunder_helper.Task(args, uid)

        if not INSIDE_INSTANCE:
            message = "[yellow]Attaching to a remote GPU from a non-managed instance - this will hurt performance. If this is not intentional, please connect to a managed CPU instance using tnr create and tnr connect <INSTANCE ID>[/yellow]"
            panel = Panel(
                message,
                title=":warning:  Warning :warning: ",
                title_align="left",
                highlight=True,
                width=100,
                box=box.ROUNDED,
            )
            rich.print(panel)

        if not task.execute_task(token):
            return

        # Close the session
        if not task.close_session(token):
            click.echo("Failed to close the session.")

    @cli.command(
        help="If device_type is empty, displays the current GPU and a list of available GPUs. If device_name has a value, switches to the specified device. Input billing information at console.thundercompute.com to use devices besides NVIDIA T4s.",
        hidden=not INSIDE_INSTANCE,
    )
    @click.argument("device_type", required=False)
    @click.option("-n", "--ngpus", type=int, help="Number of GPUs to use")
    @click.option("--raw", is_flag=True, help="Output raw device information")
    def device(device_type, ngpus, raw):
        basedir = join(os.path.expanduser("~"), ".thunder")
        device_file = join(basedir, "dev")
        num_file = join(basedir, "ngpus")
        supported_devices = set(
            [
                "cpu",
                "t4",
                "v100",
                "a100",
                "l4",
                "p4",
                "p100",
                "h100",
            ]
        )

        if device_type is None:
            with open(device_file, "r") as f:
                device = f.read().strip()

            if device not in supported_devices:
                # If not valid, set it to the default value
                device = "t4"
                with open(device_file, "w") as f:
                    f.write(device)

            if ngpus is None:
                with open(num_file, "r") as f:
                    content = f.read().strip()

                if not content.isnumeric():
                    # If not valid, set it to the default value
                    ngpus = 1
                    with open(num_file, "w") as f:
                        f.write(str(ngpus))
                else:
                    ngpus = int(content)

            if raw is not None and raw:
                if ngpus == 1:
                    click.echo(device.upper())
                else:
                    click.echo(f"{ngpus}x{device.upper()}")
                return

            if device.lower() == "cpu":
                click.echo(
                    click.style(
                        "📖 No GPU selected - use `tnr device <gpu-type>` to select a GPU",
                        fg="white",
                    )
                )
                return

            console = Console()
            if ngpus == 1:
                console.print(f'[bold cyan]📖 Current GPU:[/] {device.upper()}')
            else:
                console.print(f'[bold cyan]📖 Current GPUs:[/][white] {ngpus} x {device.upper()}[/]')

            utils.display_available_gpus()
            return

        if device_type.lower() not in supported_devices:
            raise click.ClickException(
                f"Unsupported device type: {device_type}. Please select one of CPU, T4, V100, A100, L4, P4, P100, or H100"
            )

        if device_type.lower() == "cpu":
            with open(device_file, "w") as f:
                f.write(device_type.lower())

            with open(num_file, "w") as f:
                f.write("0")

            click.echo(
                click.style(
                    f"✅ Device set to CPU, you are now disconnected from any GPUs.",
                    fg="green",
                )
            )
            return

        with open(device_file, "w") as f:
            f.write(device_type.lower())

        if ngpus is None:
            ngpus = 1

        with open(num_file, "w") as f:
            f.write(str(ngpus))

        click.echo(
            click.style(f"✅ Device set to {ngpus} x {device_type.upper()}", fg="green")
        )

else:

    @cli.command(hidden=True)
    @click.argument("args", nargs=-1, type=click.UNPROCESSED)
    def run(args):
        raise click.ClickException(
            "tnr run is only supported on Linux based systems. Please switch to a linux machine or create a Thunder Compute instance with 'tnr create' and connect using 'tnr connect <INSTANCE ID>'"
        )

    @cli.command(hidden=True)
    @click.argument("device_type", required=False)
    @click.option("-n", "--ngpus", type=int, help="Number of GPUs to use")
    @click.option("--raw", is_flag=True, help="Output raw device information")
    def device(device_type, ngpus, raw):
        raise click.ClickException(
            "tnr device is only supported on Linux based systems. Please switch to a linux machine or create a Thunder Compute instance with 'tnr create' and connect using 'tnr connect <INSTANCE ID>'"
        )


@cli.command(hidden=True)
@click.argument("args", nargs=-1, type=click.UNPROCESSED)
def launch(args):
    return run(args)


@cli.command(
    help="Gets the status of all Thunder Compute instances and GPUs that are associated with the authenticated user."
)
def status():
    token = get_token()
    success, error, instances = utils.get_instances(token)
    if not success:
        raise click.ClickException(f"Failed to list thunder compute instance: {error}")

    instances_table = Table(
        title="Thunder Compute Instances",
        title_style="bold cyan",
        title_justify="left",
        box=box.ROUNDED,
    )

    instances_table.add_column(
        "ID",
        justify="center",
    )
    instances_table.add_column(
        "Status",
        justify="center",
    )
    instances_table.add_column(
        "Address",
        justify="center",
    )
    instances_table.add_column(
        "Creation Date",
        justify="center",
    )

    for instance_id, metadata in instances.items():
        if metadata["status"] == "RUNNING":
            status_color = "green"
        elif metadata["status"] == "STOPPED":
            status_color = "red"
        else:
            status_color = "yellow"

        if metadata["ip"] is None or len(metadata["ip"]) == 0:
            ip_entry = "--"
        else:
            ip_entry = metadata["ip"]

        instances_table.add_row(
            instance_id,
            Text(
                metadata["status"],
                style="white" if CURRENT_IP == ip_entry else status_color,
            ),
            ip_entry,
            metadata["createdAt"],
            style="bold white on #1a6b2f" if CURRENT_IP == ip_entry else "",
        )

    if len(instances) == 0:
        instances_table.add_row("--", "--", "--", "--")

    gpus_table = Table(
        title="Active GPU Processes", title_style="bold cyan", box=box.ROUNDED
    )

    gpus_table.add_column("GPU Type", justify="center")
    gpus_table.add_column("Duration", justify="center")

    active_sessions = utils.get_active_sessions(token)

    for session_info in active_sessions:
        gpus_table.add_row(
            f'{session_info["count"]} x {session_info["gpu"]}',
            f"{session_info['duration']}s",
        )

    if len(active_sessions) == 0:
        gpus_table.add_row("--", "--")

    console = Console()
    console.print(instances_table)
    print()
    console.print(gpus_table)


@cli.command(
    help="Create a new Thunder Compute instance",
    hidden=INSIDE_INSTANCE,
)
def create():
    token = get_token()
    success, error = utils.create_instance(token)
    if success:
        click.echo(
            click.style(
                f"✅ Successfully created thunder compute instance!", fg="green"
            )
        )
    else:
        raise click.ClickException(
            f"Failed to create thunder compute instance: {error}"
        )


@cli.command(
    help="Deletes a Thunder Compute instance. This action is not reversible.",
    hidden=INSIDE_INSTANCE,
)
@click.argument("instance_id", required=True)
def delete(instance_id):
    token = get_token()
    success, error = utils.delete_instance(instance_id, token)
    if success:
        click.echo(
            click.style(
                f"✅ Successfully deleted thunder compute instance {instance_id}",
                fg="green",
            )
        )

        utils.remove_instance_from_ssh_config(f"tnr-{instance_id}")
    else:
        raise click.ClickException(
            f"Failed to delete thunder compute instance {instance_id}: {error}"
        )


@cli.command(
    help="Starts a stopped Thunder Compute instance",
    hidden=INSIDE_INSTANCE,
)
@click.argument("instance_id", required=True)
def start(instance_id):
    token = get_token()
    success, error = utils.start_instance(instance_id, token)
    if success:
        click.echo(
            click.style(
                f"✅ Successfully started thunder compute instance {instance_id}",
                fg="green",
            )
        )
    else:
        raise click.ClickException(
            f"Failed to start thunder compute instance {instance_id}: {error}"
        )


@cli.command(
    help="Stops a running Thunder Compute instance. Stopped instances have persistent storage and can be restarted at any time.",
    hidden=INSIDE_INSTANCE,
)
@click.argument("instance_id", required=True)
def stop(instance_id):
    token = get_token()
    success, error = utils.stop_instance(instance_id, token)
    if success:
        click.echo(
            click.style(
                f"✅ Successfully stopped thunder compute instance {instance_id}",
                fg="green",
            )
        )
    else:
        raise click.ClickException(
            f"Failed to stop thunder compute instance {instance_id}: {error}"
        )


@cli.command(
    help="Connects to the Thunder Compute instance with the specified instance_ID.",
    hidden=INSIDE_INSTANCE,
)
@click.argument("instance_id", required=True)
def connect(instance_id):
    token = get_token()
    success, error, instances = utils.get_instances(token)
    if not success:
        raise click.ClickException(f"Failed to list thunder compute instance: {error}")

    for curr_instance_id, metadata in instances.items():
        if curr_instance_id == instance_id:
            if metadata["ip"] == None:
                raise click.ClickException(
                    f"Instance {instance_id} is not available to connect"
                )

            ip = metadata["ip"]

            if ip is None or ip == "":
                raise click.ClickException(
                    f"Unable to connect to instance {instance_id}, is the instance running?"
                )

            keyfile = utils.get_key_file(metadata["uuid"])
            if not os.path.exists(keyfile):
                if not utils.add_key_to_instance(instance_id, token):
                    raise click.ClickException(
                        f"Unable to find or create ssh key file for instance {instance_id}"
                    )

            start_time = time.time()
            connection_successful = False
            while start_time + 60 > time.time():
                ssh = paramiko.SSHClient()
                ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
                try:
                    ssh.connect(ip, username="ubuntu", key_filename=keyfile, timeout=10)
                    connection_successful = True
                    break
                except Exception as e:
                    continue

            if not connection_successful:
                raise click.ClickException(
                    "Failed to connect to the thunder compute instance within a minute. Please retry this command or contant the developers at support@thundercompute.com if this issue persists."
                )

            basedir = join(os.path.expanduser("~"), ".thunder")
            ssh.exec_command("mkdir -p ~/.thunder && chmod 700 ~/.thunder")
            stdin, stdout, stderr = ssh.exec_command("pip install --upgrade tnr")
            stdout.read().decode()  # Forces the command to get executed

            scp = SCPClient(ssh.get_transport())

            if os.path.exists(join(basedir, "token")):
                scp.put(join(basedir, "token"), remote_path="~/.thunder/token")

            # Update ~/.ssh/config
            host_alias = f"tnr-{instance_id}"
            exists, prev_ip = utils.get_ssh_config_entry(host_alias)
            if not exists:
                utils.add_instance_to_ssh_config(ip, keyfile, host_alias)
            else:
                if ip != prev_ip:
                    utils.update_ssh_config_ip(host_alias, ip)

            if platform.system() == "Windows":
                subprocess.run(
                    [
                        "ssh",
                        f"ubuntu@{ip}",
                        "-o",
                        "StrictHostKeyChecking=accept-new",
                        "-i",
                        rf"{keyfile}",
                        "-t",
                        f"{'export TNR_API_TOKEN=' + os.environ['TNR_API_TOKEN'] + ';' if 'TNR_API_TOKEN' in os.environ else ''} exec /home/ubuntu/.local/bin/tnr run /bin/bash",
                    ],
                    shell=True,
                )
            else:
                subprocess.run(
                    [
                        f"ssh ubuntu@{ip} -o StrictHostKeyChecking=accept-new -o UserKnownHostsFile=/dev/null -i {keyfile} -t '{'export TNR_API_TOKEN=' + os.environ['TNR_API_TOKEN'] + ';' if 'TNR_API_TOKEN' in os.environ else ''} exec /home/ubuntu/.local/bin/tnr run /bin/bash'"
                    ],
                    shell=True,
                )
            click.echo(click.style("⚡ Exiting thunder instance ⚡", fg="cyan"))
            return

    raise click.ClickException(f"Unable to find instance {instance_id}")


@cli.command(help="Copy data from local machine to a thunder instance")
@click.argument("src", required=True)
@click.argument("dst", required=True)
def scp(src, dst):
    token = get_token()
    success, error, instances = utils.get_instances(token)
    if not success:
        raise click.ClickException(f"Failed to list thunder compute instance: {error}")

    # Determine direction and instance
    src_split = src.split(":")
    dst_split = dst.split(":")
    if len(src_split) > 1:
        src_candidate_instance = src_split[0]
    else:
        src_candidate_instance = None

    if len(dst_split) > 1:
        dst_candidate_instance = dst_split[0]
    else:
        dst_candidate_instance = None

    for instance_id, metadata in instances.items():
        if (
            instance_id == src_candidate_instance
            or instance_id == dst_candidate_instance
        ):
            local_to_remote = True if instance_id == dst_candidate_instance else False
            local_path = src if local_to_remote else dst
            remote_path = dst_split[1] if local_to_remote else src_split[1]
            if remote_path == "":
                remote_path = "~/"

            if metadata["ip"] == None:
                raise click.ClickException(
                    f"Instance {instance_id} is not available to connect"
                )

            ip = metadata["ip"]

            if ip is None or ip == "":
                raise click.ClickException(
                    f"Unable to connect to instance {instance_id}, is the instance running?"
                )

            keyfile = utils.get_key_file(metadata["uuid"])
            if not os.path.exists(keyfile):
                if not utils.add_key_to_instance(instance_id, token):
                    raise click.ClickException(
                        f"Unable to find or create ssh key file for instance {instance_id}"
                    )

            start_time = time.time()
            connection_successful = False
            while start_time + 60 > time.time():
                ssh = paramiko.SSHClient()
                ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
                try:
                    ssh.connect(ip, username="ubuntu", key_filename=keyfile, timeout=10)
                    connection_successful = True
                    break
                except Exception as e:
                    continue

            if not connection_successful:
                raise click.ClickException(
                    "Failed to connect to the thunder compute instance within a minute. Please retry this command or contant the developers at support@thundercompute.com if this issue persists."
                )

            if local_to_remote:
                waiting_text = f"Copying {local_path} to {remote_path} on remote instance {instance_id}"
            else:
                waiting_text = (
                    f"Copying {remote_path} from instance {instance_id} to {local_path}"
                )

            with yaspin(text=waiting_text, color="blue") as spinner:
                transport = ssh.get_transport()
                transport.use_compression(True)

                with SCPClient(transport) as scp:
                    if local_to_remote:
                        scp.put(local_path, remote_path, recursive=True)
                    else:
                        scp.get(remote_path, local_path, recursive=True)

                spinner.ok("✅")
            return

    raise click.ClickException(
        f"Unable to find instance to copy to/from in {src} or {dst}"
    )


@cli.command(
    help="Logs in to Thunder Compute with a token generated at console.thundercompute.com. Saves the login token to ~/.thunder/token",
    hidden=INSIDE_INSTANCE,
)
def login():
    auth.login()


@cli.command(
    help="Logs out of Thunder Compute and deletes the saved login token",
    hidden=INSIDE_INSTANCE,
)
def logout():
    auth.logout()


def does_meet_min_required_version():
    try:
        current_version = version(PACKAGE_NAME)
        response = requests.get(
            f"https://api.thundercompute.com:8443/min_version", timeout=10
        )
        json_data = response.json() if response else {}
        min_version = json_data.get("version")
        if version_parser.parse(current_version) < version_parser.parse(min_version):
            return False, (current_version, min_version)

        return True, None

    except Exception as e:
        click.echo(
            click.style(
                f"Warning: Failed to fetch minimum required tnr version, ",
                fg="yellow",
            )
        )
        return True, None


if __name__ == "__main__":
    cli()
