import csv
import os
from typing import Any

import numpy as np
import pytest

import supervision as sv
from supervision.detection.tools.csv_sink import CSVSink
from tests.helpers import _create_detections


@pytest.mark.parametrize(
    (
        "detections",
        "custom_data",
        "second_detections",
        "second_custom_data",
        "file_name",
        "expected_result",
    ),
    [
        (
            _create_detections(
                xyxy=[[10, 20, 30, 40], [50, 60, 70, 80]],
                confidence=[0.7, 0.8],
                class_id=[0, 0],
                tracker_id=[0, 1],
                data={"class_name": ["person", "person"]},
            ),
            {"frame_number": 42},
            _create_detections(
                xyxy=[[15, 25, 35, 45], [55, 65, 75, 85]],
                confidence=[0.6, 0.9],
                class_id=[1, 1],
                tracker_id=[2, 3],
                data={"class_name": ["car", "car"]},
            ),
            {"frame_number": 43},
            "test_detections.csv",
            [
                [
                    "x_min",
                    "y_min",
                    "x_max",
                    "y_max",
                    "class_id",
                    "confidence",
                    "tracker_id",
                    "class_name",
                    "frame_number",
                ],
                ["10.0", "20.0", "30.0", "40.0", "0", "0.7", "0", "person", "42"],
                ["50.0", "60.0", "70.0", "80.0", "0", "0.8", "1", "person", "42"],
                ["15.0", "25.0", "35.0", "45.0", "1", "0.6", "2", "car", "43"],
                ["55.0", "65.0", "75.0", "85.0", "1", "0.9", "3", "car", "43"],
            ],
        ),  # multiple detections
        (
            _create_detections(
                xyxy=[[60, 70, 80, 90], [100, 110, 120, 130]],
                tracker_id=[4, 5],
                data={"class_name": ["bike", "dog"]},
            ),
            {"frame_number": 44},
            _create_detections(
                xyxy=[[65, 75, 85, 95], [105, 115, 125, 135]],
                confidence=[0.5, 0.4],
                data={"class_name": ["tree", "cat"]},
            ),
            {"frame_number": 45},
            "test_detections_missing_fields.csv",
            [
                [
                    "x_min",
                    "y_min",
                    "x_max",
                    "y_max",
                    "class_id",
                    "confidence",
                    "tracker_id",
                    "class_name",
                    "frame_number",
                ],
                ["60.0", "70.0", "80.0", "90.0", "", "", "4", "bike", "44"],
                ["100.0", "110.0", "120.0", "130.0", "", "", "5", "dog", "44"],
                ["65.0", "75.0", "85.0", "95.0", "", "0.5", "", "tree", "45"],
                ["105.0", "115.0", "125.0", "135.0", "", "0.4", "", "cat", "45"],
            ],
        ),  # missing fields
        (
            _create_detections(
                xyxy=[[10, 11, 12, 13]],
                confidence=[0.95],
                data={"class_name": "unknown", "is_detected": True, "score": 1},
            ),
            {"frame_number": 46},
            _create_detections(
                xyxy=[[14, 15, 16, 17]],
                data={"class_name": "artifact", "is_detected": False, "score": 0.85},
            ),
            {"frame_number": 47},
            "test_detections_varied_data.csv",
            [
                [
                    "x_min",
                    "y_min",
                    "x_max",
                    "y_max",
                    "class_id",
                    "confidence",
                    "tracker_id",
                    "class_name",
                    "frame_number",
                    "is_detected",
                    "score",
                ],
                [
                    "10.0",
                    "11.0",
                    "12.0",
                    "13.0",
                    "",
                    "0.95",
                    "",
                    "unknown",
                    "46",
                    "True",
                    "1",
                ],
                [
                    "14.0",
                    "15.0",
                    "16.0",
                    "17.0",
                    "",
                    "",
                    "",
                    "artifact",
                    "47",
                    "False",
                    "0.85",
                ],
            ],
        ),  # Inconsistent Data Types
        (
            _create_detections(
                xyxy=[[20, 21, 22, 23]],
            ),
            {
                "metadata": {"sensor_id": 101, "location": "north"},
                "tags": ["urgent", "review"],
            },
            _create_detections(
                xyxy=[[14, 15, 16, 17]],
            ),
            {
                "metadata": {"sensor_id": 104, "location": "west"},
                "tags": ["not-urgent", "done"],
            },
            "test_detections_complex_data.csv",
            [
                [
                    "x_min",
                    "y_min",
                    "x_max",
                    "y_max",
                    "class_id",
                    "confidence",
                    "tracker_id",
                    "metadata",
                    "tags",
                ],
                [
                    "20.0",
                    "21.0",
                    "22.0",
                    "23.0",
                    "",
                    "",
                    "",
                    "{'sensor_id': 101, 'location': 'north'}",
                    "['urgent', 'review']",
                ],
                [
                    "14.0",
                    "15.0",
                    "16.0",
                    "17.0",
                    "",
                    "",
                    "",
                    "{'sensor_id': 104, 'location': 'west'}",
                    "['not-urgent', 'done']",
                ],
            ],
        ),  # Complex Data
        (
            _create_detections(
                xyxy=[[10, 20, 30, 40], [50, 60, 70, 80]],
                confidence=[0.9, 0.8],
                class_id=[0, 1],
            ),
            {"area": np.array([400.0, 400.0]), "frame": 5},
            _create_detections(
                xyxy=[[15, 25, 35, 45]],
                confidence=[0.7],
                class_id=[2],
            ),
            {"area": np.array([100.0]), "frame": 6},
            "test_detections_mixed_custom_data.csv",
            [
                [
                    "x_min",
                    "y_min",
                    "x_max",
                    "y_max",
                    "class_id",
                    "confidence",
                    "tracker_id",
                    "area",
                    "frame",
                ],
                ["10.0", "20.0", "30.0", "40.0", "0", "0.9", "", "400.0", "5"],
                ["50.0", "60.0", "70.0", "80.0", "1", "0.8", "", "400.0", "5"],
                ["15.0", "25.0", "35.0", "45.0", "2", "0.7", "", "100.0", "6"],
            ],
        ),  # mixed custom_data: ndarray sliced per row, scalar broadcast to all rows
        (
            _create_detections(
                xyxy=[[10, 20, 30, 40], [50, 60, 70, 80]],
                confidence=[0.9, 0.8],
                class_id=[0, 1],
            ),
            {"ids": ["a", "b"], "tags": ("x", "y"), "frame": 7},
            _create_detections(
                xyxy=[[15, 25, 35, 45]],
                confidence=[0.7],
                class_id=[2],
            ),
            {"ids": ["c"], "tags": ("z",), "frame": 8},
            "test_detections_list_custom_data.csv",
            [
                [
                    "x_min",
                    "y_min",
                    "x_max",
                    "y_max",
                    "class_id",
                    "confidence",
                    "tracker_id",
                    "frame",
                    "ids",
                    "tags",
                ],
                ["10.0", "20.0", "30.0", "40.0", "0", "0.9", "", "7", "a", "x"],
                ["50.0", "60.0", "70.0", "80.0", "1", "0.8", "", "7", "b", "y"],
                ["15.0", "25.0", "35.0", "45.0", "2", "0.7", "", "8", "c", "z"],
            ],
        ),  # list/tuple custom_data matching detection count is sliced per row
        (
            sv.Detections(
                xyxy=np.array([[10, 20, 30, 40], [50, 60, 70, 80]]),
                data={"labels": ["person", "car"]},
            ),
            None,
            sv.Detections(
                xyxy=np.array([[15, 25, 35, 45]]),
                data={"labels": ["bus"]},
            ),
            None,
            "test_detections_plain_list_data.csv",
            [
                [
                    "x_min",
                    "y_min",
                    "x_max",
                    "y_max",
                    "class_id",
                    "confidence",
                    "tracker_id",
                    "labels",
                ],
                ["10", "20", "30", "40", "", "", "", "person"],
                ["50", "60", "70", "80", "", "", "", "car"],
                ["15", "25", "35", "45", "", "", "", "bus"],
            ],
        ),  # plain Python list in detections.data is sliced per row without custom_data
    ],
)
def test_csv_sink(
    detections: sv.Detections,
    custom_data: dict[str, Any] | None,
    second_detections: sv.Detections,
    second_custom_data: dict[str, Any] | None,
    file_name: str,
    expected_result: list[list[Any]],
) -> None:
    with sv.CSVSink(file_name) as sink:
        if custom_data is None:
            sink.append(detections)
        else:
            sink.append(detections, custom_data)
        if second_custom_data is None:
            sink.append(second_detections)
        else:
            sink.append(second_detections, second_custom_data)

    assert_csv_equal(file_name, expected_result)


