"""Unit tests for zone file I/O.

Tests:
- read_iwfm_zone_file
- write_iwfm_zone_file
- read_geojson_zones
- write_geojson_zones
- auto_detect_zone_file
- read_zone_file / write_zone_file
"""

from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import pytest

from pyiwfm.core.zones import Zone, ZoneDefinition
from pyiwfm.io.zones import (
    auto_detect_zone_file,
    read_geojson_zones,
    read_iwfm_zone_file,
    read_zone_file,
    write_geojson_zones,
    write_iwfm_zone_file,
    write_zone_file,
)

# =============================================================================
# Test IWFM Zone File Reading
# =============================================================================


class TestReadIwfmZoneFile:
    """Tests for read_iwfm_zone_file function."""

    def test_read_basic_file(self, tmp_path: Path) -> None:
        """Test reading a basic IWFM zone file."""
        zone_file = tmp_path / "zones.dat"
        zone_file.write_text(
            """C Zone Definition File
C For testing
1                           / ZExtent: horizontal
1  North Region
2  South Region
/
1    1
2    1
3    1
4    2
5    2
"""
        )

        zone_def = read_iwfm_zone_file(zone_file)

        assert zone_def.n_zones == 2
        assert zone_def.extent == "horizontal"
        assert zone_def.get_zone(1).name == "North Region"
        assert zone_def.get_zone(2).name == "South Region"
        assert zone_def.get_elements_in_zone(1) == [1, 2, 3]
        assert zone_def.get_elements_in_zone(2) == [4, 5]

    def test_read_vertical_extent(self, tmp_path: Path) -> None:
        """Test reading file with vertical extent."""
        zone_file = tmp_path / "zones.dat"
        zone_file.write_text(
            """0                           / ZExtent: vertical
1  Layer 1
/
1    1
2    1
"""
        )

        zone_def = read_iwfm_zone_file(zone_file)

        assert zone_def.extent == "vertical"

    def test_read_with_comments(self, tmp_path: Path) -> None:
        """Test reading file with various comment styles."""
        zone_file = tmp_path / "zones.dat"
        zone_file.write_text(
            """C This is a C comment
* This is an asterisk comment
1                           / ZExtent
C Zone definitions
1  Test Zone
C Separator
/
C Element assignments
1    1
"""
        )

        zone_def = read_iwfm_zone_file(zone_file)

        assert zone_def.n_zones == 1
        assert zone_def.get_zone(1).name == "Test Zone"

    def test_read_with_element_areas(self, tmp_path: Path) -> None:
        """Test reading with element areas for zone area calculation."""
        zone_file = tmp_path / "zones.dat"
        zone_file.write_text(
            """1
1  Zone A
/
1    1
2    1
3    1
"""
        )

        element_areas = {1: 100.0, 2: 150.0, 3: 200.0}
        zone_def = read_iwfm_zone_file(zone_file, element_areas=element_areas)

        assert zone_def.get_zone(1).area == pytest.approx(450.0)

    def test_read_zone_without_name(self, tmp_path: Path) -> None:
        """Test reading zone with only ID (no name)."""
        zone_file = tmp_path / "zones.dat"
        zone_file.write_text(
            """1
1
2
/
1    1
2    2
"""
        )

        zone_def = read_iwfm_zone_file(zone_file)

        assert zone_def.get_zone(1).name == "Zone 1"
        assert zone_def.get_zone(2).name == "Zone 2"


# =============================================================================
# Test IWFM Zone File Writing
# =============================================================================


