import io
import os
import pathlib
import re
import sys
import tempfile
import unittest.mock
import urllib.parse
from typing import Any, Dict, Mapping, Optional, Set, Union

import requests
from expandvars import UnboundVariable, expand

from datahub.configuration.common import ConfigurationError, ConfigurationMechanism
from datahub.configuration.json_loader import JsonConfigurationMechanism
from datahub.configuration.toml import TomlConfigurationMechanism
from datahub.configuration.yaml import YamlConfigurationMechanism
from datahub.masking.secret_registry import SecretRegistry, is_masking_enabled

Environ = Mapping[str, str]


def _extract_env_var_names(text: str) -> Set[str]:
    """Extract environment variable names from ${VAR} or $VAR patterns."""
    var_names = set()

    # Match ${VAR} and bash parameter expansion patterns
    # Pattern breakdown:
    #   ([A-Za-z_][A-Za-z0-9_]*)  - capture variable name
    #   (?::[+\-=?][^}]*)?        - optional bash operators with content
    for match in re.finditer(r"\$\{([A-Za-z_][A-Za-z0-9_]*)(?::[+\-=?][^}]*)?\}", text):
        var_names.add(match.group(1))

    # Match $VAR patterns (without braces)
    for match in re.finditer(r"\$([A-Za-z_][A-Za-z0-9_]*)", text):
        var_name = match.group(1)
        # Only add if not already captured by ${} pattern
        if var_name not in var_names:
            var_names.add(var_name)

    return var_names


def resolve_env_variables(config: dict, environ: Environ) -> dict:
    # TODO: This is kept around for backwards compatibility.
    return EnvResolver(environ).resolve(config)


def list_referenced_env_variables(config: dict) -> Set[str]:
    # TODO: This is kept around for backwards compatibility.
    return EnvResolver(environ=os.environ).list_referenced_variables(config)


class EnvResolver:
    """Resolves environment variable references in configuration dictionaries."""

    def __init__(
        self,
        environ: Environ,
        strict_env_syntax: bool = False,
        register_secrets: bool = True,
    ):
        """
        Initialize the environment variable resolver.

        Args:
            environ: Environment variable mapping (os.environ, custom dict, or external secrets)
            strict_env_syntax: If True, only match ${VAR} syntax (not $VAR)
            register_secrets: If True, register resolved values with masking registry
                            (only if masking is globally enabled)
        """
        self.environ = environ
        self.strict_env_syntax = strict_env_syntax
        self.register_secrets = register_secrets

    def resolve(self, config: dict) -> dict:
        return self._resolve_dict(config)

    @classmethod
    def list_referenced_variables(
        cls,
        config: dict,
        strict_env_syntax: bool = False,
    ) -> Set[str]:
        # This is a bit of a hack, but expandvars does a bunch of escaping
        # and other logic that we don't want to duplicate here.

        vars = set()

        def mock_get_env(key: str, default: Optional[str] = None) -> str:
            vars.add(key)
            if default is not None:
                return default
            return "mocked_value"

        mock = unittest.mock.MagicMock()
        mock.get.side_effect = mock_get_env

        resolver = EnvResolver(
            environ=mock, strict_env_syntax=strict_env_syntax, register_secrets=False
        )
        resolver._resolve_dict(config)

        return vars

    def _register_env_vars_from_element(self, element: str) -> None:
        """Register environment variables found in element for secret masking."""
        if not self.register_secrets or not is_masking_enabled():
            return

        # Extract variable names from the pattern
        var_names = _extract_env_var_names(element)

        # Collect secrets for batch registration
        # We register all ${VAR} references from recipe as secrets
        secrets = {}
        for var_name in var_names:
            value = self.environ.get(var_name)
            if value:
                secrets[var_name] = value

        # Batch register all secrets from this element
        if secrets:
            SecretRegistry.get_instance().register_secrets_batch(secrets)

    def _resolve_element(self, element: str) -> str:
        if re.search(r"(\$\{).+(\})", element):
            # Register secrets before expansion
            self._register_env_vars_from_element(element)
            return expand(element, nounset=True, environ=self.environ)
        elif not self.strict_env_syntax and element.startswith("$"):
            try:
                # Register secrets before expansion
                self._register_env_vars_from_element(element)
                return expand(element, nounset=True, environ=self.environ)
            except UnboundVariable:
                # TODO: This fallback is kept around for backwards compatibility, but
                # doesn't make a ton of sense from first principles.
                return element
        else:
            return element

    def _resolve_list(self, ele_list: list) -> list:
        new_v: list = []
        for ele in ele_list:
            if isinstance(ele, str):
                new_v.append(self._resolve_element(ele))
            elif isinstance(ele, list):
                new_v.append(self._resolve_list(ele))
            elif isinstance(ele, dict):
                new_v.append(self._resolve_dict(ele))
            else:
                new_v.append(ele)
        return new_v

    def _resolve_dict(self, config: dict) -> dict:
        new_dict: Dict[Any, Any] = {}
        for k, v in config.items():
            if isinstance(v, dict):
                new_dict[k] = self._resolve_dict(v)
            elif isinstance(v, list):
                new_dict[k] = self._resolve_list(v)
            elif isinstance(v, str):
                new_dict[k] = self._resolve_element(v)
            else:
                new_dict[k] = v
        return new_dict


