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

import numpy as np
import pandas as pd

from emerge_res_integration_tool import config, constants
from emerge_res_integration_tool.network_model import BatteryFleet, NetworkModel
from emerge_res_integration_tool.timeseries_manager import TimeSeriesManager

try:
    from metrix_launcher.core.metrix_simulation import MetrixSimulation
except ImportError:  # pragma: no cover
    MetrixSimulation = None


_TEMPLATES_DIR = pathlib.Path(__file__).parent / "templates"

def _read_template(name: str) -> str:
    return (_TEMPLATES_DIR / name).read_text(encoding="utf-8")

def prepare_working_directory(folder_path: pathlib.Path, work_dir: pathlib.Path) -> None:
    """
    Create a working directory by copying <folder_path>/input to <work_dir>/input.

    The intent is to never modify user-provided input files in-place.
    """
    user_input_dir = folder_path / constants.INPUT_DATA_FOLDER
    if not user_input_dir.exists():
        raise FileNotFoundError(f"Input directory not found: {user_input_dir}")

    if work_dir.exists():
        shutil.rmtree(work_dir)

    work_input_dir = work_dir / constants.INPUT_DATA_FOLDER
    work_output_dir = folder_path / constants.OUTPUT_DATA_FOLDER
    work_input_dir.mkdir(parents=True, exist_ok=True)
    work_output_dir.mkdir(parents=True, exist_ok=True)

    shutil.copytree(
        user_input_dir,
        work_input_dir,
        dirs_exist_ok=True,
    )


def ensure_metrix_inputs(metrix_simulator) -> None:
    """
    Ensure required Metrix input files exist in the working directory.

    Creates missing:
      - contingencies.groovy (empty)
      - mapping.groovy (simple per-element mapping + battery mappings + PST fix)
      - conf.groovy (simple OPF configuration, enforcing computationType OPF)
    """
    # Contingencies
    contingencies_path = pathlib.Path(metrix_simulator.contingencies_file_path)
    if not contingencies_path.exists():
        config.logger.info("Creating empty contingencies file")
        if hasattr(metrix_simulator, "write_empty_contingency_file"):
            metrix_simulator.write_empty_contingency_file()
        else:
            contingencies_path.write_text("", encoding="utf-8")

    # Mapping
    mapping_path = pathlib.Path(metrix_simulator.mapping_file_path)
    if not mapping_path.exists():
        config.logger.info("Creating simple per-element mapping file")
        if hasattr(metrix_simulator, "write_per_element_simple_mapping_file"):
            metrix_simulator.write_per_element_simple_mapping_file()
        else:
            raise AttributeError(
                "MetrixSimulation has no write_per_element_simple_mapping_file method"
            )

    # Post-process mapping file: add battery mappings + fix PST mappings
    if mapping_path.exists():
        mapping_text = mapping_path.read_text(encoding="utf-8")

        # 1) Add BAT generator mappings if not already present
        if "g.id.contains('BAT')" not in mapping_text:
            mapping_text += "\n" + _read_template(constants.TEMPLATE_MAPPING_FILE)

        # 2) Replace mapToPhaseTapChangers NE_NO_1 block with direct assignment
        pattern = r"mapToPhaseTapChangers\s*\{\s*timeSeriesName\s*'NE_NO_1'\s*\}"
        replacement = (
            "network.getTwoWindingsTransformer('NE_NO_1').phaseTapChanger.tapPosition=17"
        )
        mapping_text = re.sub(
            pattern, replacement, mapping_text, flags=re.MULTILINE
        )

        mapping_path.write_text(mapping_text, encoding="utf-8")

    # Config (OPF)
    conf_path = pathlib.Path(metrix_simulator.conf_dsl_file_path)
    if not conf_path.exists():
        config.logger.info("Creating simple OPF configuration file")
        config.logger.warning("As the conf.groovy file was not provided, the line ratings are all set to 170 MW. This may cause issues for Metrix or result in no redispatching.")
        if hasattr(metrix_simulator, "write_simple_lf_conf_file"):
            metrix_simulator.write_simple_lf_conf_file() 
            # Replace everything because the function in metrix_simulator does not work well (temporary, until fix)
            text = _read_template(constants.TEMPLATE_COSTS_FILE)
            conf_path.write_text(text, encoding="utf-8")
        else:
            raise AttributeError(
                "MetrixSimulation has no write_simple_lf_conf_file method"
            )


