#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
from __future__ import annotations

import contextlib
import logging
import textwrap
import time
from datetime import datetime as dt, timedelta, timezone
from pathlib import Path
from unittest import mock
from unittest.mock import ANY, call

import boto3
import pendulum
import pytest
import time_machine
from botocore.exceptions import ClientError
from moto import mock_aws
from pydantic import TypeAdapter
from watchtower import CloudWatchLogHandler

from airflow.models import DAG, DagRun, TaskInstance
from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook
from airflow.providers.amazon.aws.log.cloudwatch_task_handler import (
    CloudWatchRemoteLogIO,
    CloudwatchTaskHandler,
)
from airflow.providers.amazon.aws.utils import datetime_to_epoch_utc_ms
from airflow.utils.state import State
from airflow.utils.timezone import datetime

from tests_common.test_utils.compat import EmptyOperator
from tests_common.test_utils.config import conf_vars
from tests_common.test_utils.dag import sync_dag_to_db
from tests_common.test_utils.db import clear_db_dag_bundles, clear_db_dags, clear_db_runs
from tests_common.test_utils.taskinstance import create_task_instance
from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_2_PLUS


def get_time_str(time_in_milliseconds):
    dt_time = dt.fromtimestamp(time_in_milliseconds / 1000.0, tz=timezone.utc)
    return dt_time.strftime("%Y-%m-%dT%H:%M:%SZ")


@pytest.fixture(autouse=True)
def logmock():
    with mock_aws():
        yield


@pytest.fixture(autouse=True)
def _cleanup_cloudwatch_handlers():
    # Watchtower's CloudWatchLogHandler spawns a queue worker thread that keeps
    # the handler alive even after the test fixture tears down, so it stays
    # registered in logging._handlerList. Several tests here never call close()
    # or mock watchtower.CloudWatchLogHandler.close, leaking the handler across
    # tests. When a later test in the same xdist worker calls
    # logging.config.dictConfig (e.g. settings.configure_logging()), its
    # _clearExistingHandlers path runs close() on the leaked handler and blocks
    # up to FLUSH_TIMEOUT waiting for an undrainable queue — mock_aws is gone.
    # Closing here while mock_aws is still active drains cleanly.
    yield
    for handler_ref in logging._handlerList[:]:
        handler = handler_ref()
        if handler is not None and isinstance(handler, CloudWatchLogHandler):
            with contextlib.suppress(Exception):
                handler.close()
            logging._removeHandlerRef(handler_ref)


