import importlib.util
import os
import re
import sys
import time
import traceback
from functools import update_wrapper
from typing import Dict, Optional

# Module-level sets to track printed exception messages
# to prevent duplicate messages as exceptions propagate through decorators
_printed_uninstalled_packages_errors = set()
_printed_adapter_missing_errors = set()
_printed_project_not_found_errors = set()
_printed_general_errors = set()


def _extract_clean_adapter_error(exc_str: str) -> str:
    """Extract a clean, user-friendly message from adapter error."""
    # Find adapter type (e.g., "postgres", "databricks")
    adapter_match = re.search(r"Could not find adapter type (\w+)", exc_str)
    adapter_name = adapter_match.group(1) if adapter_match else "unknown"
    return f"❌ Missing adapter '{adapter_name}'. Run 'dvt sync' to install it."


def _is_project_not_found_error(exc: Exception) -> bool:
    """Check if exception is a project-not-found error."""
    exc_str = str(exc)
    return (
        "No dbt_project.yml" in exc_str
        or "No dvt_project.yml" in exc_str
        or "not found at expected path" in exc_str
    )


def _extract_clean_project_error(exc_str: str) -> str:
    """Extract a clean, user-friendly message from project-not-found error."""
    return "❌ Not in a DVT project directory. Run from a folder with dbt_project.yml or use --project-dir."


from click import Context

import dvt.tracking
from dvt.adapters.factory import adapter_management, get_adapter, register_adapter
from dvt.cli.exceptions import ExceptionExit, ResultExit
from dvt.cli.flags import Flags
from dvt.config import RuntimeConfig
from dvt.config.catalogs import get_active_write_integration, load_catalogs
from dvt.config.runtime import UnsetProfile, load_profile, load_project
from dvt.context.providers import generate_runtime_macro_context
from dvt.context.query_header import generate_query_header_context
from dvt.deprecations import show_deprecations_summary
from dvt.env_vars import KNOWN_ENGINE_ENV_VARS, validate_engine_env_vars
from dvt.events.logging import setup_event_logger
from dvt.events.types import (
    ArtifactUploadError,
    CommandCompleted,
    MainEncounteredError,
    MainReportArgs,
    MainReportVersion,
    MainStackTrace,
    MainTrackingUserState,
    ResourceReport,
)
from dvt.exceptions import DvtProfileError, DvtProjectError, FailFastError
from dvt.flags import get_flag_dict, get_flags, set_flags
from dvt.mp_context import get_mp_context
from dvt.parser.manifest import parse_manifest
from dvt.plugins import set_up_plugin_manager
from dvt.profiler import profiler
from dvt.tracking import active_user, initialize_from_flags, track_run
from dvt.utils import try_get_max_rss_kb
from dvt.utils.artifact_upload import upload_artifacts
from dvt.version import installed as installed_version
from dbt_common.clients.system import get_env
from dbt_common.context import get_invocation_context, set_invocation_context
from dbt_common.events.base_types import EventLevel
from dbt_common.events.event_manager_client import get_event_manager
from dbt_common.events.functions import LOG_VERSION, fire_event
from dbt_common.events.helpers import get_json_string_utcnow
from dbt_common.exceptions import DbtBaseException as DbtException
from dbt_common.invocation import reset_invocation_id
from dbt_common.record import (
    Recorder,
    RecorderMode,
    get_record_mode_from_env,
    get_record_types_from_dict,
    get_record_types_from_env,
)
from dbt_common.utils import cast_dict_to_dict_of_strings


def _cross_propagate_engine_env_vars(env_dict: Dict[str, str]) -> None:
    for env_var in KNOWN_ENGINE_ENV_VARS:
        if env_var.old_name is not None:
            # If the old name is in the env dict, and not the new name, set the new name based on the old name
            if env_var.old_name in env_dict and env_var.name not in env_dict:
                env_dict[env_var.name] = env_dict[env_var.old_name]
            # If the new name is in the env dict, override the old name with it
            elif env_var.name in env_dict:
                env_dict[env_var.old_name] = env_dict[env_var.name]


