# Copyright 2021 Acryl Data, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import json
import logging
import os
import re
import shutil
from collections.abc import Sequence
from typing import Any

import pydantic
import yaml
from datahub.masking.bootstrap import initialize_secret_masking
from datahub.masking.secret_registry import SecretRegistry

from acryl.executor.common.config import PermissiveConfigModel
from acryl.executor.context.execution_context import ExecutionContext
from acryl.executor.context.executor_context import ExecutorContext
from acryl.executor.execution import venv_utils

logger = logging.getLogger(__name__)


class SubProcessTaskUtil:
    MAX_LOG_LINES = 2000

    # The original value is 64kb (https://github.com/python/cpython/blob/7528e2c06c8baf809b56f406bcc50e8436c9647c/Lib/asyncio/streams.py#L23).
    # Increasing it should improve performance.
    SUBPROCESS_BUFFER_SIZE = 2**20  # 1mb

    # GMS / mysql has a 4mb limit on the size of a data packet.
    # Doing 90% of that so we have some buffer.
    MAX_LOG_SIZE_BYTES = int(0.9 * 2**22)  # 90% of 4mb

    # We want to truncate long lines so that we can show more lines in the logs.
    MAX_BYTES_PER_LINE = 2**12  # 4kb

    @staticmethod
    def _format_log_lines(lines: Sequence[str]) -> str:
        text = "".join(lines)

        # Python slices are super permissive on index bounds, so this works.
        text = text[-SubProcessTaskUtil.MAX_LOG_SIZE_BYTES :]

        if len(lines) >= SubProcessTaskUtil.MAX_LOG_LINES:
            # lines is a deque, so len(lines) won't be larger than MAX_LOG_LINES.
            text = f"[earlier logs truncated...]\n{text}"

        return text

    @staticmethod
    def _resolve_secrets(secret_names: list[str], ctx: ExecutorContext) -> dict:
        # Attempt to resolve secret using by checking each configured secret store.
        secret_stores = ctx.get_secret_stores()
        final_secret_values = dict({})

        for secret_store in secret_stores:
            try:
                # Retrieve secret values from the store.
                secret_values_dict = secret_store.get_secret_values(secret_names)
                # Overlay secret values from each store, if not None.
                for secret_name, secret_value in secret_values_dict.items():
                    if secret_value is not None:
                        final_secret_values[secret_name] = secret_value
            except Exception:
                logger.exception(
                    f"Failed to fetch secret values from secret store with id {secret_store.get_id()}"
                )
        return final_secret_values

    @staticmethod
    def _warn_on_bad_secret_value(ctx: ExecutionContext, key: str, val: str) -> None:
        # Log a warning if the value is a valid JSON document (dict or list)
        # to hint AWS Secret Manager users of a wrong type.
        # We only warn for complex structures, not simple scalar values like numbers.
        try:
            parsed = json.loads(val)
            # Only warn if it's a dict or list (actual JSON documents), not scalars
            if isinstance(parsed, (dict, list)):
                ctx.get_report().report_error(
                    f"Secret variable ${{{key}}} appears to contain a JSON document while string is expected. "
                    "If you are using AWS Secret Manager, make sure to pass secret as plain text and not as a key/value pair."
                )
        except Exception:
            pass

    @staticmethod
    def _resolve_recipe(
        recipe: str, execution_ctx: ExecutionContext, executor_ctx: ExecutorContext
    ) -> tuple[dict, list[str], set[str]]:
        # Now attempt to find and replace all secrets inside the recipe.
        secret_pattern = re.compile(r".*?\${(\w+)}.*?")

        resolved_recipe = recipe
        secret_matches = secret_pattern.findall(resolved_recipe)

        # 1. Extract all secrets needing resolved.
        secrets_to_resolve = []
        if secret_matches:
            for match in secret_matches:
                secrets_to_resolve.append(match)

        # 2. Resolve secret values
        secret_values_dict = SubProcessTaskUtil._resolve_secrets(
            secrets_to_resolve, executor_ctx
        )

        # 3. Fall back to os.environ for any secrets not found in secret stores
        # Track which secrets we need to add (vs. already present)
        secrets_to_add_to_environ = set()

        for secret_name in secrets_to_resolve:
            # Check if secret already exists in environment BEFORE we modify it
            was_already_in_environ = secret_name in os.environ

            if (
                secret_name not in secret_values_dict
                or secret_values_dict[secret_name] is None
            ):
                env_value = os.environ.get(secret_name)
                if env_value is not None:
                    logger.info(
                        f"Secret '{secret_name}' not found in secret stores, using value from environment variable"
                    )
                    secret_values_dict[secret_name] = env_value
                    # Don't add to cleanup - it was already in environment
                else:
                    logger.warning(
                        f"Secret '{secret_name}' not found in secret stores or environment, using empty string"
                    )
                    secret_values_dict[secret_name] = ""
                    secrets_to_add_to_environ.add(secret_name)
            else:
                # Secret came from secret store
                # Only add to cleanup if it wasn't already in environment
                if not was_already_in_environ:
                    secrets_to_add_to_environ.add(secret_name)

        # Set environment variables for SecretRegistry and DataHub CLI
        for name, value in secret_values_dict.items():
            if value is not None and name in secrets_to_add_to_environ:
                os.environ[name] = value

        # Ensure environment cleanup on error
        try:
            # Set up secret masking (must happen after env vars are set)
            if secrets_to_resolve:
                try:
                    initialize_secret_masking(force=True)
                    registry = SecretRegistry.get_instance()
                    for secret_name in secrets_to_resolve:
                        secret_value = os.environ.get(secret_name)
                        if secret_value:
                            registry.register_secret(secret_name, secret_value)

                    logger.info(
                        f"Secret masking enabled for {registry.get_count()} secret(s)"
                    )
                except Exception as e:
                    logger.warning(
                        f"Failed to set up secret masking: {e}. Continuing without masking."
                    )

            # Validate secret values and warn on potential issues
            if secret_matches:
                for match in secret_matches:
                    secret_value = secret_values_dict.get(match, "")
                    SubProcessTaskUtil._warn_on_bad_secret_value(
                        execution_ctx, match, secret_value
                    )

            json_recipe = json.loads(resolved_recipe, strict=False)
            json_recipe["run_id"] = execution_ctx.exec_id

            # Return recipe, all secret names, and secrets we added to environment
            return json_recipe, secrets_to_resolve, secrets_to_add_to_environ
        except Exception:
            # Clean up only the secrets we added (not ones already in environment)
            for secret_name in secrets_to_add_to_environ:
                os.environ.pop(secret_name, None)
            raise

    @staticmethod
    def _get_plugin_from_recipe(recipe: dict) -> str:
        # The source type -- ASSUMPTION ALERT: This should always correspond to the plugin name.
        return recipe["source"]["type"]

    @staticmethod
    def _write_recipe_to_file(
        dir_path: str, recipe: dict, file_name: str = "recipe.yml"
    ) -> str:
        # 1. Create directories to the path
        os.makedirs(dir_path, mode=0o777, exist_ok=True)

        # 2. Dump recipe dictionary to a YAML string
        yaml_recipe = yaml.dump(recipe)

        # 3. Write YAML to file
        file_path = dir_path + "/" + file_name
        with open(file_path, "w") as file_handle:
            file_handle.write(yaml_recipe)

        return file_path

    @staticmethod
    def _remove_directory(dir_path: str) -> None:
        shutil.rmtree(dir_path)