# We only test this directly on Airflow 3
@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="This path only works on Airflow 3")
class TestCloudRemoteLogIO:
    # We use the cap_structlog so that our changes get reverted for us
    @pytest.fixture(autouse=True)
    def setup_tests(self, create_runtime_ti, tmp_path, monkeypatch):
        import structlog

        import airflow.logging_config
        import airflow.sdk.log
        from airflow.sdk import BaseOperator

        task = BaseOperator(task_id="task_1")
        self.ti = create_runtime_ti(task)

        self.remote_log_group = "log_group_name"
        self.region_name = "us-west-2"
        self.task_log_path = "dag_id=a/0:0.log"
        self.local_log_location = tmp_path / "local-cloudwatch-log-location"
        self.local_log_location.mkdir()
        # Create the local log file structure
        task_log_path_parts = self.task_log_path.split("/")
        dag_dir = self.local_log_location / task_log_path_parts[0]
        dag_dir.mkdir()
        task_log_file = dag_dir / task_log_path_parts[1]
        task_log_file.touch()

        # The subject under test
        self.subject = CloudWatchRemoteLogIO(
            base_log_folder=self.local_log_location,
            log_group_arn=f"arn:aws:logs:{self.region_name}:11111111:log-group:{self.remote_log_group}",
        )

        conn = boto3.client("logs", region_name=self.region_name)
        conn.create_log_group(logGroupName=self.remote_log_group)

        processors = structlog.get_config()["processors"]
        logger_factory = structlog.get_config()["logger_factory"]
        old_processors = processors.copy()

        try:
            # Modify `_Configuration.default_processors` set via `configure` but always
            # keep the list instance intact to not break references held by bound
            # loggers.

            # Set up the right chain of processors so the event looks like we want for our full test
            if AIRFLOW_V_3_2_PLUS:
                airflow.sdk.log._ActiveLoggingConfig.set(self.subject, None)
            else:
                monkeypatch.setattr(airflow.logging_config, "REMOTE_TASK_LOG", self.subject)
            # Clear @cache'd processors so this test picks up self.subject.
            airflow.sdk.log.logging_processors.cache_clear()
            try:
                procs = airflow.sdk.log.logging_processors(colors=False, json_output=False)
            except TypeError:
                # Compat issue only comes up in the tests, not in the real code
                procs, _ = airflow.sdk.log.logging_processors(enable_pretty_log=False)
            processors.clear()
            processors.extend(procs)

            # Replace the last "output" one with a DropEvent one instead - else we get the output on stdout
            def drop(*args):
                raise structlog.DropEvent()

            processors[-1] = drop
            structlog.configure(
                processors=processors,
                # Create a logger factory and pass in the file path we want it to use
                # This is because we use the logger to determine the streamname/filepath
                # in the CloudWatchRemoteLogIO processor.
                logger_factory=structlog.PrintLoggerFactory(file=task_log_file.open("w+")),
            )
            yield
        finally:
            # remove LogCapture and restore original processors
            processors.clear()
            processors.extend(old_processors)
            structlog.configure(processors=old_processors, logger_factory=logger_factory)
            # Clear @cache to avoid cross-test contamination.
            airflow.sdk.log.logging_processors.cache_clear()

    def test_upload_deletes_local_log_dir_when_delete_local_copy_true(self):
        """upload() with delete_local_copy=True deletes the local log parent directory."""
        assert self.subject.delete_local_copy is True  # attrs default
        dag_dir = self.local_log_location / "dag_id=a"
        assert dag_dir.exists()

        with mock.patch.object(self.subject, "close"):
            self.subject.upload(self.task_log_path, self.ti)

        assert not dag_dir.exists()
        # Close watchtower handler; conf_vars needed because base_log_folder differs from tmp_path.
        with conf_vars({("logging", "base_log_folder"): self.local_log_location.as_posix()}):
            self.subject.handler.close()

    def test_upload_does_not_delete_when_delete_local_copy_false(self):
        """upload() with delete_local_copy=False leaves local files untouched."""
        self.subject.delete_local_copy = False
        dag_dir = self.local_log_location / "dag_id=a"
        assert dag_dir.exists()

        with mock.patch.object(self.subject, "close"):
            self.subject.upload(self.task_log_path, self.ti)

        assert dag_dir.exists()
        with conf_vars({("logging", "base_log_folder"): self.local_log_location.as_posix()}):
            self.subject.handler.close()

    def test_upload_handles_absolute_path(self):
        """upload() resolves an absolute path to the correct parent directory."""
        dag_dir = self.local_log_location / "dag_id=a"
        assert dag_dir.exists()
        absolute_path = self.local_log_location / self.task_log_path

        with mock.patch.object(self.subject, "close"):
            self.subject.upload(absolute_path, self.ti)

        assert not dag_dir.exists()
        with conf_vars({("logging", "base_log_folder"): self.local_log_location.as_posix()}):
            self.subject.handler.close()

    def test_upload_skips_deletion_outside_base_log_folder(self):
        """upload() logs a warning and skips deletion for paths outside base_log_folder."""
        outside_path = "/tmp/evil/../../../etc/passwd"
        with conf_vars({("logging", "base_log_folder"): self.local_log_location.as_posix()}):
            with mock.patch.object(self.subject, "close"):
                self.subject.upload(outside_path, self.ti)
        dag_dir = self.local_log_location / "dag_id=a"
        assert dag_dir.exists()

    @time_machine.travel(datetime(2025, 3, 27, 21, 58, 1, 2345), tick=False)
    def test_log_message(self):
        # Use a context instead of a decorator on the test method because we need access to self to
        # get the path from the setup method.
        with conf_vars({("logging", "base_log_folder"): self.local_log_location.as_posix()}):
            import structlog

            log = structlog.get_logger()
            log.info("Hi", foo="bar")
            # We need to close in order to flush the logs etc.
            self.subject.close()

            # close call should only flush the logs and not set shutting_down to True
            assert self.subject.handler.shutting_down is False

            # Inside the Cloudwatch logger we swap colons for underscores since colons are not allowed in
            # stream names.
            stream_name = self.task_log_path.replace(":", "_")
            logs = self.subject.read(stream_name, self.ti)

            metadata, logs = logs

            assert metadata == [
                f"Reading remote log from Cloudwatch log_group: log_group_name log_stream: {stream_name}"
            ]
            assert logs == [
                '{"foo": "bar", "event": "Hi", "level": "info", "timestamp": "2025-03-27T21:58:01.002000+00:00"}\n'
            ]