class TestWriteIwfmZoneFile:
    """Tests for write_iwfm_zone_file function."""

    @pytest.fixture
    def sample_zone_def(self) -> ZoneDefinition:
        """Create sample zone definition for testing."""
        zones = {
            1: Zone(id=1, name="North", elements=[1, 2, 3], area=300.0),
            2: Zone(id=2, name="South", elements=[4, 5], area=200.0),
        }
        element_zones = np.array([1, 1, 1, 2, 2], dtype=np.int32)
        return ZoneDefinition(
            zones=zones,
            extent="horizontal",
            element_zones=element_zones,
            name="Test Zones",
            description="Test description",
        )

    def test_write_basic(self, sample_zone_def: ZoneDefinition, tmp_path: Path) -> None:
        """Test basic zone file writing."""
        output_file = tmp_path / "output_zones.dat"

        write_iwfm_zone_file(sample_zone_def, output_file)

        assert output_file.exists()

        content = output_file.read_text()
        assert "Zone Definition File" in content
        assert "North" in content
        assert "South" in content
        assert "ZExtent" in content

    def test_write_with_header(self, sample_zone_def: ZoneDefinition, tmp_path: Path) -> None:
        """Test writing with custom header comment."""
        output_file = tmp_path / "output_zones.dat"

        write_iwfm_zone_file(sample_zone_def, output_file, header_comment="Custom header")

        content = output_file.read_text()
        assert "Custom header" in content

    def test_roundtrip(self, sample_zone_def: ZoneDefinition, tmp_path: Path) -> None:
        """Test write and read back produces same data."""
        output_file = tmp_path / "roundtrip.dat"

        write_iwfm_zone_file(sample_zone_def, output_file)
        zone_def_back = read_iwfm_zone_file(output_file)

        assert zone_def_back.n_zones == sample_zone_def.n_zones
        assert zone_def_back.extent == sample_zone_def.extent
        assert zone_def_back.get_zone(1).name == "North"
        assert zone_def_back.get_zone(2).name == "South"
        assert zone_def_back.get_elements_in_zone(1) == [1, 2, 3]
        assert zone_def_back.get_elements_in_zone(2) == [4, 5]


# =============================================================================
# Test GeoJSON Zone File Reading
# =============================================================================


class TestReadGeojsonZones:
    """Tests for read_geojson_zones function."""

    def test_read_basic_geojson(self, tmp_path: Path) -> None:
        """Test reading basic GeoJSON zone file."""
        zone_file = tmp_path / "zones.geojson"
        geojson = {
            "type": "FeatureCollection",
            "features": [
                {
                    "type": "Feature",
                    "properties": {
                        "id": 1,
                        "name": "Zone A",
                        "element_id": [1, 2, 3],
                        "area": 500.0,
                    },
                    "geometry": None,
                },
                {
                    "type": "Feature",
                    "properties": {
                        "id": 2,
                        "name": "Zone B",
                        "element_id": [4, 5],
                        "area": 300.0,
                    },
                    "geometry": None,
                },
            ],
        }
        zone_file.write_text(json.dumps(geojson))

        zone_def = read_geojson_zones(zone_file)

        assert zone_def.n_zones == 2
        assert zone_def.get_zone(1).name == "Zone A"
        assert zone_def.get_zone(2).name == "Zone B"
        assert zone_def.get_zone(1).area == pytest.approx(500.0)
        assert zone_def.get_elements_in_zone(1) == [1, 2, 3]

    def test_read_with_custom_fields(self, tmp_path: Path) -> None:
        """Test reading GeoJSON with custom field names."""
        zone_file = tmp_path / "zones.geojson"
        geojson = {
            "type": "FeatureCollection",
            "features": [
                {
                    "type": "Feature",
                    "properties": {
                        "zone_id": 1,
                        "zone_name": "Custom Zone",
                        "elements": [1, 2],
                    },
                    "geometry": None,
                },
            ],
        }
        zone_file.write_text(json.dumps(geojson))

        zone_def = read_geojson_zones(
            zone_file,
            zone_id_field="zone_id",
            zone_name_field="zone_name",
            element_id_field="elements",
        )

        assert zone_def.n_zones == 1
        assert zone_def.get_zone(1).name == "Custom Zone"

    def test_read_missing_name(self, tmp_path: Path) -> None:
        """Test reading GeoJSON without zone name."""
        zone_file = tmp_path / "zones.geojson"
        geojson = {
            "type": "FeatureCollection",
            "features": [
                {
                    "type": "Feature",
                    "properties": {
                        "id": 1,
                        "element_id": [1, 2],
                    },
                    "geometry": None,
                },
            ],
        }
        zone_file.write_text(json.dumps(geojson))

        zone_def = read_geojson_zones(zone_file)

        assert zone_def.get_zone(1).name == "Zone 1"

    def test_read_empty_elements(self, tmp_path: Path) -> None:
        """Test reading GeoJSON with empty element list."""
        zone_file = tmp_path / "zones.geojson"
        geojson = {
            "type": "FeatureCollection",
            "features": [
                {
                    "type": "Feature",
                    "properties": {
                        "id": 1,
                        "name": "Empty Zone",
                        "element_id": [],
                    },
                    "geometry": None,
                },
            ],
        }
        zone_file.write_text(json.dumps(geojson))

        zone_def = read_geojson_zones(zone_file)

        assert zone_def.get_zone(1).elements == []