def preflight(func):
    def wrapper(*args, **kwargs):
        ctx = args[0]
        assert isinstance(ctx, Context)
        ctx.obj = ctx.obj or {}

        set_invocation_context({})

        # Record/Replay
        setup_record_replay()

        # Must be set after record/replay is set up so that the env can be
        # recorded or replayed if needed.
        env_dict = get_env()
        _cross_propagate_engine_env_vars(env_dict)
        get_invocation_context()._env = env_dict

        # Flags
        try:
            flags = Flags(ctx)
        except Exception as e:
            # Print exception immediately to stderr so it's never silent
            sys.stderr.write(f"Error initializing Flags: {e}\n")
            sys.stderr.write(traceback.format_exc())
            sys.stderr.flush()
            raise
        ctx.obj["flags"] = flags
        set_flags(flags)
        get_event_manager().require_warn_or_error_handling = (
            flags.require_all_warnings_handled_by_warn_error
        )

        # Reset invocation_id for each 'invocation' of a dbt command (can happen multiple times in a single process)
        reset_invocation_id()

        # Logging
        callbacks = ctx.obj.get("callbacks", [])
        setup_event_logger(flags=flags, callbacks=callbacks)

        # Tracking
        initialize_from_flags(flags.SEND_ANONYMOUS_USAGE_STATS, flags.PROFILES_DIR)
        ctx.with_resource(track_run(run_command=flags.WHICH))

        # Now that we have our logger, fire away!
        fire_event(
            MainReportVersion(version=str(installed_version), log_version=LOG_VERSION)
        )
        flags_dict_str = cast_dict_to_dict_of_strings(get_flag_dict())
        fire_event(MainReportArgs(args=flags_dict_str))

        # Deprecation warnings
        flags.fire_deprecations(ctx)

        if active_user is not None:  # mypy appeasement, always true
            fire_event(MainTrackingUserState(user_state=active_user.state()))

        # Profiling
        if flags.RECORD_TIMING_INFO:
            ctx.with_resource(profiler(enable=True, outfile=flags.RECORD_TIMING_INFO))

        # Adapter management
        ctx.with_resource(adapter_management())

        # Validate engine env var restricted name space
        validate_engine_env_vars()

        return func(*args, **kwargs)

    return update_wrapper(wrapper, func)


def setup_record_replay():
    rec_mode = get_record_mode_from_env()
    rec_types = get_record_types_from_env()

    recorder: Optional[Recorder] = None
    if rec_mode == RecorderMode.REPLAY:
        previous_recording_path = os.environ.get(
            "DBT_ENGINE_RECORDER_FILE_PATH"
        ) or os.environ.get("DBT_RECORDER_FILE_PATH")
        recorder = Recorder(
            RecorderMode.REPLAY,
            types=rec_types,
            previous_recording_path=previous_recording_path,
        )
    elif rec_mode == RecorderMode.DIFF:
        previous_recording_path = os.environ.get(
            "DBT_ENGINE_RECORDER_FILE_PATH"
        ) or os.environ.get("DBT_RECORDER_FILE_PATH")
        # ensure types match the previous recording
        types = get_record_types_from_dict(previous_recording_path)
        recorder = Recorder(
            RecorderMode.DIFF,
            types=types,
            previous_recording_path=previous_recording_path,
        )
    elif rec_mode == RecorderMode.RECORD:
        recorder = Recorder(RecorderMode.RECORD, types=rec_types)

    get_invocation_context().recorder = recorder


def tear_down_record_replay():
    recorder = get_invocation_context().recorder
    if recorder is not None:
        if recorder.mode == RecorderMode.RECORD:
            recorder.write()
        if recorder.mode == RecorderMode.DIFF:
            recorder.write()
            recorder.write_diffs(diff_file_name="recording_diffs.json")
        elif recorder.mode == RecorderMode.REPLAY:
            recorder.write_diffs("replay_diffs.json")