class SubProcessRecipeTaskArgs(PermissiveConfigModel):
    recipe: str
    version: str = "latest"

    extra_pip_requirements: list[str] = []
    extra_pip_plugins: list[str] = []
    extra_env_vars: dict = {}

    @pydantic.field_validator(
        "extra_pip_requirements", "extra_pip_plugins", mode="before"
    )
    @classmethod
    def parse_json_list_fields(cls, v: Any) -> list:
        if isinstance(v, str):
            # Handle corner case where UI passes an empty string
            return [] if v == "" else json.loads(v)
        return v

    @pydantic.field_validator("extra_env_vars", mode="before")
    @classmethod
    def parse_json_dict_field(cls, v: Any) -> dict:
        if isinstance(v, str):
            # Handle corner case where UI passes an empty string
            return {} if v == "" else json.loads(v)
        return v

    def get_venv_name(self, plugin: str) -> str:
        """Generate venv name using venv utilities."""
        return venv_utils.get_venv_name(
            plugin=plugin,
            version=self.version,
            extra_pip_requirements=self.extra_pip_requirements,
            extra_pip_plugins=self.extra_pip_plugins,
        )

    def should_use_bundled_venv(self) -> bool:
        """Check if this configuration should use a Bundled (pre-packaged) venv."""
        return venv_utils.should_use_bundled_venv(self.version)

    def get_combined_env_vars(self) -> dict:
        # Combines os.environ and user-provided custom env vars.
        # User's extra_env_vars will override system environment variables to allow
        # users to explicitly configure their ingestion environment.
        # Filter out empty string values from extra_env_vars to prevent them from overriding
        # non-empty system environment variables with empty values
        filtered_extra_vars = {k: v for k, v in self.extra_env_vars.items() if v != ""}
        combined = {
            **os.environ,  # System vars as base
            **filtered_extra_vars,  # User vars override (non-empty only)
        }

        return combined