# =============================================================================
# Test GeoJSON Zone File Writing
# =============================================================================


class TestWriteGeojsonZones:
    """Tests for write_geojson_zones function."""

    @pytest.fixture
    def sample_zone_def(self) -> ZoneDefinition:
        """Create sample zone definition for testing."""
        zones = {
            1: Zone(id=1, name="North", elements=[1, 2, 3], area=300.0),
            2: Zone(id=2, name="South", elements=[4, 5], area=200.0),
        }
        element_zones = np.array([1, 1, 1, 2, 2], dtype=np.int32)
        return ZoneDefinition(
            zones=zones,
            extent="horizontal",
            element_zones=element_zones,
            name="Test Zones",
        )

    def test_write_basic(self, sample_zone_def: ZoneDefinition, tmp_path: Path) -> None:
        """Test basic GeoJSON writing."""
        output_file = tmp_path / "output.geojson"

        write_geojson_zones(sample_zone_def, output_file, include_geometry=False)

        assert output_file.exists()

        with open(output_file) as f:
            geojson = json.load(f)

        assert geojson["type"] == "FeatureCollection"
        assert len(geojson["features"]) == 2

        # Check feature properties
        feature1 = next(f for f in geojson["features"] if f["properties"]["id"] == 1)
        assert feature1["properties"]["name"] == "North"
        assert feature1["properties"]["element_id"] == [1, 2, 3]
        assert feature1["properties"]["area"] == 300.0

    def test_roundtrip(self, sample_zone_def: ZoneDefinition, tmp_path: Path) -> None:
        """Test write and read back produces same data."""
        output_file = tmp_path / "roundtrip.geojson"

        write_geojson_zones(sample_zone_def, output_file, include_geometry=False)
        zone_def_back = read_geojson_zones(output_file)

        assert zone_def_back.n_zones == sample_zone_def.n_zones
        assert zone_def_back.get_zone(1).name == "North"
        assert zone_def_back.get_zone(2).name == "South"
        assert zone_def_back.get_elements_in_zone(1) == [1, 2, 3]


# =============================================================================
# Test Auto-Detection
# =============================================================================


class TestAutoDetectZoneFile:
    """Tests for auto_detect_zone_file function."""

    def test_detect_geojson_by_extension(self, tmp_path: Path) -> None:
        """Test detection by .geojson extension."""
        zone_file = tmp_path / "zones.geojson"
        zone_file.write_text("{}")

        assert auto_detect_zone_file(zone_file) == "geojson"

    def test_detect_json_by_extension(self, tmp_path: Path) -> None:
        """Test detection by .json extension."""
        zone_file = tmp_path / "zones.json"
        zone_file.write_text("{}")

        assert auto_detect_zone_file(zone_file) == "geojson"

    def test_detect_iwfm_by_dat_extension(self, tmp_path: Path) -> None:
        """Test detection by .dat extension."""
        zone_file = tmp_path / "zones.dat"
        zone_file.write_text("1\n1  Zone\n/\n")

        assert auto_detect_zone_file(zone_file) == "iwfm"

    def test_detect_iwfm_by_txt_extension(self, tmp_path: Path) -> None:
        """Test detection by .txt extension."""
        zone_file = tmp_path / "zones.txt"
        zone_file.write_text("1\n1  Zone\n/\n")

        assert auto_detect_zone_file(zone_file) == "iwfm"

    def test_detect_geojson_by_content(self, tmp_path: Path) -> None:
        """Test detection by file content (starts with {)."""
        zone_file = tmp_path / "zones.zon"  # Unknown extension
        zone_file.write_text('{"type": "FeatureCollection"}')

        assert auto_detect_zone_file(zone_file) == "geojson"

    def test_detect_iwfm_by_content(self, tmp_path: Path) -> None:
        """Test detection by file content (IWFM format)."""
        zone_file = tmp_path / "zones.zon"  # Unknown extension
        zone_file.write_text("C Comment\n1\n")

        assert auto_detect_zone_file(zone_file) == "iwfm"