@pytest.mark.parametrize(
    (
        "detections",
        "custom_data",
        "second_detections",
        "second_custom_data",
        "file_name",
        "expected_result",
    ),
    [
        (
            _create_detections(
                xyxy=[[10, 20, 30, 40], [50, 60, 70, 80]],
                confidence=[0.7, 0.8],
                class_id=[0, 0],
                tracker_id=[0, 1],
                data={"class_name": ["person", "person"]},
            ),
            {"frame_number": 42},
            _create_detections(
                xyxy=[[15, 25, 35, 45], [55, 65, 75, 85]],
                confidence=[0.6, 0.9],
                class_id=[1, 1],
                tracker_id=[2, 3],
                data={"class_name": ["car", "car"]},
            ),
            {"frame_number": 43},
            "test_detections.csv",
            [
                [
                    "x_min",
                    "y_min",
                    "x_max",
                    "y_max",
                    "class_id",
                    "confidence",
                    "tracker_id",
                    "class_name",
                    "frame_number",
                ],
                ["10.0", "20.0", "30.0", "40.0", "0", "0.7", "0", "person", "42"],
                ["50.0", "60.0", "70.0", "80.0", "0", "0.8", "1", "person", "42"],
                ["15.0", "25.0", "35.0", "45.0", "1", "0.6", "2", "car", "43"],
                ["55.0", "65.0", "75.0", "85.0", "1", "0.9", "3", "car", "43"],
            ],
        ),  # multiple detections
        (
            _create_detections(
                xyxy=[[60, 70, 80, 90], [100, 110, 120, 130]],
                tracker_id=[4, 5],
                data={"class_name": ["bike", "dog"]},
            ),
            {"frame_number": 44},
            _create_detections(
                xyxy=[[65, 75, 85, 95], [105, 115, 125, 135]],
                confidence=[0.5, 0.4],
                data={"class_name": ["tree", "cat"]},
            ),
            {"frame_number": 45},
            "test_detections_missing_fields.csv",
            [
                [
                    "x_min",
                    "y_min",
                    "x_max",
                    "y_max",
                    "class_id",
                    "confidence",
                    "tracker_id",
                    "class_name",
                    "frame_number",
                ],
                ["60.0", "70.0", "80.0", "90.0", "", "", "4", "bike", "44"],
                ["100.0", "110.0", "120.0", "130.0", "", "", "5", "dog", "44"],
                ["65.0", "75.0", "85.0", "95.0", "", "0.5", "", "tree", "45"],
                ["105.0", "115.0", "125.0", "135.0", "", "0.4", "", "cat", "45"],
            ],
        ),  # missing fields
        (
            _create_detections(
                xyxy=[[10, 11, 12, 13]],
                confidence=[0.95],
                data={"class_name": "unknown", "is_detected": True, "score": 1},
            ),
            {"frame_number": 46},
            _create_detections(
                xyxy=[[14, 15, 16, 17]],
                data={"class_name": "artifact", "is_detected": False, "score": 0.85},
            ),
            {"frame_number": 47},
            "test_detections_varied_data.csv",
            [
                [
                    "x_min",
                    "y_min",
                    "x_max",
                    "y_max",
                    "class_id",
                    "confidence",
                    "tracker_id",
                    "class_name",
                    "frame_number",
                    "is_detected",
                    "score",
                ],
                [
                    "10.0",
                    "11.0",
                    "12.0",
                    "13.0",
                    "",
                    "0.95",
                    "",
                    "unknown",
                    "46",
                    "True",
                    "1",
                ],
                [
                    "14.0",
                    "15.0",
                    "16.0",
                    "17.0",
                    "",
                    "",
                    "",
                    "artifact",
                    "47",
                    "False",
                    "0.85",
                ],
            ],
        ),  # Inconsistent Data Types
        (
            _create_detections(
                xyxy=[[20, 21, 22, 23]],
            ),
            {
                "metadata": {"sensor_id": 101, "location": "north"},
                "tags": ["urgent", "review"],
            },
            _create_detections(
                xyxy=[[14, 15, 16, 17]],
            ),
            {
                "metadata": {"sensor_id": 104, "location": "west"},
                "tags": ["not-urgent", "done"],
            },
            "test_detections_complex_data.csv",
            [
                [
                    "x_min",
                    "y_min",
                    "x_max",
                    "y_max",
                    "class_id",
                    "confidence",
                    "tracker_id",
                    "metadata",
                    "tags",
                ],
                [
                    "20.0",
                    "21.0",
                    "22.0",
                    "23.0",
                    "",
                    "",
                    "",
                    "{'sensor_id': 101, 'location': 'north'}",
                    "['urgent', 'review']",
                ],
                [
                    "14.0",
                    "15.0",
                    "16.0",
                    "17.0",
                    "",
                    "",
                    "",
                    "{'sensor_id': 104, 'location': 'west'}",
                    "['not-urgent', 'done']",
                ],
            ],
        ),  # Complex Data
    ],
)
def test_csv_sink_manual(
    detections: sv.Detections,
    custom_data: dict[str, Any],
    second_detections: sv.Detections,
    second_custom_data: dict[str, Any],
    file_name: str,
    expected_result: list[list[Any]],
) -> None:
    sink = sv.CSVSink(file_name)
    sink.open()
    sink.append(detections, custom_data)
    sink.append(second_detections, second_custom_data)
    sink.close()

    assert_csv_equal(file_name, expected_result)


