# 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 glob
import logging
import os
import shutil
import tempfile
import unittest
import uuid

import pyarrow as pa

from pypaimon import CatalogFactory, Schema
from pypaimon.api.api_response import ConfigResponse, Partition
from pypaimon.api.auth import BearTokenAuthProvider
from pypaimon.common.options import Options
from pypaimon.catalog.catalog_context import CatalogContext
from pypaimon.catalog.rest.rest_catalog import RESTCatalog
from pypaimon.catalog.rest.table_metadata import TableMetadata
from pypaimon.common.identifier import Identifier
from pypaimon.schema.data_types import (ArrayType, AtomicType, DataField,
                                        MapType)
from pypaimon.schema.table_schema import TableSchema
from pypaimon.tests.rest.rest_server import RESTCatalogServer


class RESTBaseTest(unittest.TestCase):
    def setUp(self):
        self.temp_dir = tempfile.mkdtemp(prefix="unittest_")
        self.warehouse = os.path.join(self.temp_dir, 'warehouse')

        self.config = ConfigResponse(defaults={"prefix": "mock-test"})
        self.token = str(uuid.uuid4())
        self.server = RESTCatalogServer(
            data_path=self.temp_dir,
            auth_provider=BearTokenAuthProvider(self.token),
            config=self.config,
            warehouse="warehouse"
        )
        self.server.start()
        print(f"\nServer started at: {self.server.get_url()}")

        self.options = {
            'metastore': 'rest',
            'uri': f"http://localhost:{self.server.port}",
            'warehouse': "warehouse",
            'dlf.region': 'cn-hangzhou',
            "token.provider": "bear",
            'token': self.token,
        }
        self.rest_catalog = CatalogFactory.create(self.options)
        self.rest_catalog.create_database("default", False)

        self.pa_schema = pa.schema([
            ('user_id', pa.int64()),
            ('item_id', pa.int64()),
            ('behavior', pa.string()),
            ('dt', pa.string()),
            ('long-dt', pa.string())
        ])
        self.pk_pa_schema = pa.schema([
            pa.field('user_id', pa.int64(), nullable=False),
            ('item_id', pa.int64()),
            ('behavior', pa.string()),
            pa.field('dt', pa.string(), nullable=False),
            ('long-dt', pa.string())
        ])
        self.raw_data = {
            'user_id': [1, 2, 3, 4, 5, 6, 7, 8],
            'item_id': [1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008],
            'behavior': ['a', 'b', 'c', None, 'e', 'f', 'g', 'h'],
            'dt': ['p1', 'p1', 'p2', 'p1', 'p2', 'p1', 'p2', 'p2'],
            'long-dt': ['2024-10-10', '2024-10-10', '2024-10-10', '2024-01-01', '2024-10-10', '2025-01-23',
                        'abcdefghijklmnopk', '2025-08-08']
        }
        self.expected = pa.Table.from_pydict(self.raw_data, schema=self.pa_schema)
        self.pk_expected = pa.Table.from_pydict(self.raw_data, schema=self.pk_pa_schema)

        schema = Schema.from_pyarrow_schema(self.pa_schema)
        self.rest_catalog.create_table('default.test_reader_iterator', schema, False)
        self.table = self.rest_catalog.get_table('default.test_reader_iterator')
        write_builder = self.table.new_batch_write_builder()
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        table_write.write_arrow(self.expected)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

    def tearDown(self):
        # Shutdown server
        self.server.shutdown()
        print("Server stopped")
        import gc
        gc.collect()
        shutil.rmtree(self.temp_dir, ignore_errors=True)

    def test_rest_catalog(self):
        """Example usage of RESTCatalogServer"""
        # Setup logging
        logging.basicConfig(level=logging.INFO)

        # Create config
        config = ConfigResponse(defaults={"prefix": "mock-test"})
        token = str(uuid.uuid4())
        # Create server
        server = RESTCatalogServer(
            data_path="/tmp/test_warehouse",
            auth_provider=BearTokenAuthProvider(token),
            config=config,
            warehouse="test_warehouse"
        )
        try:
            # Start server
            server.start()
            print(f"Server started at: {server.get_url()}")
            test_databases = {
                "default": server.mock_database("default", {"env": "test"}),
                "test_db1": server.mock_database("test_db1", {"env": "test"}),
                "test_db2": server.mock_database("test_db2", {"env": "test"}),
                "prod_db": server.mock_database("prod_db", {"env": "prod"})
            }
            data_fields = [
                DataField(0, "name", AtomicType('INT'), 'desc  name'),
                DataField(1, "arr11", ArrayType(True, AtomicType('INT')), 'desc  arr11'),
                DataField(2, "map11", MapType(False, AtomicType('INT'),
                                              MapType(False, AtomicType('INT'), AtomicType('INT'))),
                          'desc  arr11'),
            ]
            schema = TableSchema(TableSchema.CURRENT_VERSION, len(data_fields), data_fields, len(data_fields),
                                 [], [], {}, "")
            test_tables = {
                "default.user": TableMetadata(uuid=str(uuid.uuid4()), is_external=True, schema=schema),
            }
            server.table_metadata_store.update(test_tables)
            server.database_store.update(test_databases)
            options = {
                'uri': f"http://localhost:{server.port}",
                'warehouse': 'test_warehouse',
                'dlf.region': 'cn-hangzhou',
                "token.provider": "bear",
                'token': token
            }
            rest_catalog = RESTCatalog(CatalogContext.create_from_options(Options(options)))
            self.assertSetEqual(set(rest_catalog.list_databases()), {*test_databases})
            self.assertEqual(rest_catalog.get_database('default').name, test_databases.get('default').name)
            table = rest_catalog.get_table(Identifier.from_string('default.user'))
            self.assertEqual(table.identifier.get_full_name(), 'default.user')
        finally:
            # Shutdown server
            server.shutdown()
            print("Server stopped")

    def test_write(self):
        # Setup logging
        logging.basicConfig(level=logging.INFO)

        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'])
        self.rest_catalog.drop_table("default.test_table", True)
        self.rest_catalog.create_table("default.test_table", schema, False)
        table = self.rest_catalog.get_table("default.test_table")

        data = {
            'user_id': [1, 2, 3, 4],
            'item_id': [1001, 1002, 1003, 1004],
            'behavior': ['a', 'b', 'c', None],
            'dt': ['p1', 'p1', 'p2', 'p1'],
            'long-dt': ['2024-10-10', '2024-10-10', '2024-10-10', '2024-01-01']
        }
        expect = pa.Table.from_pydict(data, schema=self.pa_schema)

        write_builder = table.new_batch_write_builder()
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        table_write.write_arrow(expect)
        commit_messages = table_write.prepare_commit()
        table_commit.commit(commit_messages)
        table_write.close()
        table_commit.close()

        self.assertTrue(os.path.exists(self.warehouse + "/default/test_table/snapshot/LATEST"))
        self.assertTrue(os.path.exists(self.warehouse + "/default/test_table/snapshot/snapshot-1"))
        self.assertTrue(os.path.exists(self.warehouse + "/default/test_table/manifest"))
        self.assertTrue(os.path.exists(self.warehouse + "/default/test_table/dt=p1"))
        self.assertEqual(len(glob.glob(self.warehouse + "/default/test_table/manifest/*")), 3)

    def _write_test_table(self, table):
        write_builder = table.new_batch_write_builder()
        table_pa_schema = self.pk_pa_schema if table.primary_keys else self.pa_schema

        # first write
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data1 = {
            'user_id': [1, 2, 3, 4],
            'item_id': [1001, 1002, 1003, 1004],
            'behavior': ['a', 'b', 'c', None],
            'dt': ['p1', 'p1', 'p2', 'p1'],
            'long-dt': ['2024-10-10', '2024-10-10', '2024-10-10', '2024-01-01'],
        }
        pa_table = pa.Table.from_pydict(data1, schema=table_pa_schema)
        table_write.write_arrow(pa_table)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        # second write
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data2 = {
            'user_id': [5, 6, 7, 8],
            'item_id': [1005, 1006, 1007, 1008],
            'behavior': ['e', 'f', 'g', 'h'],
            'dt': ['p2', 'p1', 'p2', 'p2'],
            'long-dt': ['2024-10-10', '2025-01-23', 'abcdefghijklmnopk', '2025-08-08'],
        }
        pa_table = pa.Table.from_pydict(data2, schema=table_pa_schema)
        table_write.write_arrow(pa_table)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

    def _read_test_table(self, read_builder):
        table_read = read_builder.new_read()
        splits = read_builder.new_scan().plan().splits()
        return table_read.to_arrow(splits)

    def _write_test_table_with_schema(self, table, pa_schema):
        """Write test data using the specified PyArrow schema."""
        write_builder = table.new_batch_write_builder()

        # first write
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data1 = {
            'user_id': [1, 2, 3, 4],
            'item_id': [1001, 1002, 1003, 1004],
            'behavior': ['a', 'b', 'c', None],
            'dt': ['p1', 'p1', 'p2', 'p1'],
            'long-dt': ['2024-10-10', '2024-10-10', '2024-10-10', '2024-01-01'],
        }
        pa_table = pa.Table.from_pydict(data1, schema=pa_schema)
        table_write.write_arrow(pa_table)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        # second write
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data2 = {
            'user_id': [5, 6, 7, 8],
            'item_id': [1005, 1006, 1007, 1008],
            'behavior': ['e', 'f', 'g', 'h'],
            'dt': ['p2', 'p1', 'p2', 'p2'],
            'long-dt': ['2024-10-10', '2025-01-23', 'abcdefghijklmnopk', '2025-08-08'],
        }
        pa_table = pa.Table.from_pydict(data2, schema=pa_schema)
        table_write.write_arrow(pa_table)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

    def test_list_partitions_paged(self):
        """Test list_partitions_paged returns partitions from server store."""
        identifier = Identifier.from_string('default.test_reader_iterator')

        # Add test partitions to server store
        p1 = Partition(
            spec={"dt": "p1"},
            record_count=4,
            file_size_in_bytes=1024,
            file_count=1,
            last_file_creation_time=1000,
            total_buckets=1,
            done=False,
        )
        p2 = Partition(
            spec={"dt": "p2"},
            record_count=4,
            file_size_in_bytes=2048,
            file_count=1,
            last_file_creation_time=2000,
            total_buckets=1,
            done=True,
        )
        p3 = Partition(
            spec={"dt": "p3"},
            record_count=2,
            file_size_in_bytes=512,
            file_count=1,
            last_file_creation_time=3000,
            total_buckets=1,
            done=False,
        )
        self.server.table_partitions_store[identifier.get_full_name()] = [p1, p2, p3]

        # Test: list all partitions
        result = self.rest_catalog.list_partitions_paged(identifier)
        self.assertEqual(len(result.elements), 3)

        # Test: list with max_results
        result = self.rest_catalog.list_partitions_paged(identifier, max_results=2)
        self.assertEqual(len(result.elements), 2)
        self.assertIsNotNone(result.next_page_token)

        # Test: list next page using page_token
        result2 = self.rest_catalog.list_partitions_paged(
            identifier, max_results=2, page_token=result.next_page_token
        )
        self.assertEqual(len(result2.elements), 1)
        self.assertIsNone(result2.next_page_token)

        # Test: list with pattern filter
        result = self.rest_catalog.list_partitions_paged(
            identifier, partition_name_pattern="dt=p1"
        )
        self.assertEqual(len(result.elements), 1)
        self.assertEqual(result.elements[0].spec, {"dt": "p1"})

        # Test: list with pattern using wildcard
        result = self.rest_catalog.list_partitions_paged(
            identifier, partition_name_pattern="dt=p%"
        )
        self.assertEqual(len(result.elements), 3)

    def test_alter_database(self):
        """Test alter_database sets and removes properties."""
        from pypaimon.catalog.rest.property_change import PropertyChange
        db_name = "alter_db_test"
        self.rest_catalog.create_database(db_name, True)

        # set property
        self.rest_catalog.alter_database(
            db_name,
            [PropertyChange.set_property("key1", "value1"),
             PropertyChange.set_property("key2", "value2")])
        db = self.rest_catalog.get_database(db_name)
        self.assertEqual(db.options.get("key1"), "value1")
        self.assertEqual(db.options.get("key2"), "value2")

        # remove property
        self.rest_catalog.alter_database(
            db_name,
            [PropertyChange.remove_property("key1")])
        db = self.rest_catalog.get_database(db_name)
        self.assertNotIn("key1", db.options)
        self.assertEqual(db.options.get("key2"), "value2")

    def test_list_partitions_paged_empty(self):
        """Test list_partitions_paged returns empty when no partitions."""
        identifier = Identifier.from_string('default.test_reader_iterator')
        result = self.rest_catalog.list_partitions_paged(identifier)
        self.assertEqual(len(result.elements), 0)
        self.assertIsNone(result.next_page_token)
