# Copyright (c) 2026, Artelys (https://www.artelys.com)
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
# SPDX-License-Identifier: MPL-2.0
#
#
# @author Léo Rossignol <leo.rossignol at artelys.com>
# @author Alexis Godefroy <alexis.godefroy at artelys.com>
#
import shutil
import unittest
import xml.etree.ElementTree as ET

from emerge_res_integration_tool import config_test, constants, utils
from emerge_res_integration_tool.network_model import NetworkModel


class TestAddBatteryToNetwork(unittest.TestCase):

    def setUp(self):
        """Work on a fresh copy of the network file for each test."""
        self.original_path = (
            config_test.TEST_DATA_DIR
            / constants.OPTIM_MOZAMBIQUE_FOLDER
            / constants.INPUT_DATA_FOLDER
            / constants.INPUT_NETWORK_FILE
        )
        self.test_network_path = self.original_path.parent / "test_network_tmp.xiidm"
        shutil.copy2(self.original_path, self.test_network_path)

        self.model = NetworkModel(network_filepath=str(self.test_network_path))

    def tearDown(self):
        """Clean up all files generated during the test."""
        for f in self.test_network_path.parent.glob("test_network_tmp*"):
            f.unlink(missing_ok=True)

    def test_add_single_battery_creates_output_file(self):
        """add_battery_to_network should produce a valid *_with_battery.xiidm file."""
        output_path = self.model.add_battery_to_network(substation_id="S14")

        self.assertTrue(
            ET.parse(output_path),
            "Output network file is not valid XML",
        )

    def test_add_single_battery_present_in_xml(self):
        """The battery generator node should appear in the output XML."""
        output_path = self.model.add_battery_to_network(
            substation_id="S14", battery_id="BAT_TEST"
        )

        root = ET.parse(output_path).getroot()
        generators = [el for el in root.iter() if el.get("id") == "BAT_TEST"]
        self.assertEqual(
            len(generators), 1, "Expected exactly one BAT_TEST generator in XML"
        )

    def test_add_batteries_multiple_substations(self):
        """add_batteries_to_network should add one battery per substation."""
        substation_ids = ["S14", "S4"]
        self.model.add_batteries_to_network(substation_ids=substation_ids)

        root = ET.parse(self.model.network_filepath).getroot()
        for sub_id in substation_ids:
            bat_id = f"BAT_{sub_id}"
            matches = [el for el in root.iter() if el.get("id") == bat_id]
            self.assertEqual(len(matches), 1, f"Battery {bat_id} not found in XML")


class TestFindSubstationId(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.network_path = (
            config_test.TEST_DATA_DIR
            / constants.OPTIM_MOZAMBIQUE_FOLDER
            / constants.INPUT_DATA_FOLDER
            / constants.INPUT_NETWORK_FILE
        )

    def test_find_substation_matola(self):
        """'79 Matola_220' should resolve to substation S14."""
        result = utils.find_substation_id(self.network_path, "79 Matola_220")
        self.assertEqual(result, "S14")

    def test_find_substation_unknown_bus(self):
        """An unknown bus name should return None."""
        result = utils.find_substation_id(self.network_path, "99 Unknown_999")
        self.assertIsNone(result)