def assert_csv_equal(file_name, expected_rows) -> None:
    with open(file_name, newline="") as file:
        reader = csv.reader(file)
        for i, row in enumerate(reader):
            assert [str(item) for item in expected_rows[i]] == row, (
                f"Row in CSV didn't match expected output: {row} != {expected_rows[i]}"
            )

    os.remove(file_name)


@pytest.mark.parametrize(
    ("value", "i", "n", "expected"),
    [
        (["x"], 0, 1, "x"),
        (["a", "b", "c"], 0, 2, ["a", "b", "c"]),
        ([42, "hello", None], 1, 3, "hello"),
        ([["a", "b"], ["c", "d"]], 0, 2, ["a", "b"]),
        ("ab", 0, 2, "ab"),
        (("z",), 0, 1, "z"),
    ],
)
def test_csv_sink_slice_value(value: Any, i: int, n: int, expected: Any) -> None:
    assert CSVSink._slice_value(value, i, n) == expected


@pytest.mark.parametrize(
    ("custom_data", "expected_value"),
    [
        pytest.param(
            {"embedding": np.array([1, 2, 3])},
            np.array([1, 2, 3]),
            id="custom_data_ndarray_length_mismatch",
        )
    ],
)
def test_csv_sink_broadcasts_ndarray_when_length_mismatches_detection_count(
    custom_data: dict[str, Any] | None, expected_value: np.ndarray
) -> None:
    """Mismatched ndarray data is broadcast instead of indexed per row."""
    detections = sv.Detections(
        xyxy=np.array([[0, 0, 10, 10], [20, 20, 30, 30]]),
    )

    rows = CSVSink.parse_detection_data(detections, custom_data=custom_data)

    assert len(rows) == 2
    np.testing.assert_array_equal(rows[0]["embedding"], expected_value)
    np.testing.assert_array_equal(rows[1]["embedding"], expected_value)


