# Copyright 2021-present MongoDB, 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 concurrent.futures
import threading
import unittest
import unittest.mock as mock
import uuid
from datetime import datetime
from test import client_context
from test.utils import AllowListEventListener

import bson
import pyarrow as pa
import pytest
from pyarrow import int32, int64
from pymongo import DESCENDING, WriteConcern
from pymongo.collection import Collection

from pymongoarrow import api
from pymongoarrow.api import Schema, aggregate_polars_all, find_arrow_all, find_polars_all, write
from pymongoarrow.errors import ArrowWriteError
from pymongoarrow.types import (
    _TYPE_NORMALIZER_FACTORY,
    BinaryType,
    CodeType,
    Decimal128Type,
    ObjectIdType,
)

try:
    import polars as pl
    from polars.testing import assert_frame_equal
except ImportError:
    pl = None

if pl is None:
    pytest.skip("Requires polars.", allow_module_level=True)


class PolarsTestBase(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        if not client_context.connected:
            raise unittest.SkipTest("cannot connect to MongoDB")
        cls.cmd_listener = AllowListEventListener("find", "aggregate")
        cls.getmore_listener = AllowListEventListener("getMore")
        cls.client = client_context.get_client(
            event_listeners=[cls.getmore_listener, cls.cmd_listener],
            uuidRepresentation="standard",
        )
        cls.addClassCleanup(cls.client.close)


class TestExplicitPolarsApi(PolarsTestBase):
    @classmethod
    def setUpClass(cls):
        PolarsTestBase.setUpClass()
        cls.schema = Schema({"_id": int32(), "data": int64()})
        cls.coll = cls.client.pymongoarrow_test.get_collection(
            "test", write_concern=WriteConcern(w="majority")
        )

    def setUp(self):
        """Insert simple use case data."""
        self.coll.drop()
        self.data = [
            {"_id": 1, "data": 10},
            {"_id": 2, "data": 20},
            {"_id": 3, "data": 30},
            {"_id": 4},
        ]
        self.coll.insert_many(self.data)
        self.cmd_listener.reset()
        self.getmore_listener.reset()

    def round_trip(self, df_in, schema=None, **kwargs):
        """Helper tests pl.DataFrame written matches that found."""
        self.coll.drop()
        res = write(self.coll, df_in)
        self.assertEqual(len(df_in), res.raw_result["insertedCount"])
        df_out = find_polars_all(self.coll, {}, schema=schema, **kwargs)
        pl.testing.assert_frame_equal(df_in, df_out)
        return res

    def test_find_simple(self):
        expected = pl.DataFrame(
            data={
                "_id": pl.Series(values=[1, 2, 3, 4], dtype=pl.Int32),
                "data": pl.Series(values=[10, 20, 30, None], dtype=pl.Int64),
            }
        )
        table = find_polars_all(self.coll, {}, schema=self.schema)
        self.assertEqual(expected.dtypes, table.dtypes)
        self.assertTrue(table.equals(expected))

        expected = pl.DataFrame(
            data={
                "_id": pl.Series(values=[4, 3], dtype=pl.Int32),
                "data": pl.Series(values=[None, 30], dtype=pl.Int64),
            }
        )
        table = find_polars_all(
            self.coll,
            {"_id": {"$gt": 2}},
            schema=self.schema,
            sort=[("_id", DESCENDING)],
        )
        self.assertEqual(expected.dtypes, table.dtypes)
        self.assertTrue(table.equals(expected))

        find_cmd = self.cmd_listener.results["started"][-1]
        self.assertEqual(find_cmd.command_name, "find")
        self.assertEqual(find_cmd.command["projection"], {"_id": True, "data": True})

    def _check_aggregation_simple(self, coll):
        expected = pl.DataFrame(
            data={
                "_id": pl.Series(values=[1, 2, 3, 4], dtype=pl.Int32),
                "data": pl.Series(values=[20, 40, 60, None], dtype=pl.Int64),
            }
        )
        coll.drop()
        coll.insert_many(self.data)
        projection = {"_id": True, "data": {"$multiply": [2, "$data"]}}
        table = aggregate_polars_all(coll, [{"$project": projection}], schema=self.schema)
        self.assertTrue(table.equals(expected))

        agg_cmd = self.cmd_listener.results["started"][-1]
        self.assertEqual(agg_cmd.command_name, "aggregate")
        assert len(agg_cmd.command["pipeline"]) == 2
        self.assertEqual(agg_cmd.command["pipeline"][0]["$project"], projection)
        self.assertEqual(agg_cmd.command["pipeline"][1]["$project"], {"_id": True, "data": True})

    def test_aggregate_simple(self):
        self._check_aggregation_simple(self.coll)

    @mock.patch.object(Collection, "insert_many", side_effect=Collection.insert_many, autospec=True)
    def test_write_batching(self, mock):
        data = pl.DataFrame(data={"_id": pl.Series(values=range(100040), dtype=pl.Int64)})
        self.round_trip(data, Schema(dict(_id=int64())))
        self.assertEqual(mock.call_count, 2)

    def test_duplicate_key_error(self):
        """Confirm expected error is raised, simple duplicate key case."""
        n = 3
        data = pl.DataFrame(
            data={
                "_id": pl.Series(values=list(range(n)) * 2, dtype=pl.Int32),
                "data": pl.Series(values=range(n * 2), dtype=pl.Int64),
            }
        )
        with self.assertRaises(ArrowWriteError):
            try:
                self.round_trip(data, Schema({"_id": int32(), "data": int64()}))
            except ArrowWriteError as awe:
                self.assertEqual(n, awe.details["writeErrors"][0]["index"])
                self.assertEqual(n, awe.details["nInserted"])
                raise awe

    def test_polars_types(self):
        """Test round-trip of DataFrame consisting of Polar.DataTypes.

        This does NOT include ExtensionTypes as Polars doesn't support them (yet).
        """
        pl_typenames = ["Int64", "Int32", "Float64", "Datetime", "String", "Boolean"]
        pl_types = [pl.Int64, pl.Int32, pl.Float64, pl.Datetime("ms"), pl.String, pl.Boolean]
        pa_types = [
            pa.int64(),
            pa.int32(),
            pa.float64(),
            pa.timestamp("ms"),
            pa.string(),
            pa.bool_(),
        ]

        pl_schema = dict(zip(pl_typenames, pl_types))
        pa_schema = dict(zip(pl_typenames, pa_types))

        data = {
            "Int64": pl.Series([i for i in range(2)] + [None]),
            "Int32": pl.Series([i for i in range(2)] + [None]),
            "Float64": pl.Series([i for i in range(2)] + [None]),
            "Datetime": pl.Series([datetime(1970 + i, 1, 1) for i in range(2)] + [None]),
            "String": pl.Series([f"a{i}" for i in range(2)] + [None]),
            "Boolean": pl.Series([True, False, None]),
        }

        df_in = pl.from_dict(data=data, schema=pl_schema)
        self.coll.drop()
        write(self.coll, df_in)
        df_out = find_polars_all(self.coll, {}, schema=Schema(pa_schema))
        pl.testing.assert_frame_equal(df_in, df_out)

    def test_extension_types_fail(self):
        """Confirm failure on ExtensionTypes for Polars.DataFrame.from_arrow"""

        for ext_type, data in (
            (ObjectIdType(), [bson.ObjectId().binary, bson.ObjectId().binary]),
            (Decimal128Type(), [bson.Decimal128(str(i)).bid for i in range(2)]),
            (CodeType(), [str(i) for i in range(2)]),
        ):
            table = pa.Table.from_pydict({"foo": data}, pa.schema({"foo": ext_type}))
            with self.assertRaises(pl.exceptions.ComputeError):
                pl.from_arrow(table)

    def test_auto_schema_succeeds_on_find(self):
        """Confirms Polars can read ObjectID Extension type.

        This is inserted automatically by Collection.insert_many
        Note that the output dtype is int32
        """
        vals = [1, "2", True, 4]
        data = [{"a": v} for v in vals]

        self.coll.drop()
        self.coll.insert_many(data)  # ObjectID autogenerated here

        df_out = find_polars_all(self.coll, {})
        self.assertEqual(df_out.columns, ["_id", "a"])
        self.assertEqual(df_out.shape, (4, 2))
        self.assertEqual(df_out.dtypes, [pl.Binary, pl.Int32])

    def test_arrow_to_polars(self):
        """Test reading Polars data from written Arrow Data."""
        arrow_schema = {k.__name__: v(True) for k, v in _TYPE_NORMALIZER_FACTORY.items()}
        arrow_table_in = pa.Table.from_pydict(
            {
                "Int64": [i for i in range(2)],
                "float": [i for i in range(2)],
                "datetime": [i for i in range(2)],
                "str": [str(i) for i in range(2)],
                "int": [i for i in range(2)],
                "bool": [True, False],
                "struct": [{"objId": bson.ObjectId().binary, "str1": str(i)} for i in range(2)],
                "list": [[str(i), str(i + 1)] for i in range(2)],
                "Binary": [b"1", b"23"],
                "ObjectId": [bson.ObjectId().binary, bson.ObjectId().binary],
                "Decimal128": [bson.Decimal128(str(i)).bid for i in range(2)],
                "Code": [str(i) for i in range(2)],
            },
            pa.schema(arrow_schema),
        )

        self.coll.drop()
        res = write(self.coll, arrow_table_in)
        self.assertEqual(len(arrow_table_in), res.raw_result["insertedCount"])
        df_out = find_polars_all(self.coll, query={}, schema=Schema(arrow_schema))

        # Sanity check: compare with _arrow_to_polars
        df_actual_output = api._arrow_to_polars(arrow_table_in)
        assert_frame_equal(df_out, df_actual_output)

    def test_exceptions_for_unsupported_polar_types(self):
        """Confirm exceptions thrown are expected.

        Currently, pl.Series, and pl.Object
        Tracks future changes in any packages.
        """

        # Series: PyMongoError does not support.
        with self.assertRaises(ValueError) as exc:
            pls = pl.Series(values=range(2))
            write(self.coll, pls)
        self.assertTrue("Invalid tabular data object" in exc.exception.args[0])

        # Polars has an Object Type, similar in concept to Pandas. PyMongoError does not support them.
        class MyObject:
            pass

        with self.assertRaises(ValueError) as exc:
            df_in = pl.DataFrame(data=[MyObject()] * 2)
            write(self.coll, df_in)

    def test_polars_binary_type(self):
        """Demonstrates that binary data is not yet supported. TODO [ARROW-214]

        Will demonstrate roundtrip behavior of Polar Binary Type.
        """
        # 1. _id added by MongoDB
        self.coll.drop()
        with self.assertRaises(ValueError):
            df_in = pl.DataFrame({"Binary": [b"1", b"one"]}, schema={"Binary": pl.Binary})
            write(self.coll, df_in)
            df_out = find_polars_all(self.coll, {})
            self.assertTrue(df_out.columns == ["_id", "Binary"])
            self.assertTrue(all([isinstance(c, pl.Binary) for c in df_out.dtypes]))
            self.assertIsNone(assert_frame_equal(df_in, df_out.select("Binary")))
            # 2. Explicit Binary _id
            self.coll.drop()
            df_in = pl.DataFrame(
                data=dict(_id=[b"0", b"1"], Binary=[b"1", b"one"]),
                schema=dict(_id=pl.Binary, Binary=pl.Binary),
            )
            write(self.coll, df_in)
            df_out = find_polars_all(self.coll, {})
            self.assertEqual(df_out.columns, ["_id", "Binary"])
            self.assertTrue(all([isinstance(c, pl.Binary) for c in df_out.dtypes]))
            self.assertIsNone(assert_frame_equal(df_in, df_out))
            # 3. Explicit Int32 _id
            self.coll.drop()
            df_in = pl.DataFrame(
                data={"_id": [0, 1], "Binary": [b"1", b"one"]},
                schema={"_id": pl.Int32, "Binary": pl.Binary},
            )
            write(self.coll, df_in)
            df_out = find_polars_all(self.coll, {})
            self.assertEqual(df_out.columns, ["_id", "Binary"])
            out_types = df_out.dtypes
            self.assertTrue(isinstance(out_types[0], pl.Int32))
            self.assertTrue(isinstance(out_types[1], pl.Binary))
            self.assertTrue(assert_frame_equal(df_in, df_out) is None)

    def test_bson_types(self):
        """Test reading Polars and Arrow data from written BSON Data.

        This is meant to capture the use case of reading data in Python
        that has been written to a DB from another language.

        Note that this tests only types currently supported by Arrow.
        bson.Regex is not included, for example.
        """

        # 1. Use pymongo / bson packages to build create and write tabular data
        self.coll.drop()
        collection = self.coll

        data_type_map = [
            {"type": "int", "value": 42, "atype": pa.int32(), "ptype": pl.Int32},
            {"type": "long", "value": 1234567890123456789, "atype": pa.int64(), "ptype": pl.Int64},
            {"type": "double", "value": 10.5, "atype": pa.float64(), "ptype": pl.Float64},
            {"type": "string", "value": "hello world", "atype": pa.string(), "ptype": pl.String},
            {"type": "boolean", "value": True, "atype": pa.bool_(), "ptype": pl.Boolean},
            {
                "type": "date",
                "value": datetime(2025, 1, 21),
                "atype": pa.timestamp("ms"),
                "ptype": pl.Datetime,
            },
            {
                "type": "object",
                "value": {"a": 1, "b": 2},
                "atype": pa.struct({"a": pa.int32(), "b": pa.int32()}),
                "ptype": pl.Struct({"a": pl.Int32, "b": pl.Int32}),
            },
            {
                "type": "array",
                "value": [1, 2, 3],
                "atype": pa.list_(pa.int32()),
                "ptype": pl.List(pl.Int32),
            },
            {
                "type": "bytes",
                "value": b"\x00\x01\x02\x03\x04",
                "atype": BinaryType(pa.binary()),
                "ptype": pl.Binary,
            },
            {
                "type": "binary data",
                "value": bson.Binary(b"\x00\x01\x02\x03\x04"),
                "atype": BinaryType(pa.binary()),
                "ptype": pl.Binary,
            },
            {
                "type": "object id",
                "value": bson.ObjectId(),
                "atype": ObjectIdType(),
                "ptype": pl.Object,
            },
            {
                "type": "javascript",
                "value": bson.Code("function() { return x; }"),
                "atype": CodeType(),
                "ptype": pl.String,
            },
            {
                "type": "decimal128",
                "value": bson.Decimal128("10.99"),
                "atype": Decimal128Type(),
                "ptype": pl.Decimal,
            },
            {
                "type": "uuid",
                "value": uuid.uuid4(),
                "atype": BinaryType(pa.binary()),
                "ptype": pl.Binary,
            },
        ]

        # Iterate over types
        for data_type in data_type_map:
            collection.insert_one({"data_type": data_type["type"], "value": data_type["value"]})
            table = find_arrow_all(collection=collection, query={"data_type": data_type["type"]})
            assert table.shape == (1, 3)
            assert table["value"].type == data_type["atype"]
            try:
                dfpl = pl.from_arrow(table.drop("_id"))
                assert dfpl["value"].dtype == data_type["ptype"]
            except pl.exceptions.ComputeError:
                assert isinstance(table["value"].type, pa.ExtensionType)

    def test_threading(self):
        def run_test():
            client = client_context.get_client(
                event_listeners=[self.getmore_listener, self.cmd_listener]
            )
            name = f"test-{threading.current_thread().name}"
            coll = client.pymongoarrow_test.get_collection(
                name, write_concern=WriteConcern(w="majority")
            )
            self._check_aggregation_simple(coll)
            client.close()

        with concurrent.futures.ThreadPoolExecutor() as executor:
            futures = []
            for i in range(5):
                futures.append(executor.submit(run_test))
            concurrent.futures.wait(futures)
            for future in futures:
                future.result()