def postflight(func):
    """The decorator that handles all exception handling for the click commands.
    This decorator must be used before any other decorators that may throw an exception."""

    def wrapper(*args, **kwargs):
        ctx = args[0]
        start_func = time.perf_counter()
        success = False

        try:
            result, success = func(*args, **kwargs)
        except FailFastError as e:
            fire_event(MainEncounteredError(exc=str(e)))
            raise ResultExit(e.result)
        except DbtException as e:
            # For certain errors, don't fire event or print traceback - print once and skip event
            # to avoid duplicate messages as exception propagates through decorators
            from dvt.exceptions import (
                UninstalledPackagesFoundError,
                PySparkNotInstalledError,
            )
            from dvt.cli.exceptions import ExceptionExit as ExceptionExitType

            # Check if this is a special error type (either directly or wrapped in ExceptionExit)
            actual_exception = e.exception if isinstance(e, ExceptionExitType) else e
            if isinstance(actual_exception, PySparkNotInstalledError):
                # PySpark not installed: print clean message once, NO traceback
                exc_message = str(actual_exception)
                if exc_message not in _printed_uninstalled_packages_errors:
                    sys.stderr.write(exc_message)
                    sys.stderr.flush()
                    _printed_uninstalled_packages_errors.add(exc_message)
            elif isinstance(actual_exception, UninstalledPackagesFoundError):
                # Use module-level set to track printed exception messages to prevent duplicates
                # across multiple decorator catches (use message as key since exception objects may differ)
                exc_message = str(actual_exception)
                if exc_message not in _printed_uninstalled_packages_errors:
                    sys.stderr.write(f"{actual_exception}\n")
                    sys.stderr.flush()
                    _printed_uninstalled_packages_errors.add(exc_message)
            elif isinstance(actual_exception, DvtProfileError) and (
                "Could not find adapter" in str(actual_exception)
                or "Run 'dvt sync'" in str(actual_exception)
            ):
                # Adapter missing: print clean message once, NO traceback
                clean_msg = _extract_clean_adapter_error(str(actual_exception))
                if clean_msg not in _printed_adapter_missing_errors:
                    sys.stderr.write(clean_msg + "\n")
                    sys.stderr.flush()
                    _printed_adapter_missing_errors.add(clean_msg)
            elif isinstance(
                actual_exception, DvtProjectError
            ) and _is_project_not_found_error(actual_exception):
                # Project not found: print clean message once, NO traceback
                clean_msg = _extract_clean_project_error(str(actual_exception))
                if clean_msg not in _printed_project_not_found_errors:
                    sys.stderr.write(clean_msg + "\n")
                    sys.stderr.flush()
                    _printed_project_not_found_errors.add(clean_msg)
            else:
                exc_message = str(e)
                if exc_message not in _printed_general_errors:
                    sys.stderr.write(f"\n{exc_message}\n")
                    sys.stderr.flush()
                    _printed_general_errors.add(exc_message)
            # Only wrap if not already wrapped
            if isinstance(e, ExceptionExitType):
                raise e
            raise ExceptionExit(e)
        except BaseException as e:
            # Check if this is UninstalledPackagesFoundError, PySparkNotInstalledError, or adapter-missing DvtProfileError wrapped in ExceptionExit
            from dvt.exceptions import (
                UninstalledPackagesFoundError,
                PySparkNotInstalledError,
            )
            from dvt.cli.exceptions import ExceptionExit as ExceptionExitType

            actual_exception = e.exception if isinstance(e, ExceptionExitType) else e
            if isinstance(actual_exception, PySparkNotInstalledError):
                # PySpark not installed: print clean message once, NO traceback
                exc_message = str(actual_exception)
                if exc_message not in _printed_uninstalled_packages_errors:
                    sys.stderr.write(exc_message)
                    sys.stderr.flush()
                    _printed_uninstalled_packages_errors.add(exc_message)
            elif isinstance(actual_exception, UninstalledPackagesFoundError):
                # Use module-level set to track printed exception messages to prevent duplicates
                # across multiple decorator catches (use message as key since exception objects may differ)
                exc_message = str(actual_exception)
                if exc_message not in _printed_uninstalled_packages_errors:
                    sys.stderr.write(f"{actual_exception}\n")
                    sys.stderr.flush()
                    _printed_uninstalled_packages_errors.add(exc_message)
            elif isinstance(actual_exception, DvtProfileError) and (
                "Could not find adapter" in str(actual_exception)
                or "Run 'dvt sync'" in str(actual_exception)
            ):
                # Adapter missing: print clean message once, NO traceback
                clean_msg = _extract_clean_adapter_error(str(actual_exception))
                if clean_msg not in _printed_adapter_missing_errors:
                    sys.stderr.write(clean_msg + "\n")
                    sys.stderr.flush()
                    _printed_adapter_missing_errors.add(clean_msg)
            elif isinstance(
                actual_exception, DvtProjectError
            ) and _is_project_not_found_error(actual_exception):
                # Project not found: print clean message once, NO traceback
                clean_msg = _extract_clean_project_error(str(actual_exception))
                if clean_msg not in _printed_project_not_found_errors:
                    sys.stderr.write(clean_msg + "\n")
                    sys.stderr.flush()
                    _printed_project_not_found_errors.add(clean_msg)
            else:
                exc_message = (
                    str(actual_exception) if actual_exception is not e else str(e)
                )
                if exc_message not in _printed_general_errors:
                    sys.stderr.write(f"\n{exc_message}\n")
                    sys.stderr.flush()
                    _printed_general_errors.add(exc_message)
            # Only wrap if not already wrapped
            if isinstance(e, ExceptionExitType):
                raise e
            raise ExceptionExit(e)
        finally:
            # Fire ResourceReport, but only on systems which support the resource
            # module. (Skip it on Windows).
            try:
                if get_flags().upload_to_artifacts_ingest_api:
                    upload_artifacts(
                        get_flags().project_dir,
                        get_flags().target_path,
                        ctx.command.name,
                    )

            except Exception as e:
                fire_event(ArtifactUploadError(msg=str(e)))

            show_deprecations_summary()

            if importlib.util.find_spec("resource") is not None:
                import resource

                rusage = resource.getrusage(resource.RUSAGE_SELF)
                fire_event(
                    ResourceReport(
                        command_name=ctx.command.name,
                        command_success=success,
                        command_wall_clock_time=time.perf_counter() - start_func,
                        process_user_time=rusage.ru_utime,
                        process_kernel_time=rusage.ru_stime,
                        process_mem_max_rss=try_get_max_rss_kb() or rusage.ru_maxrss,
                        process_in_blocks=rusage.ru_inblock,
                        process_out_blocks=rusage.ru_oublock,
                    ),
                    (
                        EventLevel.INFO
                        if "flags" in ctx.obj and ctx.obj["flags"].SHOW_RESOURCE_REPORT
                        else None
                    ),
                )

            fire_event(
                CommandCompleted(
                    command=ctx.command_path,
                    success=success,
                    completed_at=get_json_string_utcnow(),
                    elapsed=time.perf_counter() - start_func,
                )
            )

            tear_down_record_replay()

        if not success:
            raise ResultExit(result)

        return (result, success)

    return update_wrapper(wrapper, func)


