from __future__ import annotations

import json
import shlex
import subprocess
import tempfile
import os
from typing import Any, Dict, List, Optional


class AzCliError(RuntimeError):
    pass


class AzCliBackend:
    """Minimal wrapper around `az rest` to perform policy operations via the Azure CLI.

    This backend prefers `az rest` to call ARM endpoints because it allows passing
    arbitrary JSON payloads (safer for complex properties). It requires `az`
    being installed and authenticated.
    """

    def __init__(self, subscription_id: Optional[str] = None):
        self.subscription_id = subscription_id

    def _run_az(self, args: List[str]) -> Dict[str, Any]:
        cmd = ["az"] + args + ["--output", "json"]
        try:
            proc = subprocess.run(
                cmd,
                check=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
            )
            out = proc.stdout.strip()
            if not out:
                return {}
            try:
                return json.loads(out)
            except Exception:
                return {"raw": out}
        except subprocess.CalledProcessError as ex:
            raise AzCliError(
                f"az command failed: {ex.stderr.strip() or ex.stdout.strip()}"
            )

    def _az_rest(self, method: str, uri: str, body: Optional[Dict] = None) -> Any:
        args = ["rest", "--method", method, "--uri", uri]
        if body is not None:
            # write temp file and pass with --body @file. Ensure the temp file is removed after use.
            tf = tempfile.NamedTemporaryFile(
                "w+", suffix=".json", delete=False, encoding="utf-8"
            )
            try:
                try:
                    json.dump(body, tf, indent=2)
                    tf.flush()
                    tf_name = tf.name
                finally:
                    tf.close()

                args += ["--body", f"@{tf_name}"]
                result = self._run_az(args)
                return result
            finally:
                # attempt to remove the temporary file regardless of success or failure
                try:
                    if tf and os.path.exists(tf_name):
                        os.remove(tf_name)
                except Exception:
                    # best-effort cleanup; do not mask original exceptions
                    pass
        return self._run_az(args)

    # Policy definitions (list/get/create)
    def list_policy_definitions(self) -> List[Dict[str, Any]]:
        # Try subscription-scoped list when subscription provided, otherwise tenant-scoped
        if self.subscription_id:
            uri = f"https://management.azure.com/subscriptions/{self.subscription_id}/providers/Microsoft.Authorization/policyDefinitions?api-version=2021-06-01"
        else:
            uri = f"https://management.azure.com/providers/Microsoft.Authorization/policyDefinitions?api-version=2021-06-01"
        res = self._az_rest("GET", uri)
        # az rest returns dict with 'value' for list responses
        if isinstance(res, dict) and "value" in res:
            return res.get("value", [])
        if isinstance(res, list):
            return res
        return []

    def get_policy_definition(self, name: str) -> Dict[str, Any]:
        if self.subscription_id:
            uri = f"https://management.azure.com/subscriptions/{self.subscription_id}/providers/Microsoft.Authorization/policyDefinitions/{name}?api-version=2021-06-01"
        else:
            uri = f"https://management.azure.com/providers/Microsoft.Authorization/policyDefinitions/{name}?api-version=2021-06-01"
        return self._az_rest("GET", uri)

    def create_or_update_policy_definition(
        self, name: str, policy: Dict[str, Any]
    ) -> Dict[str, Any]:
        if self.subscription_id:
            uri = f"https://management.azure.com/subscriptions/{self.subscription_id}/providers/Microsoft.Authorization/policyDefinitions/{name}?api-version=2021-06-01"
        else:
            uri = f"https://management.azure.com/providers/Microsoft.Authorization/policyDefinitions/{name}?api-version=2021-06-01"
        return self._az_rest("PUT", uri, body=policy)

    # Assignments
    def list_assignments(self, scope: Optional[str] = None) -> List[Dict[str, Any]]:
        scope_path = scope or (
            f"/subscriptions/{self.subscription_id}" if self.subscription_id else ""
        )
        scope_path = scope_path.rstrip("/")
        if not scope_path:
            # list at tenant scope
            uri = f"https://management.azure.com/providers/Microsoft.Authorization/policyAssignments?api-version=2021-06-01"
        else:
            uri = f"https://management.azure.com{scope_path}/providers/Microsoft.Authorization/policyAssignments?api-version=2021-06-01"
        res = self._az_rest("GET", uri)
        if isinstance(res, dict) and "value" in res:
            return res.get("value", [])
        if isinstance(res, list):
            return res
        return []

    def create_or_update_assignment(
        self, name: str, assignment: Dict[str, Any], scope: Optional[str] = None
    ) -> Dict[str, Any]:
        scope_path = scope or (
            f"/subscriptions/{self.subscription_id}" if self.subscription_id else ""
        )
        if not scope_path:
            raise AzCliError(
                "scope must be provided when subscription id is not configured"
            )
        uri = f"https://management.azure.com{scope_path}/providers/Microsoft.Authorization/policyAssignments/{name}?api-version=2021-06-01"
        return self._az_rest("PUT", uri, body=assignment)

    def delete_assignment(
        self, name: str, scope: Optional[str] = None
    ) -> Dict[str, Any]:
        scope_path = scope or (
            f"/subscriptions/{self.subscription_id}" if self.subscription_id else ""
        )
        if not scope_path:
            raise AzCliError(
                "scope must be provided when subscription id is not configured"
            )
        uri = f"https://management.azure.com{scope_path}/providers/Microsoft.Authorization/policyAssignments/{name}?api-version=2021-06-01"
        return self._az_rest("DELETE", uri)

    # Initiatives (policy set definitions)
    def list_initiatives(self) -> List[Dict[str, Any]]:
        if self.subscription_id:
            uri = f"https://management.azure.com/subscriptions/{self.subscription_id}/providers/Microsoft.Authorization/policySetDefinitions?api-version=2021-06-01"
        else:
            uri = f"https://management.azure.com/providers/Microsoft.Authorization/policySetDefinitions?api-version=2021-06-01"
        res = self._az_rest("GET", uri)
        if isinstance(res, dict) and "value" in res:
            return res.get("value", [])
        if isinstance(res, list):
            return res
        return []

    def get_initiative(self, name: str) -> Dict[str, Any]:
        if self.subscription_id:
            uri = f"https://management.azure.com/subscriptions/{self.subscription_id}/providers/Microsoft.Authorization/policySetDefinitions/{name}?api-version=2021-06-01"
        else:
            uri = f"https://management.azure.com/providers/Microsoft.Authorization/policySetDefinitions/{name}?api-version=2021-06-01"
        return self._az_rest("GET", uri)

    def create_or_update_initiative(
        self, name: str, initiative: Dict[str, Any]
    ) -> Dict[str, Any]:
        if self.subscription_id:
            uri = f"https://management.azure.com/subscriptions/{self.subscription_id}/providers/Microsoft.Authorization/policySetDefinitions/{name}?api-version=2021-06-01"
        else:
            uri = f"https://management.azure.com/providers/Microsoft.Authorization/policySetDefinitions/{name}?api-version=2021-06-01"
        return self._az_rest("PUT", uri, body=initiative)
