import io
import json
import shutil
import subprocess
import sys
import tarfile
from pathlib import Path
from typing import Optional, Union

import srsly
from radicli import Arg
from wasabi import msg

from ..cli import cli
from ..cloud import gcp as gcp_utils
from ._state import get_auth_state
from .clusters import Config


@cli.subcommand(
    "gcp",
    "create-project",
    project_id=Arg(help="The unique string name of the project"),
    project_name=Arg(
        help="The human-readable name of the project",
    ),
    no_create_ip=Arg("--no-create-ip", help="Don't create an IP address"),
    owners=Arg("--owners", help="Email addresses of users to add as owners"),
)
def create_project(
    project_id: str,
    project_name: str,
    no_create_ip: bool,
    owners: Optional[str] = None,
) -> None:
    """
    Prepare everything to deploy a Prodigy Teams Cluster.
    """
    service_account_name = "prodigy-cluster-sa"
    service_account_display_name = "Prodigy Teams service account"
    ip_name = f"{project_id}-web"
    create_tf_bucket = True
    gcp_utils.check_gcloud_exists()
    msg.info("Initializing...")
    project_name = gcp_utils.create_project(
        project_id=project_id, project_name=project_name, exist_ok=True
    )
    gcp_utils.create_service_account(
        project_name=project_name,
        service_account_name=service_account_name,
        display_name=service_account_display_name,
        exist_ok=True,
    )
    gcp_utils.add_service_account_member(
        project_name=project_name, service_account_name=service_account_name
    )
    gcp_utils.setup_billing(project_name=project_name)
    gcp_utils.enable_apis(project_name=project_name)
    if owners:
        gcp_utils.add_owners(project_id=project_id, owners=owners.split(","))
    if create_tf_bucket:
        gcp_utils.create_backend_bucket(project_id=project_id)
    if not no_create_ip:
        gcp_utils.create_ip_address(project_id, ip_name or f"{project_id}-web")


@cli.subcommand(
    "gcp",
    "clone-files",
    dest=Arg(help="Path to save cluster definition files to"),
    name=Arg("--name", help="The name of your cluster"),
    project=Arg("--project", help="The GCP project you're using"),
    domain=Arg("--domain", help="The domain to access your cluster"),
    zone=Arg("--zone", help="The GCP zone for your cluster"),
)
def init_gcloud_cluster(
    dest: Path, name: str, project: str, domain: str, zone: str
) -> None:
    """Download Terraform and Nomad configuration files to launch a cluster.

    After you've downloaded the files and updated the config, use 'ptc clusters launch'
    to create the infrastructure.
    """
    auth = get_auth_state()
    dest.mkdir(exist_ok=True, parents=True)
    tar_stream = auth.client.broker.download_gcloud_files()
    # Can't stream directly to the tar file. Need to download first,
    # then wrap with io.BytesIO
    tar_bytes = tar_stream.read()
    with tarfile.open(fileobj=io.BytesIO(tar_bytes)) as file_:
        file_.extractall(dest)
    if (dest / "gcp").exists():
        shutil.rmtree(dest / "gcp")
    if (dest / "tasks").exists():
        shutil.rmtree(dest / "tasks")
    shutil.copytree(dest / "cluster" / "gcp", dest / "gcp")
    shutil.copytree(dest / "cluster" / "tasks", dest / "tasks")
    shutil.rmtree(dest / "cluster")
    # 2. Write config
    config = Config(
        name=name,
        terraform=Config.Terraform(
            gcs=Config.Terraform.GCS(bucket=f"{project}-terraform", prefix="")
        ),
        deployment=Config.Deployment(
            project=project,
            zone=zone,
            domain=domain,
            pam_domain=auth.pam_url.netloc,
            container_registry="",
        ),
        node_pools=[
            Config.NodePool(
                name="base",
                node_class="base",
                machine_type="n1-standard-1",
                target_size=1,
                gpu=None,
                per_job={},
            ),
            Config.NodePool(
                name="small",
                node_class="small",
                machine_type="g1-small",
                target_size=1,
                gpu=None,
                per_job={},
            ),
            Config.NodePool(
                name="medium",
                node_class="medium",
                machine_type="n1-standard-1",
                target_size=1,
                gpu=None,
                per_job={},
            ),
            Config.NodePool(
                name="gpu",
                node_class="gpu",
                machine_type="n1-standard-1",
                target_size=1,
                gpu=None,
                per_job={},
            ),
        ],
    )
    srsly.write_yaml(dest / "config.yaml", config.model_dump())