# =============================================================================
# Test Generic Read/Write Functions
# =============================================================================


class TestReadZoneFile:
    """Tests for read_zone_file function."""

    def test_read_iwfm_file(self, tmp_path: Path) -> None:
        """Test reading IWFM format by extension."""
        zone_file = tmp_path / "zones.dat"
        zone_file.write_text(
            """1
1  Test Zone
/
1    1
"""
        )

        zone_def = read_zone_file(zone_file)

        assert zone_def.n_zones == 1

    def test_read_geojson_file(self, tmp_path: Path) -> None:
        """Test reading GeoJSON format by extension."""
        zone_file = tmp_path / "zones.geojson"
        geojson = {
            "type": "FeatureCollection",
            "features": [
                {
                    "type": "Feature",
                    "properties": {"id": 1, "name": "Test", "element_id": [1]},
                    "geometry": None,
                }
            ],
        }
        zone_file.write_text(json.dumps(geojson))

        zone_def = read_zone_file(zone_file)

        assert zone_def.n_zones == 1


class TestWriteZoneFile:
    """Tests for write_zone_file function."""

    @pytest.fixture
    def sample_zone_def(self) -> ZoneDefinition:
        """Create sample zone definition."""
        zones = {1: Zone(id=1, name="Test", elements=[1, 2], area=100.0)}
        element_zones = np.array([1, 1], dtype=np.int32)
        return ZoneDefinition(zones=zones, element_zones=element_zones)

    def test_write_iwfm_by_extension(self, sample_zone_def: ZoneDefinition, tmp_path: Path) -> None:
        """Test writing IWFM format by extension."""
        output_file = tmp_path / "output.dat"

        write_zone_file(sample_zone_def, output_file)

        assert output_file.exists()
        content = output_file.read_text()
        assert "Zone Definition File" in content

    def test_write_geojson_by_extension(
        self, sample_zone_def: ZoneDefinition, tmp_path: Path
    ) -> None:
        """Test writing GeoJSON format by extension."""
        output_file = tmp_path / "output.geojson"

        write_zone_file(sample_zone_def, output_file)

        assert output_file.exists()
        with open(output_file) as f:
            geojson = json.load(f)
        assert geojson["type"] == "FeatureCollection"

    def test_write_txt_as_iwfm(self, sample_zone_def: ZoneDefinition, tmp_path: Path) -> None:
        """Test .txt extension writes IWFM format."""
        output_file = tmp_path / "output.txt"

        write_zone_file(sample_zone_def, output_file)

        content = output_file.read_text()
        assert "Zone Definition File" in content  # IWFM format indicator


# =============================================================================
# Additional coverage tests
# =============================================================================


class TestReadIwfmZoneFileEdgeCases:
    """Edge case tests for read_iwfm_zone_file."""

    def test_empty_lines_ignored(self, tmp_path: Path) -> None:
        """Test that empty lines in zone file are ignored."""
        zone_file = tmp_path / "zones.dat"
        zone_file.write_text(
            """1

1  Zone A

/

1    1

"""
        )
        zone_def = read_iwfm_zone_file(zone_file)
        assert zone_def.n_zones == 1

    def test_in_extension(self, tmp_path: Path) -> None:
        """Test auto-detect with .in extension."""
        zone_file = tmp_path / "zones.in"
        zone_file.write_text("1\n1 Z\n/\n1 1\n")
        assert auto_detect_zone_file(zone_file) == "iwfm"

    def test_invalid_zone_id_skipped(self, tmp_path: Path) -> None:
        """Test that non-integer zone IDs are skipped."""
        zone_file = tmp_path / "zones.dat"
        zone_file.write_text(
            """1
1  Zone A
abc  Invalid
/
1    1
"""
        )
        zone_def = read_iwfm_zone_file(zone_file)
        assert zone_def.n_zones == 1

    def test_invalid_element_assignment_skipped(self, tmp_path: Path) -> None:
        """Test invalid element assignments are skipped."""
        zone_file = tmp_path / "zones.dat"
        zone_file.write_text(
            """1
1  Zone A
/
1    1
abc  xyz
2    1
"""
        )
        zone_def = read_iwfm_zone_file(zone_file)
        assert zone_def.get_elements_in_zone(1) == [1, 2]

    def test_single_column_element_line_skipped(self, tmp_path: Path) -> None:
        """Test element lines with fewer than 2 columns are skipped."""
        zone_file = tmp_path / "zones.dat"
        zone_file.write_text(
            """1
1  Zone A
/
1    1
3
2    1
"""
        )
        zone_def = read_iwfm_zone_file(zone_file)
        assert zone_def.get_elements_in_zone(1) == [1, 2]


