# 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 datetime import date

import numpy as np
import pandas as pd
import pytest

from tsfile import ColumnSchema, TableSchema, TSDataType, TIME_COLUMN
from tsfile import TsFileTableWriter, ColumnCategory
from tsfile import to_dataframe
from tsfile.exceptions import ColumnNotExistError, TypeMismatchError
from tsfile.tsfile_table_writer import (
    validate_dataframe_for_tsfile,
    infer_object_column_type,
)


def convert_to_nullable_types(df):
    """
    Convert DataFrame columns to nullable types to match returned DataFrame from to_dataframe.
    This handles the fact that returned DataFrames use nullable types (Int64, Float64, etc.)
    to support Null values.
    """
    df = df.copy()
    for col in df.columns:
        dtype = df[col].dtype
        if dtype == "int64":
            df[col] = df[col].astype("Int64")
        elif dtype == "int32":
            df[col] = df[col].astype("Int32")
        elif dtype == "float64":
            df[col] = df[col].astype("Float64")
        elif dtype == "float32":
            df[col] = df[col].astype("Float32")
        elif dtype == "bool":
            df[col] = df[col].astype("boolean")
        elif pd.api.types.is_object_dtype(df[col]):
            non_null = df[col].dropna()
            if len(non_null) and non_null.map(lambda x: isinstance(x, str)).all():
                df[col] = df[col].astype("string")
        elif pd.api.types.is_string_dtype(df[col]):
            # NumPy/pandas str dtype (e.g. Py3.14) vs object: unify for Series.equals
            df[col] = df[col].astype("string")
    return df


def test_infer_object_column_type_bool():
    """infer_object_column_type should infer BOOLEAN for object column containing bool values."""
    s_true = pd.Series([True, False], dtype=object)
    assert infer_object_column_type(s_true) == TSDataType.BOOLEAN
    s_false = pd.Series([False], dtype=object)
    assert infer_object_column_type(s_false) == TSDataType.BOOLEAN


def test_write_dataframe_basic():
    table = TableSchema(
        "test_table",
        [
            ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG),
            ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD),
            ColumnSchema("value2", TSDataType.INT64, ColumnCategory.FIELD),
        ],
    )
    tsfile_path = "test_write_dataframe_basic.tsfile"
    try:
        if os.path.exists(tsfile_path):
            os.remove(tsfile_path)

        with TsFileTableWriter(tsfile_path, table) as writer:
            df = pd.DataFrame(
                {
                    "time": [i for i in range(100)],
                    "device": [f"device{i}" for i in range(100)],
                    "value": [i * 1.5 for i in range(100)],
                    "value2": [i * 10 for i in range(100)],
                }
            )
            writer.write_dataframe(df)

        df_read = to_dataframe(tsfile_path, table_name="test_table")
        df_read = df_read.sort_values(TIME_COLUMN).reset_index(drop=True)
        df_sorted = convert_to_nullable_types(
            df.sort_values("time").reset_index(drop=True)
        )
        df_read = convert_to_nullable_types(df_read)
        assert df_read.shape == (100, 4)
        assert df_read[TIME_COLUMN].equals(df_sorted["time"])
        assert df_read["device"].equals(df_sorted["device"])
        assert df_read["value"].equals(df_sorted["value"])
        assert df_read["value2"].equals(df_sorted["value2"])
    finally:
        if os.path.exists(tsfile_path):
            os.remove(tsfile_path)


def test_write_dataframe_with_index():
    table = TableSchema(
        "test_table",
        [
            ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG),
            ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD),
        ],
    )
    tsfile_path = "test_write_dataframe_index.tsfile"
    try:
        if os.path.exists(tsfile_path):
            os.remove(tsfile_path)

        with TsFileTableWriter(tsfile_path, table) as writer:
            df = pd.DataFrame(
                {
                    "device": [f"device{i}" for i in range(50)],
                    "value": [i * 2.5 for i in range(50)],
                }
            )
            df.index = [i * 10 for i in range(50)]  # Set index as timestamps
            writer.write_dataframe(df)
        df_read = to_dataframe(tsfile_path, table_name="test_table")
        df_read = df_read.sort_values(TIME_COLUMN).reset_index(drop=True)
        df_sorted = df.sort_index()
        df_sorted = convert_to_nullable_types(df_sorted.reset_index(drop=True))
        df_read = convert_to_nullable_types(df_read)
        time_series = pd.Series(df.sort_index().index.values, dtype="Int64")
        assert df_read.shape == (50, 3)
        assert df_read[TIME_COLUMN].equals(time_series)
        assert df_read["device"].equals(df_sorted["device"])
        assert df_read["value"].equals(df_sorted["value"])
    finally:
        if os.path.exists(tsfile_path):
            os.remove(tsfile_path)


