"""ptc infra get-terraform / ptc infra terraform — terraform operations."""

import io
import json
import shutil
import subprocess
import sys
import tarfile
import tempfile
from contextlib import contextmanager
from pathlib import Path
from typing import Iterator, Optional, Union

import srsly
from radicli import Arg
from wasabi import msg

from ...cli import cli
from ...cluster_config import ClusterConfig

CLOUD_PROVIDERS = ("gcp", "k3s", "aws")


# ---------------------------------------------------------------------------
# Download terraform files from PAM
# ---------------------------------------------------------------------------


@cli.subcommand(
    "infra",
    "get-terraform",
    dest=Arg("--dest", help="Directory to save terraform files to"),
    cloud_provider=Arg(
        "--cloud-provider",
        help=f"Cloud provider: {', '.join(CLOUD_PROVIDERS)}",
    ),
)
def get_terraform(
    dest: Union[Path, str] = ".",
    cloud_provider: str = "gcp",
) -> None:
    """Download terraform files for cluster infrastructure."""
    from .._state import get_auth_state

    if isinstance(dest, str):
        dest = Path(dest)
    auth = get_auth_state()
    dest.mkdir(exist_ok=True, parents=True)
    tar_stream = auth.client.broker.download_gcloud_files()
    tar_bytes = tar_stream.read()
    with tarfile.open(fileobj=io.BytesIO(tar_bytes)) as tf:
        tf.extractall(dest)
    cluster_dir = dest / "cluster"
    if cluster_dir.exists():
        provider_src = cluster_dir / cloud_provider
        provider_dest = dest / cloud_provider
        if provider_src.exists() and not provider_dest.exists():
            shutil.copytree(provider_src, provider_dest)
        shutil.rmtree(cluster_dir)
    msg.good(f"Terraform files saved to {dest / cloud_provider}")


# ---------------------------------------------------------------------------
# Apply terraform (init / plan / apply)
# ---------------------------------------------------------------------------


@cli.subcommand(
    "infra",
    "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.",
    ),
    yes_tf_init=Arg(
        "--tf-init",
        help="Force run 'terraform init'",
    ),
    yes_tf_apply=Arg(
        "--tf-apply",
        help="Apply the terraform changes",
    ),
    yes_tf_auto_accept=Arg(
        "--tf-auto-accept", help="Accept the terraform changes without interaction"
    ),
)
def apply_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:
    """Apply terraform to create or update cluster infrastructure."""
    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 = ClusterConfig(**config_dict)
    tf_dir = src / "gcp"
    if yes_tf_init:
        _terraform_init(config, tf_dir)
    if yes_tf_init or yes_tf_apply:
        with _write_terraform_vars(config) as var_file:
            plan_file = "plan.tfplan"
            cmd = [
                "terraform",
                "plan",
                "-input=false",
                "-out",
                plan_file,
                f"-var-file={var_file}",
            ]
            _ = subprocess.run(cmd, check=True, text=True, cwd=tf_dir)
            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=tf_dir)


def _terraform_init(config: ClusterConfig, 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)


@contextmanager
def _write_terraform_vars(config: ClusterConfig) -> Iterator[str]:
    """Write terraform variables to a temp file, yielding its path."""
    data = {
        "gcp_project": config.deployment.gcp.project,
        "gcp_zone": config.deployment.gcp.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: w.to_terraform_var() for w in config.node_pools},
    }
    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".tfvars.json", delete=True
    ) as tmp:
        tmp.write(json.dumps(data, indent=2))
        tmp.flush()
        yield tmp.name