class TestCSVSinkLifecycle:
    """Tests for opening and reopening a CSV sink."""

    def test_reopening_starts_fresh_output_session(self, tmp_path: Any) -> None:
        """Reopening a sink writes a new header and field schema."""
        first_path = tmp_path / "first.csv"
        second_path = tmp_path / "second.csv"
        first_detections = sv.Detections(
            xyxy=np.array([[0, 0, 10, 10]]),
            data={"first_label": ["first"]},
        )
        second_detections = sv.Detections(
            xyxy=np.array([[20, 20, 30, 30]]),
            data={"second_label": ["second"]},
        )
        sink = sv.CSVSink(str(first_path))

        with sink:
            sink.append(first_detections)
        sink.file_name = str(second_path)
        with sink:
            sink.append(second_detections)

        with open(second_path, newline="") as file:
            reader = csv.DictReader(file)
            field_names = reader.fieldnames
            rows = list(reader)

        assert field_names == [
            "x_min",
            "y_min",
            "x_max",
            "y_max",
            "class_id",
            "confidence",
            "tracker_id",
            "second_label",
        ]
        assert rows == [
            {
                "x_min": "20",
                "y_min": "20",
                "x_max": "30",
                "y_max": "30",
                "class_id": "",
                "confidence": "",
                "tracker_id": "",
                "second_label": "second",
            }
        ]