def test_write_dataframe_case_insensitive():
    table = TableSchema(
        "test_table",
        [
            ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG),
            ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD),
        ],
    )
    tsfile_path = "test_write_dataframe_case.tsfile"
    try:
        if os.path.exists(tsfile_path):
            os.remove(tsfile_path)

        with TsFileTableWriter(tsfile_path, table) as writer:
            df = pd.DataFrame(
                {
                    "Time": [i for i in range(30)],  # Capital T
                    "Device": [f"device{i}" for i in range(30)],  # Capital D
                    "VALUE": [i * 3.0 for i in range(30)],  # All caps
                }
            )
            writer.write_dataframe(df)

        df_read = to_dataframe(tsfile_path, table_name="test_table")
        df_read = df_read.sort_values(TIME_COLUMN).reset_index(drop=True)
        df_sorted = convert_to_nullable_types(
            df.sort_values("Time").reset_index(drop=True)
        )
        df_read = convert_to_nullable_types(df_read)
        assert df_read.shape == (30, 3)
        assert df_read[TIME_COLUMN].equals(df_sorted["Time"])
        assert df_read["device"].equals(df_sorted["Device"])
        assert df_read["value"].equals(df_sorted["VALUE"])
    finally:
        if os.path.exists(tsfile_path):
            os.remove(tsfile_path)


def test_write_dataframe_column_not_in_schema():
    table = TableSchema(
        "test_table",
        [
            ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG),
            ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD),
        ],
    )
    tsfile_path = "test_write_dataframe_extra_col.tsfile"
    try:
        if os.path.exists(tsfile_path):
            os.remove(tsfile_path)

        with TsFileTableWriter(tsfile_path, table) as writer:
            df = pd.DataFrame(
                {
                    "time": [i for i in range(10)],
                    "device": [f"device{i}" for i in range(10)],
                    "value": [i * 1.0 for i in range(10)],
                    "extra_column": [i for i in range(10)],  # Not in schema
                }
            )
            with pytest.raises(ColumnNotExistError):
                writer.write_dataframe(df)
    finally:
        if os.path.exists(tsfile_path):
            os.remove(tsfile_path)


def test_write_dataframe_type_mismatch():
    table = TableSchema(
        "test_table", [ColumnSchema("value", TSDataType.STRING, ColumnCategory.FIELD)]
    )
    tsfile_path = "test_write_dataframe_type_mismatch.tsfile"
    try:
        if os.path.exists(tsfile_path):
            os.remove(tsfile_path)

        with TsFileTableWriter(tsfile_path, table) as writer:
            df = pd.DataFrame(
                {"time": [i for i in range(10)], "value": [i for i in range(10)]}
            )
            with pytest.raises(TypeMismatchError) as exc_info:
                writer.write_dataframe(df)
    finally:
        if os.path.exists(tsfile_path):
            os.remove(tsfile_path)


def test_write_dataframe_all_datatypes():
    table = TableSchema(
        "test_table",
        [
            ColumnSchema("bool_col", TSDataType.BOOLEAN, ColumnCategory.FIELD),
            ColumnSchema("int32_col", TSDataType.INT32, ColumnCategory.FIELD),
            ColumnSchema("int64_col", TSDataType.INT64, ColumnCategory.FIELD),
            ColumnSchema("float_col", TSDataType.FLOAT, ColumnCategory.FIELD),
            ColumnSchema("double_col", TSDataType.DOUBLE, ColumnCategory.FIELD),
            ColumnSchema("string_col", TSDataType.STRING, ColumnCategory.FIELD),
            ColumnSchema("blob_col", TSDataType.BLOB, ColumnCategory.FIELD),
            ColumnSchema("text_col", TSDataType.TEXT, ColumnCategory.FIELD),
            ColumnSchema("date_col", TSDataType.DATE, ColumnCategory.FIELD),
            ColumnSchema("timestamp_col", TSDataType.TIMESTAMP, ColumnCategory.FIELD),
        ],
    )
    tsfile_path = "test_write_dataframe_all_types.tsfile"
    try:
        if os.path.exists(tsfile_path):
            os.remove(tsfile_path)

        with TsFileTableWriter(tsfile_path, table) as writer:
            df = pd.DataFrame(
                {
                    "time": [i for i in range(50)],
                    "bool_col": [i % 2 == 0 for i in range(50)],
                    "int32_col": pd.Series([i for i in range(50)], dtype="int32"),
                    "int64_col": [i * 10 for i in range(50)],
                    "float_col": pd.Series(
                        [i * 1.5 for i in range(50)], dtype="float32"
                    ),
                    "double_col": [i * 2.5 for i in range(50)],
                    "string_col": [f"str{i}" for i in range(50)],
                    "blob_col": [f"blob{i}".encode("utf-8") for i in range(50)],
                    "text_col": [f"text{i}" for i in range(50)],
                    "date_col": [date(2025, i % 11 + 1, i % 20 + 1) for i in range(50)],
                    "timestamp_col": [i for i in range(50)],
                }
            )
            writer.write_dataframe(df)

        df_read = to_dataframe(tsfile_path, table_name="test_table")
        df_read = df_read.sort_values(TIME_COLUMN).reset_index(drop=True)
        df_sorted = convert_to_nullable_types(
            df.sort_values("time").reset_index(drop=True)
        )
        df_read = convert_to_nullable_types(df_read)
        assert df_read.shape == (50, 11)
        assert df_read["bool_col"].equals(df_sorted["bool_col"])
        assert df_read["int32_col"].equals(df_sorted["int32_col"])
        assert df_read["int64_col"].equals(df_sorted["int64_col"])
        assert np.allclose(df_read["float_col"], df_sorted["float_col"])
        assert np.allclose(df_read["double_col"], df_sorted["double_col"])
        assert df_read["string_col"].equals(df_sorted["string_col"])
        assert df_read["blob_col"].equals(df_sorted["blob_col"])
        assert df_read["text_col"].equals(df_sorted["text_col"])
        assert df_read["date_col"].equals(df_sorted["date_col"])
        assert df_read["timestamp_col"].equals(df_sorted["timestamp_col"])
        for i in range(50):
            assert df_read["blob_col"].iloc[i] == df_sorted["blob_col"].iloc[i]
    finally:
        if os.path.exists(tsfile_path):
            os.remove(tsfile_path)