WRITE_TO_FILE_DIRECTIVE_PREFIX = "__DATAHUB_TO_FILE_"


def _process_directives(config: dict) -> dict:
    def _process(obj: Any) -> Any:
        if isinstance(obj, dict):
            new_obj = {}
            for k, v in obj.items():
                if isinstance(k, str) and k.startswith(WRITE_TO_FILE_DIRECTIVE_PREFIX):
                    # This writes the value to a temporary file and replaces the value with the path to the file.
                    config_option = k[len(WRITE_TO_FILE_DIRECTIVE_PREFIX) :]

                    with tempfile.NamedTemporaryFile("w", delete=False) as f:
                        filepath = f.name
                        f.write(v)

                    new_obj[config_option] = filepath
                else:
                    new_obj[k] = _process(v)

            return new_obj
        else:
            return obj

    return _process(config)


def load_config_file(
    config_file: Union[str, pathlib.Path],
    squirrel_original_config: bool = False,
    squirrel_field: str = "__orig_config",
    allow_stdin: bool = False,
    allow_remote: bool = True,  # TODO: Change the default to False.
    resolve_env_vars: bool = True,  # TODO: Change the default to False.
    process_directives: bool = False,
    extra_env_vars: Optional[Dict[str, str]] = None,
) -> dict:
    config_mech: ConfigurationMechanism
    if allow_stdin and config_file == "-":
        # Reading from stdin. Supports two formats:
        # 1. Plain YAML/JSON recipe (backward compatible)
        # 2. JSON envelope: {"__recipe_yaml__": "...", "__secrets__": {...}}
        #    Secrets are merged into the env resolver without touching os.environ.
        config_mech = YamlConfigurationMechanism()
        raw_stdin = sys.stdin.read()

        try:
            import json as _json

            envelope = _json.loads(raw_stdin)
            if isinstance(envelope, dict) and "__recipe_yaml__" in envelope:
                raw_config_file = envelope["__recipe_yaml__"]
                stdin_secrets = envelope.get("__secrets__", {})
                if stdin_secrets:
                    extra_env_vars = {**(extra_env_vars or {}), **stdin_secrets}
            else:
                # Plain JSON (which is valid YAML) — treat as recipe
                raw_config_file = raw_stdin
        except (ValueError, _json.JSONDecodeError):
            # Not JSON — treat as plain YAML
            raw_config_file = raw_stdin
    else:
        config_file_path = pathlib.Path(config_file)
        if config_file_path.suffix in {".yaml", ".yml"}:
            config_mech = YamlConfigurationMechanism()
        elif config_file_path.suffix == ".json":
            config_mech = JsonConfigurationMechanism()
        elif config_file_path.suffix == ".toml":
            config_mech = TomlConfigurationMechanism()
        else:
            raise ConfigurationError(
                f"Only .toml, .yml, and .json are supported. Cannot process file type {config_file_path.suffix}"
            )

        url_parsed = urllib.parse.urlparse(str(config_file))
        if allow_remote and url_parsed.scheme in (
            "http",
            "https",
        ):  # URLs will return http/https
            # If the URL is remote, we need to fetch it.
            try:
                response = requests.get(
                    str(config_file),
                    auth=(
                        urllib.parse.unquote(url_parsed.username or ""),
                        urllib.parse.unquote(url_parsed.password or ""),
                    )
                    if url_parsed.username or url_parsed.password
                    else None,
                )
                response.raise_for_status()
                raw_config_file = response.text
            except Exception as e:
                raise ConfigurationError(
                    f"Cannot read remote file {config_file_path}: {e}"
                ) from e
        else:
            if not config_file_path.is_file():
                raise ConfigurationError(
                    f"Cannot open config file {config_file_path.resolve()}"
                )
            raw_config_file = config_file_path.read_text()

    config_fp = io.StringIO(raw_config_file)
    raw_config = config_mech.load_config(config_fp)

    config = raw_config.copy()
    if resolve_env_vars:
        environ: Environ = os.environ
        if extra_env_vars:
            environ = {**os.environ, **extra_env_vars}
        config = EnvResolver(environ=environ).resolve(config)
    if process_directives:
        config = _process_directives(config)

    if squirrel_original_config:
        config[squirrel_field] = raw_config
    return config
