# 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 pathlib
import pypowsybl as pp
import shutil
import xml.etree.ElementTree as ET
import re

from emerge_res_integration_tool import config
from emerge_res_integration_tool.constants import ( 
    BAT_COST_ADEQUACY_DOWN,
    BAT_COST_ADEQUACY_UP,
    BAT_COST_REDISPATCHING_DOWN,
    BAT_COST_REDISPATCHING_UP,
    BAT_EMAX, 
    BAT_PMAX,
    BAT_TYPE,
    HYDRO_TYPE,
    OTHER_TYPE,
    SOLAR_TYPE,
    WIND_TYPE,
    )


class NetworkModel:
    """
    Wrapper around the network providing typed accessors for
    generators and helper methods to discover asset IDs (batteries, hydro,
    thermal, renewables...).
    """

    def __init__(self, network_filepath: str):
        self.network_filepath = network_filepath
        self.network = None

    def load(self):
        self.network = pp.network.load(self.network_filepath)
        config.logger.info(f"Loaded network from {self.network_filepath}")
        return self.network

    def save_temporary(self, filepath: str):
        if self.network is None:
            raise RuntimeError("Network not loaded before save_temporary")
        self.network.save(file=filepath)
        config.logger.info(f"Saved temporary network to {filepath}")

    @property
    def generators(self):
        if self.network is None:
            raise RuntimeError("Network not loaded before accessing generators")
        return self.network.get_generators()

    def get_battery_ids(self):
        generators = self.generators
        battery_ids = generators.index[generators.index.str.contains(BAT_TYPE)].tolist()
        config.logger.debug(f"Detected batteries on network: {battery_ids}")
        return battery_ids

    def get_hydro_ids(self):
        generators = self.generators
        hydro_ids = generators.index[
            generators["energy_source"].str.contains(HYDRO_TYPE)
        ].tolist()
        config.logger.debug(f"Detected hydro power plants on network: {hydro_ids}")
        return hydro_ids

    def get_solar_ids(self):
        """Get solar generator IDs from the network."""
        generators = self.generators
        solar_ids = generators.index[
            generators["energy_source"].str.contains(SOLAR_TYPE)
        ].tolist()
        config.logger.debug(f"Detected solar power plants on network: {solar_ids}")
        return solar_ids

    def get_wind_ids(self):
        """Get wind generator IDs from the network."""
        generators = self.generators
        wind_ids = generators.index[
            generators["energy_source"].str.contains(WIND_TYPE)
        ].tolist()
        config.logger.debug(f"Detected wind power plants on network: {wind_ids}")
        return wind_ids

    def get_thermal_ids(self):
        generators = self.generators
        thermal_ids = generators.index[
            generators["energy_source"].str.contains(OTHER_TYPE)
        ].tolist()
        config.logger.debug(f"Detected thermal power plants on network: {thermal_ids}")
        return thermal_ids
    
    def add_battery_to_network(
        self,
        substation_id: str,
        battery_id: str = "BAT",
        min_p: float = -BAT_PMAX,
        max_p: float = BAT_PMAX,
        target_p: float = 0.0,
        target_v: float = 406.45,
        backup_original: bool = True,
    ):
        """
        Add a battery to the network XIIDM file.
        
        Args:
            substation_id: ID of the substation where to add the battery
            battery_id: ID for the new battery generator
            min_p: Minimum power (MW), negative for charging
            max_p: Maximum power (MW), positive for discharging
            target_p: Target power setpoint (MW)
            target_v: Target voltage (kV)
            backup_original: If True, save original network as *_without_battery.xiidm
        
        Returns:
            Path to the new network file with battery
        """
        original_path = pathlib.Path(self.network_filepath)
        
        # Backup original if requested
        if backup_original:
            backup_path = original_path.parent / f"{original_path.stem}_without_battery{original_path.suffix}"
            if not backup_path.exists():
                shutil.copy2(original_path, backup_path)
        
        # Parse XML
        tree = ET.parse(self.network_filepath)
        root = tree.getroot()
        
        # Get namespace
        ns = self._get_iidm_namespace(root)
        
        # Find substation and voltage level
        substation, voltage_level = self._find_voltage_level(root, substation_id, ns)
        
        # Detect topology type
        topology_kind = voltage_level.get("topologyKind", "")
        
        if topology_kind == "BUS_BREAKER":
            self._add_battery_bus_breaker(
                voltage_level, ns, battery_id, min_p, max_p, target_p, target_v
            )
        elif topology_kind == "NODE_BREAKER":
            self._add_battery_node_breaker(
                voltage_level, ns, battery_id, min_p, max_p, target_p, target_v
            )
        else:
            raise ValueError(f"Unsupported topology kind: {topology_kind}")
        
        # Save modified network
        output_path = original_path.parent / f"{original_path.stem}_with_battery{original_path.suffix}"
        tree.write(output_path, encoding="unicode", xml_declaration=True)
        
        config.logger.info(f"Added battery '{battery_id}' to substation '{substation_id}'")
        config.logger.info(f"Saved network with battery to {output_path}")
        
        # Update internal path to point to new file
        self.network_filepath = str(output_path)
        
        return str(output_path)

    def add_batteries_to_network(
        self,
        substation_ids,
        backup_original: bool = True,
    ) -> str:
        """
        Add one battery per substation ID to the network XIIDM file.

        Batteries are created with unique IDs derived from the substation IDs.

        Args:
            substation_ids: Iterable of substation IDs where batteries are added.
            backup_original: If True, save original network as *_without_battery.xiidm
                before the first modification.

        Returns:
            Path to the final network file with all batteries.
        """
        substation_ids = list(substation_ids)
        final_path = None
        for idx, sub_id in enumerate(substation_ids):
            bat_id = "BAT" if len(substation_ids) == 1 else f"BAT_{sub_id}"
            final_path = self.add_battery_to_network(
                substation_id=sub_id,
                battery_id=bat_id,
                backup_original=backup_original and idx == 0,
            )
        if final_path is None:
            raise ValueError("No substation IDs provided for battery creation")
        return final_path
    
    def _get_iidm_namespace(self, root: ET.Element) -> str:
        """Extract IIDM namespace from root element."""
        raw = ET.tostring(root, encoding="unicode")
        for prefix, uri in re.findall(r'xmlns(?::(\w+))?="([^"]+)"', raw):
            if "iidm" in uri.lower() or "itesla" in uri.lower():
                return uri
        return ""  # No namespace
    
    def _find_voltage_level(self, root: ET.Element, substation_id: str, ns: str):
        """Find voltage level in specified substation."""
        tag = f"{{{ns}}}substation" if ns else "substation"
        vl_tag = f"{{{ns}}}voltageLevel" if ns else "voltageLevel"
        
        for substation in root.iter(tag):
            if substation.get("id") == substation_id:
                voltage_level = substation.find(vl_tag)
                if voltage_level is None:
                    raise ValueError(f"No voltageLevel found in substation {substation_id}")
                return substation, voltage_level
        
        raise ValueError(f"Substation {substation_id} not found in network")
    
    def _add_battery_bus_breaker(
        self, vl: ET.Element, ns: str, battery_id: str,
        min_p: float, max_p: float, target_p: float, target_v: float
    ):
        """Add battery to BUS_BREAKER topology."""
        topo_tag = f"{{{ns}}}busBreakerTopology" if ns else "busBreakerTopology"
        bus_tag = f"{{{ns}}}bus" if ns else "bus"
        gen_tag = f"{{{ns}}}generator" if ns else "generator"
        cap_tag = f"{{{ns}}}reactiveCapabilityCurve" if ns else "reactiveCapabilityCurve"
        pt_tag = f"{{{ns}}}point" if ns else "point"
        
        topo = vl.find(topo_tag)
        if topo is None:
            raise ValueError("busBreakerTopology not found")
        
        # Find first bus
        bus = topo.find(bus_tag)
        if bus is None:
            raise ValueError("No bus found in voltageLevel")
        bus_id = bus.get("id")
        
        # Create generator (battery)
        gen = ET.SubElement(vl, gen_tag)
        gen.set("id", battery_id)
        gen.set("energySource", "NUCLEAR")
        gen.set("minP", str(min_p))
        gen.set("maxP", str(max_p))
        gen.set("voltageRegulatorOn", "true")
        gen.set("targetP", str(target_p))
        gen.set("targetQ", "0.0")
        gen.set("targetV", str(target_v))
        gen.set("bus", bus_id)
        gen.set("connectableBus", bus_id)
        gen.set("p", "0.0")
        
        # Reactive capability curve
        cap = ET.SubElement(gen, cap_tag)
        for p_val in [min_p, max_p]:
            pt = ET.SubElement(cap, pt_tag)
            pt.set("p", str(p_val))
            pt.set("minQ", "-200")
            pt.set("maxQ", "200")
        
        config.logger.debug(f"Added battery {battery_id} to bus {bus_id}")
    
    def _add_battery_node_breaker(
        self, vl: ET.Element, ns: str, battery_id: str,
        min_p: float, max_p: float, target_p: float, target_v: float
    ):
        """Add battery to NODE_BREAKER topology with breaker and disconnectors."""
        topo_tag = f"{{{ns}}}nodeBreakerTopology" if ns else "nodeBreakerTopology"
        sw_tag = f"{{{ns}}}switch" if ns else "switch"
        gen_tag = f"{{{ns}}}generator" if ns else "generator"
        cap_tag = f"{{{ns}}}reactiveCapabilityCurve" if ns else "reactiveCapabilityCurve"
        pt_tag = f"{{{ns}}}point" if ns else "point"
        
        topo = vl.find(topo_tag)
        if topo is None:
            raise ValueError("nodeBreakerTopology not found")
        
        # Get next available nodes
        node_count = int(topo.get("nodeCount", "0"))
        gen_node = node_count
        mid_node = node_count + 1
        
        topo.set("nodeCount", str(node_count + 2))
        
        # Main busbar (node 0)
        main_bus_node = "0"
        
        # Disconnector (to busbar)
        sw_disc = ET.SubElement(topo, sw_tag)
        sw_disc.set("id", f"{vl.get('id')}_{battery_id}_SC_0")
        sw_disc.set("kind", "DISCONNECTOR")
        sw_disc.set("retained", "false")
        sw_disc.set("open", "false")
        sw_disc.set("node1", main_bus_node)
        sw_disc.set("node2", str(mid_node))
        
        # Breaker
        sw_breaker = ET.SubElement(topo, sw_tag)
        sw_breaker.set("id", f"{vl.get('id')}_{battery_id}_DJ")
        sw_breaker.set("kind", "BREAKER")
        sw_breaker.set("retained", "false")
        sw_breaker.set("open", "false")
        sw_breaker.set("node1", str(mid_node))
        sw_breaker.set("node2", str(gen_node))
        
        # Generator (battery)
        gen = ET.SubElement(vl, gen_tag)
        gen.set("id", battery_id)
        gen.set("energySource", "NUCLEAR")
        gen.set("minP", str(min_p))
        gen.set("maxP", str(max_p))
        gen.set("voltageRegulatorOn", "true")
        gen.set("targetP", str(target_p))
        gen.set("targetQ", "0.0")
        gen.set("targetV", str(target_v))
        gen.set("node", str(gen_node))
        gen.set("p", "0.0")
        
        # Reactive capability curve
        cap = ET.SubElement(gen, cap_tag)
        for p_val in [min_p, max_p]:
            pt = ET.SubElement(cap, pt_tag)
            pt.set("p", str(p_val))
            pt.set("minQ", "-200")
            pt.set("maxQ", "200")
        
        config.logger.debug(f"Added battery {battery_id} at nodes {gen_node}/{mid_node}")