class TestReadGeojsonZonesEdgeCases:
    """Edge case tests for read_geojson_zones."""

    def test_missing_zone_id_skipped(self, tmp_path: Path) -> None:
        """Test features without zone ID are skipped."""
        zone_file = tmp_path / "zones.geojson"
        geojson = {
            "type": "FeatureCollection",
            "features": [
                {
                    "type": "Feature",
                    "properties": {"name": "No ID Zone", "element_id": [1]},
                    "geometry": None,
                },
                {
                    "type": "Feature",
                    "properties": {"id": 1, "name": "Valid", "element_id": [2]},
                    "geometry": None,
                },
            ],
        }
        zone_file.write_text(json.dumps(geojson))

        zone_def = read_geojson_zones(zone_file)
        assert zone_def.n_zones == 1

    def test_non_list_elements(self, tmp_path: Path) -> None:
        """Test that non-list element_id results in empty elements."""
        zone_file = tmp_path / "zones.geojson"
        geojson = {
            "type": "FeatureCollection",
            "features": [
                {
                    "type": "Feature",
                    "properties": {"id": 1, "name": "Test", "element_id": "not_a_list"},
                    "geometry": None,
                },
            ],
        }
        zone_file.write_text(json.dumps(geojson))

        zone_def = read_geojson_zones(zone_file)
        assert zone_def.get_zone(1).elements == []

    def test_empty_features_list(self, tmp_path: Path) -> None:
        """Test reading GeoJSON with no features."""
        zone_file = tmp_path / "zones.geojson"
        geojson = {"type": "FeatureCollection", "features": []}
        zone_file.write_text(json.dumps(geojson))

        zone_def = read_geojson_zones(zone_file)
        assert zone_def.n_zones == 0


class TestWriteGeojsonZonesWithGrid:
    """Tests for write_geojson_zones with grid geometry."""

    def test_write_with_convex_hull(self, tmp_path: Path) -> None:
        """Test writing GeoJSON with grid geometry (convex hull)."""
        from pyiwfm.core.mesh import AppGrid, Element, Node

        # Arrange elements so centroids are NOT collinear
        #   elem 1 centroid: (50, 50)
        #   elem 2 centroid: (150, 50)
        #   elem 3 centroid: (100, 200)
        nodes = {
            1: Node(id=1, x=0.0, y=0.0),
            2: Node(id=2, x=100.0, y=0.0),
            3: Node(id=3, x=100.0, y=100.0),
            4: Node(id=4, x=0.0, y=100.0),
            5: Node(id=5, x=200.0, y=0.0),
            6: Node(id=6, x=200.0, y=100.0),
            7: Node(id=7, x=50.0, y=200.0),
            8: Node(id=8, x=150.0, y=200.0),
            9: Node(id=9, x=150.0, y=300.0),
            10: Node(id=10, x=50.0, y=300.0),
        }
        elements = {
            1: Element(id=1, vertices=(1, 2, 3, 4), subregion=1, area=10000.0),
            2: Element(id=2, vertices=(2, 5, 6, 3), subregion=1, area=10000.0),
            3: Element(id=3, vertices=(7, 8, 9, 10), subregion=1, area=10000.0),
        }
        grid = AppGrid(nodes=nodes, elements=elements)

        zones = {1: Zone(id=1, name="Region", elements=[1, 2, 3], area=30000.0)}
        element_zones = np.array([1, 1, 1], dtype=np.int32)
        zone_def = ZoneDefinition(zones=zones, element_zones=element_zones)

        output_file = tmp_path / "with_geom.geojson"
        write_geojson_zones(zone_def, output_file, grid=grid, include_geometry=True)

        assert output_file.exists()
        with open(output_file) as f:
            geojson = json.load(f)

        feature = geojson["features"][0]
        assert feature["geometry"] is not None
        assert feature["geometry"]["type"] == "Polygon"

    def test_write_single_element_zone(self, tmp_path: Path) -> None:
        """Test writing zone with single element (point geometry)."""
        from pyiwfm.core.mesh import AppGrid, Element, Node

        nodes = {
            1: Node(id=1, x=0.0, y=0.0),
            2: Node(id=2, x=100.0, y=0.0),
            3: Node(id=3, x=50.0, y=100.0),
        }
        elements = {1: Element(id=1, vertices=(1, 2, 3), subregion=1, area=5000.0)}
        grid = AppGrid(nodes=nodes, elements=elements)

        zones = {1: Zone(id=1, name="Single", elements=[1], area=5000.0)}
        element_zones = np.array([1], dtype=np.int32)
        zone_def = ZoneDefinition(zones=zones, element_zones=element_zones)

        output_file = tmp_path / "single_elem.geojson"
        write_geojson_zones(zone_def, output_file, grid=grid, include_geometry=True)

        with open(output_file) as f:
            geojson = json.load(f)

        feature = geojson["features"][0]
        assert feature["geometry"] is not None
        assert feature["geometry"]["type"] == "Point"

    def test_write_no_geometry(self, tmp_path: Path) -> None:
        """Test writing GeoJSON without geometry."""
        zones = {1: Zone(id=1, name="Test", elements=[1], area=100.0)}
        zone_def = ZoneDefinition(zones=zones)

        output_file = tmp_path / "no_geom.geojson"
        write_geojson_zones(zone_def, output_file, include_geometry=False)

        with open(output_file) as f:
            geojson = json.load(f)

        assert geojson["features"][0]["geometry"] is None