@pytest.mark.db_test
class TestCloudwatchTaskHandler:
    def clear_db(self):
        if AIRFLOW_V_3_0_PLUS:
            clear_db_runs()
            clear_db_dags()
            clear_db_dag_bundles()

    @pytest.fixture(autouse=True)
    def setup(self, create_log_template, tmp_path_factory, session, testing_dag_bundle):
        # self.clear_db()
        with conf_vars({("logging", "remote_log_conn_id"): "aws_default"}):
            self.remote_log_group = "log_group_name"
            self.region_name = "us-west-2"
            self.local_log_location = str(tmp_path_factory.mktemp("local-cloudwatch-log-location"))
            if AIRFLOW_V_3_0_PLUS:
                create_log_template("{dag_id}/{task_id}/{logical_date}/{try_number}.log")
            else:
                create_log_template("{dag_id}/{task_id}/{execution_date}/{try_number}.log")
            self.cloudwatch_task_handler = CloudwatchTaskHandler(
                self.local_log_location,
                f"arn:aws:logs:{self.region_name}:11111111:log-group:{self.remote_log_group}",
            )

        date = datetime(2020, 1, 1)
        dag_id = "dag_for_testing_cloudwatch_task_handler"
        task_id = "task_for_testing_cloudwatch_log_handler"
        self.dag = DAG(dag_id=dag_id, schedule=None, start_date=date)
        task = EmptyOperator(task_id=task_id, dag=self.dag)
        if AIRFLOW_V_3_0_PLUS:
            sync_dag_to_db(self.dag)
            dag_run = DagRun(
                dag_id=self.dag.dag_id,
                logical_date=date,
                run_id="test",
                run_type="scheduled",
            )
        else:
            dag_run = DagRun(
                dag_id=self.dag.dag_id,
                execution_date=date,
                run_id="test",
                run_type="scheduled",
            )
        session.add(dag_run)
        session.commit()
        session.refresh(dag_run)

        if AIRFLOW_V_3_0_PLUS:
            from airflow.models.dag_version import DagVersion

            dag_version = DagVersion.get_latest_version(self.dag.dag_id, session=session)
            self.ti = create_task_instance(task=task, run_id=dag_run.run_id, dag_version_id=dag_version.id)
        else:
            self.ti = TaskInstance(task=task, run_id=dag_run.run_id)
        self.ti.dag_run = dag_run
        self.ti.try_number = 1
        self.ti.state = State.RUNNING
        session.add(self.ti)
        session.commit()

        self.remote_log_stream = (f"{dag_id}/{task_id}/{date.isoformat()}/{self.ti.try_number}.log").replace(
            ":", "_"
        )
        self.conn = boto3.client("logs", region_name=self.region_name)

        yield

        self.cloudwatch_task_handler.handler = None
        del self.cloudwatch_task_handler

        self.clear_db()

    def test_hook(self):
        assert isinstance(self.cloudwatch_task_handler.hook, AwsLogsHook)

    def test_handler(self):
        self.cloudwatch_task_handler.set_context(self.ti)
        assert isinstance(self.cloudwatch_task_handler.handler, CloudWatchLogHandler)

    def test_write(self):
        handler = self.cloudwatch_task_handler
        handler.set_context(self.ti)
        messages = [str(i) for i in range(10)]

        with mock.patch("watchtower.CloudWatchLogHandler.emit") as mock_emit:
            for message in messages:
                handler.handle(message)
            mock_emit.assert_has_calls([call(message) for message in messages])

    # TODO: Remove when we stop testing for 2.11 compatibility
    @conf_vars({("core", "use_historical_filename_templates"): "True"})
    @time_machine.travel(datetime(2025, 3, 27, 21, 58, 1, 2345), tick=False)
    def test_read(self, monkeypatch):
        # Confirmed via AWS Support call:
        # CloudWatch events must be ordered chronologically otherwise
        # boto3 put_log_event API throws InvalidParameterException
        # (moto does not throw this exception)
        current_time = int(time.time()) * 1000
        generate_log_events(
            self.conn,
            self.remote_log_group,
            self.remote_log_stream,
            [
                {"timestamp": current_time - 2000, "message": "First"},
                {"timestamp": current_time - 1000, "message": "Second"},
                {"timestamp": current_time, "message": "Third"},
            ],
        )
        monkeypatch.setattr(
            self.cloudwatch_task_handler,
            "_read_from_logs_server",
            lambda ti, worker_log_rel_path: ([], []),
        )
        msg_template = textwrap.dedent("""
             INFO - ::group::Log message source details
            *** Reading remote log from Cloudwatch log_group: {} log_stream: {}
             INFO - ::endgroup::
            {}
        """)[1:][:-1]  # Strip off leading and trailing new lines, but not spaces

        logs, metadata = self.cloudwatch_task_handler.read(self.ti)
        if AIRFLOW_V_3_0_PLUS:
            from airflow.utils.log.file_task_handler import StructuredLogMessage

            logs = list(logs)
            results = TypeAdapter(list[StructuredLogMessage]).dump_python(logs)
            assert results[-4:] == [
                {"event": "::endgroup::", "timestamp": None},
                {
                    "event": "[2025-03-27T21:57:59Z] First",
                    "timestamp": pendulum.datetime(2025, 3, 27, 21, 57, 59),
                },
                {
                    "event": "[2025-03-27T21:58:00Z] Second",
                    "timestamp": pendulum.datetime(2025, 3, 27, 21, 58, 0),
                },
                {
                    "event": "[2025-03-27T21:58:01Z] Third",
                    "timestamp": pendulum.datetime(2025, 3, 27, 21, 58, 1),
                },
            ]
            assert metadata == {"end_of_log": False, "log_pos": 3}
        else:
            events = "\n".join(
                [
                    f"[{get_time_str(current_time - 2000)}] First",
                    f"[{get_time_str(current_time - 1000)}] Second",
                    f"[{get_time_str(current_time)}] Third",
                ]
            )
            assert logs == [
                [
                    (
                        "",
                        msg_template.format(self.remote_log_group, self.remote_log_stream, events),
                    )
                ]
            ]

    @pytest.mark.parametrize(
        ("end_date", "expected_end_time"),
        [
            (None, None),
            (
                datetime(2020, 1, 2),
                datetime_to_epoch_utc_ms(datetime(2020, 1, 2) + timedelta(seconds=30)),
            ),
        ],
    )
    @mock.patch.object(AwsLogsHook, "get_log_events")
    def test_get_cloudwatch_logs(self, mock_get_log_events, end_date, expected_end_time):
        self.ti.end_date = end_date
        self.cloudwatch_task_handler.io.get_cloudwatch_logs(self.remote_log_stream, self.ti)
        mock_get_log_events.assert_called_once_with(
            log_group=self.remote_log_group,
            log_stream_name=self.remote_log_stream,
            end_time=expected_end_time,
        )

    @mock.patch.object(AwsLogsHook, "get_log_events")
    def test_get_cloudwatch_logs_missing_stream_yields_hint(self, mock_get_log_events):
        # A missing log stream (no logs written for this try -- e.g. the task logged
        # to stdout instead of remote storage) must not raise (so the log reader does
        # not surface a 500) and must yield a hint instead of nothing, so the reader
        # does not show a blank view that looks like remote logging silently failed.
        def _raise_not_found(*args, **kwargs):
            raise ClientError({"Error": {"Code": "ResourceNotFoundException"}}, "GetLogEvents")
            yield  # pragma: no cover -- makes this a generator function

        mock_get_log_events.side_effect = _raise_not_found
        events = list(self.cloudwatch_task_handler.io.get_cloudwatch_logs(self.remote_log_stream, self.ti))
        assert len(events) == 1
        assert "No log stream found in CloudWatch" in events[0]["message"]
        assert self.remote_log_stream in events[0]["message"]

    @mock.patch.object(AwsLogsHook, "get_log_events")
    def test_get_cloudwatch_logs_other_client_error_propagates(self, mock_get_log_events):
        # Errors other than a missing stream must still surface.
        def _raise_access_denied(*args, **kwargs):
            raise ClientError({"Error": {"Code": "AccessDeniedException"}}, "GetLogEvents")
            yield  # pragma: no cover -- makes this a generator function

        mock_get_log_events.side_effect = _raise_access_denied
        with pytest.raises(ClientError):
            list(self.cloudwatch_task_handler.io.get_cloudwatch_logs(self.remote_log_stream, self.ti))

    @pytest.mark.parametrize(
        ("conf_json_serialize", "expected_serialized_output"),
        [
            pytest.param(
                "airflow.providers.amazon.aws.log.cloudwatch_task_handler.json_serialize_legacy",
                '{"datetime": "2023-01-01T00:00:00+00:00", "customObject": null}',
                id="json-serialize-legacy",
            ),
            pytest.param(
                "airflow.providers.amazon.aws.log.cloudwatch_task_handler.json_serialize",
                '{"datetime": "2023-01-01T00:00:00+00:00", "customObject": "SomeCustomSerialization(...)"}',
                id="json-serialize",
            ),
            pytest.param(
                None,
                '{"datetime": "2023-01-01T00:00:00+00:00", "customObject": null}',
                id="not-set",
            ),
        ],
    )
    @mock.patch.object(AwsLogsHook, "get_log_events")
    def test_write_json_logs(self, mock_get_log_events, conf_json_serialize, expected_serialized_output):
        class ToSerialize:
            def __init__(self):
                pass

            def __repr__(self):
                return "SomeCustomSerialization(...)"

        with conf_vars({("aws", "cloudwatch_task_handler_json_serializer"): conf_json_serialize}):
            handler = self.cloudwatch_task_handler
            handler.set_context(self.ti)
            message = logging.LogRecord(
                name="test_log_record",
                level=logging.DEBUG,
                pathname="fake.path",
                lineno=42,
                args=None,
                exc_info=None,
                msg={
                    "datetime": datetime(2023, 1, 1),
                    "customObject": ToSerialize(),
                },
            )
            with (
                mock.patch("watchtower.threading.Thread"),
                mock.patch("watchtower.queue.Queue") as mq,
            ):
                handler.handle(message)
                mq.return_value.put.assert_called_once_with(
                    {"message": expected_serialized_output, "timestamp": ANY}
                )

    def test_close_prevents_duplicate_calls(self):
        with mock.patch("watchtower.CloudWatchLogHandler.close") as mock_log_handler_close:
            with mock.patch("airflow.utils.log.file_task_handler.FileTaskHandler.set_context"):
                with mock.patch.object(self.cloudwatch_task_handler.io, "upload"):
                    self.cloudwatch_task_handler.set_context(self.ti)
                    for _ in range(5):
                        self.cloudwatch_task_handler.close()

                    mock_log_handler_close.assert_called_once()

    def test_set_context_stores_ti_and_log_relative_path(self):
        """set_context() stores the task instance and rendered log path for use in close()."""
        self.cloudwatch_task_handler.set_context(self.ti)

        assert self.cloudwatch_task_handler.ti is self.ti
        assert self.cloudwatch_task_handler.log_relative_path != ""
        assert (
            self.cloudwatch_task_handler.log_relative_path == self.cloudwatch_task_handler.io.log_stream_name
        )

    def test_close_calls_upload_once(self):
        """close() calls io.upload() exactly once, even when called multiple times."""
        with mock.patch("airflow.utils.log.file_task_handler.FileTaskHandler.set_context"):
            with mock.patch.object(self.cloudwatch_task_handler.io, "upload") as mock_upload:
                with mock.patch("watchtower.CloudWatchLogHandler.close"):
                    self.cloudwatch_task_handler.set_context(self.ti)
                    for _ in range(3):
                        self.cloudwatch_task_handler.close()

                    mock_upload.assert_called_once_with(
                        self.cloudwatch_task_handler.log_relative_path, self.ti
                    )

    def test_close_skips_upload_without_set_context(self):
        """close() without a prior set_context() should not call io.upload()."""
        with mock.patch.object(self.cloudwatch_task_handler.io, "upload") as mock_upload:
            self.cloudwatch_task_handler.close()
            mock_upload.assert_not_called()

    @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="delete_local_copy only relevant for Airflow 3")
    def test_close_deletes_local_log_dir(self):
        """close() after set_context() deletes the local log directory."""
        log_file_path = Path(self.local_log_location) / self.remote_log_stream
        log_file_path.parent.mkdir(parents=True, exist_ok=True)
        log_file_path.write_text("some log content")
        assert log_file_path.parent.exists()

        with conf_vars({("logging", "delete_local_logs"): "True"}):
            handler = CloudwatchTaskHandler(
                self.local_log_location,
                f"arn:aws:logs:{self.region_name}:11111111:log-group:{self.remote_log_group}",
            )
        with mock.patch("airflow.utils.log.file_task_handler.FileTaskHandler.set_context"):
            with mock.patch("watchtower.CloudWatchLogHandler.close"):
                handler.set_context(self.ti)
                handler.close()

        assert not log_file_path.parent.exists()

    def test_filename_template_for_backward_compatibility(self):
        # filename_template arg support for running the latest provider on airflow 2
        CloudwatchTaskHandler(
            self.local_log_location,
            f"arn:aws:logs:{self.region_name}:11111111:log-group:{self.remote_log_group}",
            filename_template=None,
        )

    def test_event_to_str(self):
        handler = self.cloudwatch_task_handler
        current_time = int(time.time()) * 1000
        events = [
            {"timestamp": current_time - 2000, "message": "First"},
            {"timestamp": current_time - 1000, "message": "Second"},
            {"timestamp": current_time, "message": "Third"},
        ]
        assert [handler._event_to_str(event) for event in events] == (
            [
                f"[{get_time_str(current_time - 2000)}] First",
                f"[{get_time_str(current_time - 1000)}] Second",
                f"[{get_time_str(current_time)}] Third",
            ]
        )


def generate_log_events(conn, log_group_name, log_stream_name, log_events):
    conn.create_log_group(logGroupName=log_group_name)
    conn.create_log_stream(logGroupName=log_group_name, logStreamName=log_stream_name)
    conn.put_log_events(logGroupName=log_group_name, logStreamName=log_stream_name, logEvents=log_events)
