# 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.

import os
from functools import wraps

import pytest


_NATIVE_PLAN_ENV = "PYPAIMON_TEST_NATIVE_PLAN"
_NATIVE_READ_ENV = "PYPAIMON_TEST_NATIVE_READ"
_NATIVE_WRITE_ENV = "PYPAIMON_TEST_NATIVE_WRITE"
_NATIVE_UPDATE_ENV = "PYPAIMON_TEST_NATIVE_UPDATE"
_NATIVE_COMMIT_ENV = "PYPAIMON_TEST_NATIVE_COMMIT"
_native_plan_count = 0
_native_read_count = 0
_native_write_count = 0
_native_commit_count = 0
_native_update_counts = dict.fromkeys(('row_id', 'grouped', 'predicate', 'upsert', 'incremental', 'merge', 'shard'), 0)
_force_native_for_test = False
_force_native_read_for_test = False
_force_native_write_for_test = False
_force_native_commit_for_test = False
_force_native_update_for_test = False


@pytest.fixture
def native_rest_catalog(tmp_path):
    """A local REST catalog for native writer and committer integration tests."""
    import uuid

    from pypaimon import CatalogFactory
    from pypaimon.api.api_response import ConfigResponse
    from pypaimon.api.auth import BearTokenAuthProvider
    from pypaimon.tests.rest.rest_server import RESTCatalogServer

    token = str(uuid.uuid4())
    server = RESTCatalogServer(
        data_path=str(tmp_path), auth_provider=BearTokenAuthProvider(token),
        config=ConfigResponse(defaults={'prefix': 'native-test'}), warehouse='warehouse')
    server.start()
    try:
        catalog = CatalogFactory.create({
            'metastore': 'rest', 'uri': server.get_url(), 'warehouse': 'warehouse',
            'token.provider': 'bear', 'token': token, 'data-token.enabled': 'false'})
        catalog.create_database('default', True)
        yield catalog
    finally:
        server.shutdown()


def _native_plan_enabled():
    return os.environ.get(_NATIVE_PLAN_ENV) == "1"


def _native_read_enabled():
    return os.environ.get(_NATIVE_READ_ENV) == "1"


def _native_write_enabled():
    return os.environ.get(_NATIVE_WRITE_ENV) == "1"


def _native_commit_enabled():
    return os.environ.get(_NATIVE_COMMIT_ENV) == "1"


def _native_update_enabled():
    return os.environ.get(_NATIVE_UPDATE_ENV) == "1"


def pytest_configure(config):
    config.addinivalue_line(
        "markers", "python_plan: keep Python planner assertions on the Python lane")
    config.addinivalue_line(
        "markers", "python_read: keep Python reader assertions on the Python lane")
    config.addinivalue_line(
        "markers", "native_plan: exercise the real Rust planner in the Rust main CI job")
    config.addinivalue_line(
        "markers", "python_write: keep Python writer assertions on the Python lane")
    config.addinivalue_line(
        "markers", "python_commit: keep Python committer assertions on the Python lane")
    if _native_plan_enabled():
        from pypaimon.execution.native.read.plan import NativeTableScan

        original_plan = NativeTableScan._plan

        def tracked_plan(self, *args, **kwargs):
            global _native_plan_count
            plan = original_plan(self, *args, **kwargs)
            if plan is not None and _force_native_for_test:
                _native_plan_count += 1
            return plan

        NativeTableScan._plan = tracked_plan

        from pypaimon.execution.native.read.streaming_table_scan import StreamTableScan

        original_stream_plan = StreamTableScan.plan

        def tracked_stream_plan(self, *args, **kwargs):
            global _native_plan_count
            plan = original_stream_plan(self, *args, **kwargs)
            if plan is not None and _force_native_for_test:
                _native_plan_count += 1
            return plan

        StreamTableScan.plan = tracked_stream_plan

    if _native_read_enabled():
        from pypaimon.execution.native.read.table_read import NativeTableRead

        original_read = NativeTableRead.read_batches

        @wraps(original_read)
        def tracked_read(self, splits, *args, **kwargs):
            global _native_read_count
            batches = original_read(self, splits, *args, **kwargs)
            if batches is not None and splits and _force_native_read_for_test:
                _native_read_count += 1
            return batches

        NativeTableRead.read_batches = tracked_read

    if _native_write_enabled():
        from pypaimon.execution.native.write.native_write import NativeTableWrite

        def track_write(method):
            original_write = getattr(NativeTableWrite, method)

            @wraps(original_write)
            def tracked_write(self, data, *args, **kwargs):
                global _native_write_count
                native = self._native_writer is not None
                result = original_write(self, data, *args, **kwargs)
                if native and data.num_rows and _force_native_write_for_test:
                    _native_write_count += 1
                return result

            setattr(NativeTableWrite, method, tracked_write)

        track_write('write_arrow_batch')

    if _native_commit_enabled():
        from pypaimon.write.table_commit import TableCommit

        original_prepare = TableCommit._prepare_native_commit

        def tracked_prepare(self, messages, *args, **kwargs):
            global _native_commit_count
            prepared = original_prepare(self, messages, *args, **kwargs)
            if prepared is not None and _force_native_commit_for_test:
                _native_commit_count += 1
            return prepared

        TableCommit._prepare_native_commit = tracked_prepare

    if _native_update_enabled():
        from pypaimon.execution.native.write.native_merge_into import NativeTableMergeInto
        from pypaimon.execution.native.write.native_update import (
            NativeBatchTableUpdate, NativePredicateTableUpdate,
            NativeTableUpdateByRowId, NativeTableUpsert, NativeShardTableUpdator,
        )

        def track_update(cls, method, kind):
            original = getattr(cls, method)

            def tracked(self, *args, **kwargs):
                messages = original(self, *args, **kwargs)
                if messages and _force_native_update_for_test:
                    _native_update_counts[kind] += 1
                return messages

            setattr(cls, method, tracked)

        track_update(NativeBatchTableUpdate, 'update_by_arrow_with_row_id', 'row_id')
        track_update(NativeBatchTableUpdate, 'update_by_arrow_batches_with_row_id', 'grouped')
        track_update(NativePredicateTableUpdate, 'update', 'predicate')
        track_update(NativeTableUpsert, 'upsert', 'upsert')
        track_update(NativeTableUpdateByRowId, 'update_columns', 'incremental')
        track_update(NativeTableMergeInto, 'prepare_commit', 'merge')
        track_update(NativeShardTableUpdator, 'prepare_commit', 'shard')