class TestAutoDetectEdgeCases:
    """Edge case tests for auto_detect_zone_file."""

    def test_unknown_extension_json_content(self, tmp_path: Path) -> None:
        """Test unknown extension with JSON content."""
        f = tmp_path / "zones.xyz"
        f.write_text('{"type": "FeatureCollection"}')
        assert auto_detect_zone_file(f) == "geojson"

    def test_unknown_extension_text_content(self, tmp_path: Path) -> None:
        """Test unknown extension with text content."""
        f = tmp_path / "zones.xyz"
        f.write_text("C comment\n1\n")
        assert auto_detect_zone_file(f) == "iwfm"

    def test_read_zone_file_unknown_raises(self, tmp_path: Path) -> None:
        """Test read_zone_file raises for truly unknown format."""
        # Create a file that auto_detect returns "unknown" for
        # This requires a file that cannot be read
        f = tmp_path / "zones.bin"
        f.write_bytes(b"\x00\x01\x02\x03")

        # auto_detect may return "iwfm" for binary content, which is fine
        # The key is testing the ValueError path
        fmt = auto_detect_zone_file(f)
        # If format is detected, the read might fail differently
        if fmt == "unknown":
            with pytest.raises(ValueError, match="Cannot determine"):
                read_zone_file(f)


class TestWriteIwfmZoneFileEdgeCases:
    """Edge case tests for write_iwfm_zone_file."""

    def test_write_vertical_extent(self, tmp_path: Path) -> None:
        """Test writing vertical extent zone file."""
        zones = {1: Zone(id=1, name="Layer1", elements=[1, 2])}
        element_zones = np.array([1, 1], dtype=np.int32)
        zone_def = ZoneDefinition(
            zones=zones,
            extent="vertical",
            element_zones=element_zones,
            name="Vertical Zones",
            description="Test vertical",
        )

        output = tmp_path / "vertical.dat"
        write_iwfm_zone_file(zone_def, output)

        content = output.read_text()
        # Verify content includes zone data
        assert "Vertical" in content or len(content.splitlines()) > 1

    def test_write_no_element_zones(self, tmp_path: Path) -> None:
        """Test writing zone def with no element_zones array."""
        zones = {1: Zone(id=1, name="Empty")}
        zone_def = ZoneDefinition(zones=zones)

        output = tmp_path / "empty.dat"
        write_iwfm_zone_file(zone_def, output)

        assert output.exists()