# TODO: UnsetProfile is necessary for deps and clean to load a project.
# This decorator and its usage can be removed once https://github.com/dbt-labs/dbt-core/issues/6257 is closed.
def unset_profile(func):
    def wrapper(*args, **kwargs):
        ctx = args[0]
        assert isinstance(ctx, Context)

        profile = UnsetProfile()
        ctx.obj["profile"] = profile

        return func(*args, **kwargs)

    return update_wrapper(wrapper, func)


def profile(func):
    def wrapper(*args, **kwargs):
        ctx = args[0]
        assert isinstance(ctx, Context)

        flags = ctx.obj["flags"]
        # TODO: Generalize safe access to flags.THREADS:
        # https://github.com/dbt-labs/dbt-core/issues/6259
        threads = getattr(flags, "THREADS", None)
        profile = load_profile(
            flags.PROJECT_DIR, flags.VARS, flags.PROFILE, flags.TARGET, threads
        )
        ctx.obj["profile"] = profile
        get_invocation_context().uses_adapter(profile.credentials.type)

        return func(*args, **kwargs)

    return update_wrapper(wrapper, func)


def project(func):
    def wrapper(*args, **kwargs):
        ctx = args[0]
        assert isinstance(ctx, Context)

        # TODO: Decouple target from profile, and remove the need for profile here:
        # https://github.com/dbt-labs/dbt-core/issues/6257
        if not ctx.obj.get("profile"):
            raise DvtProjectError("profile required for project")

        flags = ctx.obj["flags"]
        # TODO deprecations warnings fired from loading the project will lack
        # the project_id in the snowplow event.

        # Determine if vars should be required during project loading.
        # Commands that don't need vars evaluated (like 'deps', 'clean')
        # should use lenient mode (require_vars=False) to allow missing vars.
        # Commands that validate or execute (like 'run', 'compile', 'build', 'debug') should use
        # strict mode (require_vars=True) to show helpful "Required var X not found" errors.
        # If adding more commands to lenient mode, update this condition.
        require_vars = flags.WHICH != "deps"

        project = load_project(
            flags.PROJECT_DIR,
            flags.VERSION_CHECK,
            ctx.obj["profile"],
            flags.VARS,
            validate=True,
            require_vars=require_vars,
        )
        ctx.obj["project"] = project

        # Plugins
        set_up_plugin_manager(project_name=project.project_name)

        if dvt.tracking.active_user is not None:
            project_id = None if project is None else project.hashed_name()

            dvt.tracking.track_project_id({"project_id": project_id})

        return func(*args, **kwargs)

    return update_wrapper(wrapper, func)


