import json
import logging
import os
import tempfile
from collections import OrderedDict
from contextlib import contextmanager

from jinja2 import Environment, FileSystemLoader, select_autoescape
from python_terraform import Terraform

from dvc.exceptions import DvcException

logger = logging.getLogger(__name__)

TEMPLATE_DIR = os.path.join(os.path.dirname(__file__), "templates")
TEMPLATE_MAIN = "main.tf"

BASE_CONFIG = {
    "//": "This file auto-generated by dvc-tpi, do not edit manually.",
    "terraform": {
        "required_providers": {
            "iterative": {
                "source": "iterative/iterative",
            },
        },
    },
    "provider": {
        "iterative": {},
    },
}

_jinja_env = Environment(
    loader=FileSystemLoader(TEMPLATE_DIR),
    autoescape=select_autoescape(),
)


class TerraformError(DvcException):
    pass


class DvcTerraform(Terraform):
    def cmd(self, *args, **kwargs):
        logger.debug(" ".join(self.generate_cmd_string(*args, **kwargs)))
        kwargs["capture_output"] = False
        ret, _stdout, _stderr = super().cmd(*args, **kwargs)
        if ret != 0:
            raise TerraformError("Cmd 'terraform {cmd}' failed")

    def iter_instances(self, name: str):
        """Iterate over active iterative_machine instances."""
        self.read_state_file()
        resources = getattr(self.tfstate, "resources", [])
        for resource in resources:
            if (
                resource.get("type") == "iterative_machine"
                and resource.get("name") == name
            ):
                yield from (
                    instance.get("attributes", {})
                    for instance in resource.get("instances", [])
                )

    @staticmethod
    @contextmanager
    def pemfile(resource: dict):
        """Generate a temporary SSH key (PEM) file for the specified resource.

        The path to the file will be yielded, and the file will be removed
        upon exiting the yielded context.
        """
        with tempfile.NamedTemporaryFile(delete=False) as fobj:
            fobj.write(resource["ssh_private"].encode("ascii"))
        try:
            yield fobj.name
        finally:
            os.unlink(fobj.name)


def render_config(**config) -> str:
    """Render TPI iterative_machine config as terraform HCL."""
    assert "name" in config and "cloud" in config
    template = _jinja_env.get_template(TEMPLATE_MAIN)
    return template.render(**config)


def render_json(indent=None, **config) -> str:
    """Render TPI iterative_machine config as terraform JSON."""
    assert "name" in config and "cloud" in config
    tf_config = OrderedDict(BASE_CONFIG)
    name = config["name"]
    resource = {
        name: {key: value for key, value in config.items()},
    }
    tf_config["resource"] = {"iterative_machine": resource}
    return json.dumps(tf_config, indent=indent)