def test_write_dataframe_schema_time_column():
    table = TableSchema(
        "test_table",
        [
            ColumnSchema("time", TSDataType.TIMESTAMP, ColumnCategory.TIME),
            ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG),
            ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD),
        ],
    )
    tsfile_path = "test_write_dataframe_schema_time.tsfile"
    try:
        if os.path.exists(tsfile_path):
            os.remove(tsfile_path)

        with TsFileTableWriter(tsfile_path, table) as writer:
            df = pd.DataFrame(
                {
                    "time": [i * 100 for i in range(50)],
                    "device": [f"device{i}" for i in range(50)],
                    "value": [i * 1.5 for i in range(50)],
                }
            )
            writer.write_dataframe(df)

        df_read = to_dataframe(tsfile_path, table_name="test_table")
        df_read = df_read.sort_values(TIME_COLUMN).reset_index(drop=True)
        df_sorted = convert_to_nullable_types(
            df.sort_values("time").reset_index(drop=True)
        )
        df_read = convert_to_nullable_types(df_read)
        assert df_read.shape == (50, 3)
        assert df_read[TIME_COLUMN].equals(df_sorted[TIME_COLUMN])
        assert df_read["device"].equals(df_sorted["device"])
        assert df_read["value"].equals(df_sorted["value"])
    finally:
        if os.path.exists(tsfile_path):
            os.remove(tsfile_path)


def test_write_dataframe_schema_time_and_dataframe_time():
    table = TableSchema(
        "test_table",
        [
            ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG),
            ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD),
        ],
    )
    tsfile_path = "test_write_dataframe_schema_and_df_time.tsfile"
    try:
        if os.path.exists(tsfile_path):
            os.remove(tsfile_path)

        with TsFileTableWriter(tsfile_path, table) as writer:
            df = pd.DataFrame(
                {
                    "Time": [i for i in range(30)],
                    "device": [f"dev{i}" for i in range(30)],
                    "value": [float(i) for i in range(30)],
                }
            )
            writer.write_dataframe(df)

        df_read = to_dataframe(tsfile_path, table_name="test_table")
        df_read = df_read.sort_values(TIME_COLUMN).reset_index(drop=True)
        df_sorted = convert_to_nullable_types(
            df.sort_values("Time").rename(columns=str.lower).reset_index(drop=True)
        )
        df_read = convert_to_nullable_types(df_read)
        assert df_read.shape == (30, 3)
        assert df_read["time"].equals(df_sorted["time"])
        assert df_read["device"].equals(df_sorted["device"])
        assert df_read["value"].equals(df_sorted["value"])
    finally:
        if os.path.exists(tsfile_path):
            os.remove(tsfile_path)


def test_write_dataframe_empty():
    table = TableSchema(
        "test_table", [ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD)]
    )
    tsfile_path = "test_write_dataframe_empty.tsfile"
    try:
        if os.path.exists(tsfile_path):
            os.remove(tsfile_path)

        with TsFileTableWriter(tsfile_path, table) as writer:
            df = pd.DataFrame({"time": [], "value": []})
            with pytest.raises(ValueError):
                writer.write_dataframe(df)

    finally:
        if os.path.exists(tsfile_path):
            os.remove(tsfile_path)


def test_validate_dataframe_none_column_name():
    df = pd.DataFrame([[1, 2]], columns=[None, "value"])
    with pytest.raises(ValueError, match="Column name cannot be None or empty"):
        validate_dataframe_for_tsfile(df)