def build_components_from_workdir(work_dir: pathlib.Path):
    """
    Build NetworkModel, TimeSeriesManager and MetrixSimulation from a work dir.
    """
    network_path = work_dir / constants.INPUT_DATA_FOLDER / constants.INPUT_NETWORK_FILE
    timeseries_path = work_dir / constants.INPUT_DATA_FOLDER / constants.INPUT_TIMESERIES

    network_model = NetworkModel(str(network_path))
    timeseries_manager = TimeSeriesManager(timeseries_path=str(timeseries_path))

    if MetrixSimulation is None:  # pragma: no cover
        raise ImportError("metrix_launcher is required to run Metrix simulations")

    metrix_simulator = MetrixSimulation(
        folder=work_dir,
        variant_count=str(timeseries_manager.horizon),
    )
    ensure_metrix_inputs(metrix_simulator)

    return network_model, timeseries_manager, metrix_simulator


def ensure_battery_timeseries_columns(
    work_input_dir: pathlib.Path,
    battery_fleet: BatteryFleet,
) -> None:
    """
    Ensure battery columns exist in timeseries CSV files.

    Also adds battery cost columns (bat_cost_*) expected by downstream logic.
    """
    if not battery_fleet.batteries:
        return

    main_ts_path = work_input_dir / constants.INPUT_TIMESERIES
    if not main_ts_path.exists():
        raise FileNotFoundError(f"Timeseries file not found: {main_ts_path}")

    def _update_timeseries_file(path: pathlib.Path):
        df = pd.read_csv(path, sep=";")

        for bat in battery_fleet.batteries:
            bat_id = bat.id
            col_target = f"{bat_id}_targetP"
            col_max = f"{bat_id}_maxP"
            col_min = f"{bat_id}_minP"
            
            bat_redisp_down, bat_redisp_up = battery_fleet.redispatch_costs
            bat_adeq_down, bat_adeq_up = battery_fleet.adequacy_costs

            if col_target not in df.columns:
                df[col_target] = 0.0
            if col_max not in df.columns:
                df[col_max] = bat.pmax
            if col_min not in df.columns:
                df[col_min] = -bat.pmax
            if constants.ADEQUACY_BATTERY_UP not in df.columns:
                df[constants.ADEQUACY_BATTERY_UP] = bat_adeq_up
            if constants.ADEQUACY_BATTERY_DOWN not in df.columns:
                df[constants.ADEQUACY_BATTERY_DOWN] = bat_adeq_down
            if constants.REDISPATCHING_BATTERY_DOWN not in df.columns:
                df[constants.REDISPATCHING_BATTERY_DOWN] = bat_redisp_down
            if constants.REDISPATCHING_BATTERY_UP not in df.columns:
                df[constants.REDISPATCHING_BATTERY_UP] = bat_redisp_up

        df.to_csv(path, sep=";", index=False)

    _update_timeseries_file(main_ts_path)

    initial_ts_path = main_ts_path.parent / f"initial_{main_ts_path.name}"
    if initial_ts_path.exists():
        _update_timeseries_file(initial_ts_path)
        
def find_substation_id(xml_file, full_name):
    tree = ET.parse(xml_file)
    root = tree.getroot()

    truncated = full_name[:12]

    for substation in root.findall(".//{*}substation"):
        for bus in substation.findall(".//{*}bus"):
            name = bus.get("name")
            if name and (
                full_name.startswith(name) or
                name.startswith(truncated)
            ):
                return substation.get("id")

    return None