def pytest_collection_modifyitems(items):
    if _native_plan_enabled():
        return
    skip_native = pytest.mark.skip(reason="native plan tests run in the Rust Plan job")
    for item in items:
        if item.get_closest_marker("native_plan") is not None:
            item.add_marker(skip_native)


@pytest.fixture(autouse=True)
def enable_native_backends(request, monkeypatch):
    global _force_native_for_test, _force_native_read_for_test
    global _force_native_write_for_test, _force_native_commit_for_test
    global _force_native_update_for_test
    python_plan = request.node.get_closest_marker("python_plan") is not None
    python_read = request.node.get_closest_marker("python_read") is not None
    python_write = request.node.get_closest_marker("python_write") is not None
    python_commit = request.node.get_closest_marker("python_commit") is not None
    # request.path requires pytest 7; Python 3.6 uses pytest 6.
    native_plan_test = os.path.basename(request.node.location[0]) in (
        "native_plan_test.py", "native_plan_integration_test.py",
        "native_plan_capabilities_test.py")
    force_plan = _native_plan_enabled() and not python_plan and not native_plan_test
    force_read = (_native_read_enabled() and not python_plan and not python_read
                  and not native_plan_test)
    force_write = _native_write_enabled() and not python_write
    force_commit = _native_commit_enabled() and not python_commit
    force_update = _native_update_enabled() and not python_write
    if not (force_plan or force_read or force_write or force_commit or force_update):
        yield
        return

    from pypaimon.common.options.core_options import CoreOptions

    if force_plan:
        original_plan = CoreOptions.native_plan_enabled

        def plan_enabled(self, default=None):
            return original_plan(self, True if default is None else default)

        monkeypatch.setattr(CoreOptions, "native_plan_enabled", plan_enabled)
    if force_read:
        original_read = CoreOptions.native_read_enabled

        def read_enabled(self, default=None):
            return original_read(self, True if default is None else default)

        monkeypatch.setattr(CoreOptions, "native_read_enabled", read_enabled)
    if force_write or force_update:
        original_write = CoreOptions.native_write_enabled

        def write_enabled(self, default=None):
            return original_write(self, True if default is None else default)

        monkeypatch.setattr(CoreOptions, "native_write_enabled", write_enabled)
    if force_commit:
        original_commit = CoreOptions.native_commit_enabled

        def commit_enabled(self, default=None):
            return original_commit(self, True if default is None else default)

        monkeypatch.setattr(CoreOptions, "native_commit_enabled", commit_enabled)
    _force_native_for_test = force_plan
    _force_native_read_for_test = force_read
    _force_native_write_for_test = force_write
    _force_native_commit_for_test = force_commit
    _force_native_update_for_test = force_update
    try:
        yield
    finally:
        _force_native_for_test = False
        _force_native_read_for_test = False
        _force_native_write_for_test = False
        _force_native_commit_for_test = False
        _force_native_update_for_test = False


def pytest_sessionfinish(session, exitstatus):
    if exitstatus == 0:
        if ((_native_plan_enabled() and _native_plan_count == 0)
                or (_native_read_enabled() and _native_read_count == 0)
                or (_native_write_enabled() and _native_write_count == 0)
                or (_native_update_enabled() and not all(_native_update_counts.values()))
                or (_native_commit_enabled() and _native_commit_count == 0)):
            session.exitstatus = pytest.ExitCode.TESTS_FAILED


def pytest_terminal_summary(terminalreporter):
    if _native_plan_enabled():
        terminalreporter.write_line(
            "native plans exercised: %d" % _native_plan_count)
    if _native_read_enabled():
        terminalreporter.write_line(
            "native reads exercised: %d" % _native_read_count)
    if _native_write_enabled():
        terminalreporter.write_line(
            "native writes exercised: %d" % _native_write_count)
    if _native_commit_enabled():
        terminalreporter.write_line(
            "native commits exercised: %d" % _native_commit_count)
    if _native_update_enabled():
        terminalreporter.write_line(
            "native updates exercised: " + ', '.join(
                '%s=%d' % item for item in _native_update_counts.items()))