@cli.subcommand(
    "gcp",
    "terraform",
    src=Arg(help="Path to cluster definition files"),
    config_path=Arg(help="Path to config yaml"),
    config_key=Arg(
        "--config-key",
        help="Top-level key for the cluster config. Allows multiple cluster configurations to be kept in one file. If not set, the whole config file is used.",
    ),
    yes_tf_init=Arg(
        "--tf-init",
        help="Force run 'terraform init'. Required if you change the backend configuration",
    ),
    yes_tf_apply=Arg(
        "--tf-apply",
        help="Apply the terraform changes. Will print a plan and prompt you for acceptance",
    ),
    yes_tf_auto_accept=Arg(
        "--tf-auto-accept", help="Accept the terraform changes without interaction"
    ),
)
def gcp_terraform(
    src: Union[Path, str] = ".",
    config_path: Optional[Path] = None,
    config_key: Optional[str] = None,
    yes_tf_init: bool = False,
    yes_tf_apply: bool = False,
    yes_tf_auto_accept: bool = False,
) -> None:
    """Launch cluster infrastructure under your own cloud account.
    Use 'ptc gcp clone-files' to get the definition files.

    This command requires Hashicorp Terraform.
    """
    if src == ".":
        src = Path.cwd()
    elif isinstance(src, str):
        src = Path(src)
    if config_path is None:
        config_path = src / "config.yaml"
    config_dict = srsly.read_yaml(config_path)
    assert isinstance(config_dict, dict)
    if config_key is not None:
        config_dict = config_dict[config_key]
    config = Config(**config_dict)
    if yes_tf_init:
        _terraform_init(config, src / "gcp")
    if yes_tf_init or yes_tf_apply:
        _write_terraform_vars(config, src / "gcp" / "terraform.auto.tfvars.json")
        plan_file = "plan.tfplan"
        cmd = ["terraform", "plan", "-input=false", "-out", plan_file]
        _ = subprocess.run(cmd, check=True, text=True, cwd=src / "gcp")
        if not yes_tf_auto_accept:
            from ..ui import isatty

            if not isatty(sys.stdin):
                from ..errors import CLIError

                raise CLIError(
                    "Cannot prompt for confirmation in non-interactive mode",
                    "Use --tf-auto-accept to apply without confirmation",
                )
            accept_plan = input("Apply plan? 'yes' to accept\n")
            if accept_plan.strip() != "yes":
                return
        cmd = ["terraform", "apply", "-auto-approve", plan_file]
        _ = subprocess.run(cmd, check=True, text=True, cwd=src / "gcp")


def _terraform_init(config: Config, path: Path) -> None:
    if config.terraform.terraform_cloud:
        data = {
            "terraform": {
                "backend": {
                    "remote": {
                        "workspaces": {
                            "name": config.terraform.terraform_cloud.workspace
                        },
                        "hostname": "app.terraform.io",
                        "organization": config.terraform.terraform_cloud.organization,
                    }
                }
            }
        }
    elif config.terraform.gcs:
        data = {
            "terraform": {
                "backend": {
                    "gcs": {
                        "bucket": config.terraform.gcs.bucket,
                        "prefix": config.terraform.gcs.prefix,
                    }
                }
            }
        }
    else:
        data = {"terraform": {}}
    with (path / "backend.tf.json").open("w", encoding="utf8") as file_:
        file_.write(json.dumps(data, indent=2))
    cmd = ["terraform", "init", "-reconfigure", "-input=false"]
    _ = subprocess.run(cmd, check=True, text=True, cwd=path)


def _write_terraform_vars(config: Config, path: Path) -> None:
    data = {
        "gcp_project": config.deployment.project,
        "gcp_zone": config.deployment.zone,
        "domain": config.deployment.domain,
        "enable_ssh": True,
        "system_node_pool_machine_type": config.deployment.system_node_pool_machine_type,
        "system_node_pool_size": str(config.deployment.system_node_pool_size),
        "worker_types": {
            w.name: {
                "name": w.name,
                "node_class": w.node_class,
                "node_pool": "default",
                "machine_type": w.machine_type,
                "target_size": w.target_size,
                "preemptible": w.preemptible,
                "spot": w.spot,
                "min_size": w.min_size if w.min_size is not None else 0,
                "max_size": w.max_size
                if w.max_size is not None
                else max(w.target_size, 1),
                "volumes": [
                    {
                        "name": "nfs",
                        "source": "/mnt/nfs/",
                        "read_only": False,
                    },
                ],
                "guest_accelerator": w.gpu.model_dump() if w.gpu else None,
            }
            for w in config.node_pools
        },
    }
    with path.open("w", encoding="utf8") as file_:
        file_.write(json.dumps(data, indent=2))