class Battery:
    """Representation of a single battery."""

    def __init__(self, bat_id: str, emax: float, pmax: float):
        self.id = bat_id
        self.emax = float(emax)
        self.pmax = float(pmax)


class BatteryFleet:
    """Collection of batteries."""

    def __init__(self, batteries, redispatch_costs=None):
        self.batteries = list(batteries)
        self.redispatch_costs = redispatch_costs or (BAT_COST_REDISPATCHING_DOWN, BAT_COST_REDISPATCHING_UP)
        self.adequacy_costs = (BAT_COST_ADEQUACY_DOWN, BAT_COST_ADEQUACY_UP)

    @classmethod
    def from_ids_with_defaults(
        cls, battery_ids, default_emax=BAT_EMAX, default_pmax=BAT_PMAX, redispatch_costs=None
    ):
        batteries = [Battery(bid, default_emax, default_pmax) for bid in battery_ids]
        return cls(batteries, redispatch_costs=redispatch_costs)

    @property
    def ids(self):
        return [b.id for b in self.batteries]

    @property
    def params_dict(self):
        """
        Returns the parameter dictionary expected by the legacy
        TimeSeriesManager.update_from_results logic.
        """
        return {
            b.id: {
                "Emax": b.emax,
                "Pmax": b.pmax,
            }
            for b in self.batteries
        }