def runtime_config(func):
    """A decorator used by click command functions for generating a runtime
    config given a profile and project.
    """

    def wrapper(*args, **kwargs):
        ctx = args[0]
        assert isinstance(ctx, Context)

        req_strs = ["profile", "project"]
        reqs = [ctx.obj.get(req_str) for req_str in req_strs]

        if None in reqs:
            raise DvtProjectError("profile and project required for runtime_config")

        config = RuntimeConfig.from_parts(
            ctx.obj["project"],
            ctx.obj["profile"],
            ctx.obj["flags"],
        )

        ctx.obj["runtime_config"] = config

        if dvt.tracking.active_user is not None:
            adapter_type = (
                getattr(config.credentials, "type", None)
                if hasattr(config, "credentials")
                else None
            )
            adapter_unique_id = (
                config.credentials.hashed_unique_field()
                if hasattr(config, "credentials")
                else None
            )

            dvt.tracking.track_adapter_info(
                {
                    "adapter_type": adapter_type,
                    "adapter_unique_id": adapter_unique_id,
                }
            )

        return func(*args, **kwargs)

    return update_wrapper(wrapper, func)


def catalogs(func):
    """A decorator used by click command functions for loading catalogs"""

    def wrapper(*args, **kwargs):
        ctx = args[0]
        assert isinstance(ctx, Context)

        req_strs = ["flags", "profile", "project"]
        reqs = [ctx.obj.get(req_str) for req_str in req_strs]
        if None in reqs:
            raise DvtProjectError("profile and flags required to load catalogs")

        flags = ctx.obj["flags"]
        ctx_project = ctx.obj["project"]

        _catalogs = load_catalogs(
            flags.PROJECT_DIR, ctx_project.project_name, flags.VARS
        )
        ctx.obj["catalogs"] = _catalogs

        return func(*args, **kwargs)

    return update_wrapper(wrapper, func)


def manifest(*args0, write=True, write_perf_info=False):
    """A decorator used by click command functions for generating a manifest
    given a profile, project, and runtime config. This also registers the adapter
    from the runtime config and conditionally writes the manifest to disk.
    """

    def outer_wrapper(func):
        def wrapper(*args, **kwargs):
            ctx = args[0]
            assert isinstance(ctx, Context)

            setup_manifest(ctx, write=write, write_perf_info=write_perf_info)
            return func(*args, **kwargs)

        return update_wrapper(wrapper, func)

    # if there are no args, the decorator was used without params @decorator
    # otherwise, the decorator was called with params @decorator(arg)
    if len(args0) == 0:
        return outer_wrapper
    return outer_wrapper(args0[0])


def setup_manifest(ctx: Context, write: bool = True, write_perf_info: bool = False):
    """Load the manifest and add it to the context."""
    req_strs = ["profile", "project", "runtime_config"]
    reqs = [ctx.obj.get(dep) for dep in req_strs]

    if None in reqs:
        raise DvtProjectError(
            "profile, project, and runtime_config required for manifest"
        )

    runtime_config = ctx.obj["runtime_config"]

    catalogs = ctx.obj["catalogs"] if "catalogs" in ctx.obj else []
    active_integrations = [
        get_active_write_integration(catalog) for catalog in catalogs
    ]

    # if a manifest has already been set on the context, don't overwrite it
    if ctx.obj.get("manifest") is None:
        ctx.obj["manifest"] = parse_manifest(
            runtime_config,
            write_perf_info,
            write,
            ctx.obj["flags"].write_json,
            active_integrations,
        )
        adapter = get_adapter(runtime_config)
    else:
        register_adapter(runtime_config, get_mp_context())
        adapter = get_adapter(runtime_config)
        adapter.set_macro_context_generator(generate_runtime_macro_context)  # type: ignore[arg-type]
        adapter.set_macro_resolver(ctx.obj["manifest"])
        query_header_context = generate_query_header_context(
            adapter.config, ctx.obj["manifest"]
        )  # type: ignore[attr-defined]
        adapter.connections.set_query_header(query_header_context)
        for integration in active_integrations:
            adapter.add_catalog_integration(integration)
