"""Parse SimStadt output files into typed result objects.

Supports HeatDemand, Photovoltaic, and GreenWater workflows.
Each result type exposes: .dataframe, .kpis, .diagrams, .csv_path,
.to_json() / .from_json(), and .zip_content().
"""

import io
import json
import logging
import re
import zipfile
from abc import ABC, abstractmethod
from dataclasses import asdict, dataclass, field
from enum import Enum
from pathlib import Path
from typing import ClassVar, Dict, Sequence, Type
from xml.etree import ElementTree as et

import matplotlib.pyplot as plt
import pandas as pd

logger = logging.getLogger(__name__)


@dataclass
class KPI:
    """Metric with name, value, unit and display precision."""

    name: str
    value: float
    unit: str | None = None
    precision: int = 2
    NAME_AND_UNIT = re.compile(r"^(.*?) \[(.*?)\]$")

    def __post_init__(self):
        """Move units embedded in column names (e.g. 'Area [m²]') into the unit field."""
        if match := self.NAME_AND_UNIT.match(self.name):
            self.name, self.unit = match.groups()

    def to_dict(self) -> dict:
        return asdict(self)

    @property
    def rounded_value(self) -> float:
        return round(self.value, self.precision)

    def __str__(self):
        s = f"{self.name:35s} : {self.value:.{self.precision}f}"
        if self.unit:
            s = f"{s} {self.unit}"
        return s

    def __repr__(self):
        return str(self)

    @classmethod
    def from_dict(cls, data: dict) -> "KPI":
        return cls(**data)


class Providers(str, Enum):
    HEAT_DEMAND = "de.hftstuttgart.simstadtworkflows.energy.HeatDemandAnalysisWorkflowProvider"
    HEAT_DEMAND_WITH_REFURBISHMENT = (
        "de.hftstuttgart.simstadtworkflows.energy.HeatDemandAnalysisWithRefurbishmentStrategyWorkflowProvider"
    )
    HEAT_DEMAND_WITH_HISTORIC_REFURBISHMENT = (
        "de.hftstuttgart.simstadtworkflows.energy.HeatDemandAnalysisWithHistoricAndFutureRefurbishmentWorkflowProvider"
    )
    HEAT_DEMAND_WITH_SHADOW = (
        "de.hftstuttgart.simstadtworkflows.shadow.HeatDemandCalculationWithShadowProcessingProvider"
    )
    PHOTOVOLTAIC = "de.hftstuttgart.simstadtworkflows.energy.PhotovoltaicPotentialAnalysisWorkflowProvider"
    PHOTOVOLTAIC_WITH_SHADOW = "de.hftstuttgart.simstadtworkflows.shadow.PVPotentialWithShadowProcessingProvider"
    PHOTOVOLTAIC_FINANCE = (
        "de.hftstuttgart.simstadtworkflows.economics.PhotovoltaicPotentialFinancialAnalysisWorkflowProvider"
    )
    GREEN_WATER = "de.hftstuttgart.simstadtworkflows.greenwater.GreenWaterWorkflowProvider"
    LOAD_PROFILE = "de.hftstuttgart.simstadtworkflows.energy.LoadProfileProvider"
    SOLAR_POTENTIAL = "de.hftstuttgart.simstadtworkflows.energy.SolarPotentialAnalysisWorkflowProvider"


def detect_decimal(csv_path: Path, line_number: int, check: str = "") -> str:
    """Detect whether '.' or ',' is used as decimal separator on the given line."""
    with open(csv_path, encoding="utf-8") as csv:
        for _ in range(line_number):
            next(csv)
        line = next(csv)
        if check and (check not in line):
            raise ValueError(f"Couldn't find '{check}' in '{line}'.")
        possible_delimiters = [".", ","]
        delimiter = next(d for d in possible_delimiters if d in line)
        logger.debug("Decimal for %s is '%s'", csv_path.name, delimiter)
        return delimiter


def get_workflow_provider(output_files: list[Path]) -> str:
    """Determine the SimStadt workflow provider class name from a params.xml file."""
    if len(output_files) == 0:
        raise ValueError("No output files were provided!")
    first_output_file = output_files[0]
    workflow_path = next(d for d in first_output_file.parents if d.suffix == ".flow")
    params_xml = workflow_path / "params.xml"

    tree = et.parse(params_xml)
    root = tree.getroot()
    element = root.find(".//void[@property='workflowProvider']/object")
    assert element is not None, "workflowProvider element not found"
    return element.get("class", "UnknownWorkflowProvider")


@dataclass
class SimStadtResults(ABC):
    """Abstract base class for SimStadt workflow results.

    Subclasses register themselves via __init_subclass__ so that
    create_simstadt_results() can instantiate the correct type automatically.
    """

    description: str
    output_files: list[Path]

    supported_providers: ClassVar[list[str]]
    csv_identifier: ClassVar[str]
    _registry: ClassVar[Dict[str, Type["SimStadtResults"]]] = {}

    _df: pd.DataFrame | None = field(default=None, init=False, repr=False)

    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        for provider in cls.supported_providers:
            SimStadtResults._registry[provider] = cls

    @abstractmethod
    def _parse_results(self) -> pd.DataFrame:
        pass

    @property
    def dataframe(self) -> pd.DataFrame:
        if self._df is None:
            self._df = self._parse_results()
        return self._df

    @property
    def more_info(self) -> dict[str, list[KPI]]:
        return {}

    @property
    @abstractmethod
    def kpis(self) -> list[KPI]:
        pass

    def kpi(self, name: str) -> KPI:
        for kpi in self.kpis:
            if kpi.name == name:
                return kpi
        raise ValueError(f"No '{name}' KPI found for {self}")

    def _prepare_kpis(self, sums: Sequence, averages: Sequence) -> list[KPI]:
        kpis = []
        default_precision = 0
        for item in sums:
            name, unit, precision = (item + (default_precision,))[:3]
            kpis.append(KPI(name, self.dataframe[name].sum(), unit, precision))
        for item in averages:
            name, unit, precision = (item + (default_precision,))[:3]
            kpis.append(KPI(name, float(self.dataframe[name].mean()), unit, precision))
        return kpis

    def _params_xml(self) -> Path:
        return self.workflow_path / "params.xml"

    def _params(self) -> et.Element:
        tree = et.parse(self._params_xml())
        return tree.getroot()

    def _params_element(self, xpath: str) -> et.Element:
        element = self._params().find(xpath)
        if element is None:
            raise ValueError(f"'{xpath}' not found in XML")
        return element

    @property
    def name(self) -> str:
        return self._params_element(".//void[@property='name']/string").text or "Unknown"

    @property
    def short_name(self) -> str:
        return self._params_element(".//void[@property='shortName']/string").text or "Unknown"

    @property
    def workflow_provider(self) -> str:
        return self._params_element(".//void[@property='workflowProvider']/object").get("class", "Unknown provider")

    @property
    def workflow_path(self) -> Path:
        first_output_file = self.output_files[0]
        return next(d for d in first_output_file.parents if d.suffix == ".flow")

    @property
    def citygml(self) -> str:
        return self._params_element(".//void[@property='cityGmlFileNames']//string").text or "Unknown CityGML"

    @property
    def project_path(self) -> Path:
        return self.workflow_path.parent

    @property
    def project(self) -> str:
        return self.project_path.stem

    @property
    def repository(self) -> Path:
        return self.project_path.parent

    def relative_paths(self) -> list[str]:
        return [f.relative_to(self.workflow_path).as_posix() for f in self.output_files]

    def __str__(self) -> str:
        return f"{self.description} ({self.class_name})"

    def __repr__(self) -> str:
        files_description = "\n".join(f"  - {f}" for f in self.relative_paths())
        kpis = "\n".join(f"  {k}" for k in self.kpis)
        return (
            f"{self.description} ({self.class_name}, {self.project}, "
            f"{self.workflow_path.name}, {self.citygml})\n{files_description}\n{kpis}"
        )

    @property
    def class_name(self) -> str:
        return type(self).__name__

    @property
    def csv_path(self) -> Path:
        csvs = self.get_all_by_extension(".csv")
        relevant_csvs = [csv for csv in csvs if self.csv_identifier in csv.name]
        if len(relevant_csvs) == 0:
            raise ValueError("Workflow didn't seem to have returned any CSV file!")
        if len(relevant_csvs) > 1:
            logger.warning("Too many CSV results found for %s", self)
        return relevant_csvs[-1]

    def csv_content(self) -> bytes:
        with open(self.csv_path, "rb") as f:
            return f.read()

    def get_all_by_extension(self, ext: str) -> list[Path]:
        return [f for f in sorted(self.output_files) if f.suffix == ext]

    @property
    def diagrams(self) -> dict[str, Path]:
        return {}

    def log_content(self) -> bytes:
        logs = self.get_all_by_extension(".log")
        if len(logs) == 0:
            return b""
        with open(logs[-1], "rb") as f:
            return f.read()

    def zip_content(self) -> bytes:
        zip_buffer = io.BytesIO()
        with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED, False) as zf:
            for file_path in self.output_files:
                zf.write(file_path, arcname=file_path.name)
        zip_buffer.seek(0)
        return zip_buffer.read()

    def to_dict(self) -> dict:
        data = {
            "description": self.description,
            "output_files": [str(path) for path in self.output_files],
            "_type": self.class_name,
            "_df_dict": self.dataframe.to_dict("records"),
        }
        if self.dataframe.attrs:
            data["_df_attrs"] = self.dataframe.attrs
        return data

    @classmethod
    def from_dict(cls, data: dict) -> "SimStadtResults":
        data.pop("_type", None)
        df_dict = data.pop("_df_dict", None)
        df_attrs = data.pop("_df_attrs", {})
        data["output_files"] = [Path(p) for p in data["output_files"]]

        if cls is SimStadtResults:
            provider = get_workflow_provider(data["output_files"])
            result = cls.create(provider, **data)
        else:
            result = cls(**data)

        result._df = pd.DataFrame(df_dict)
        result._df.attrs = df_attrs
        return result

    def to_json(self) -> str:
        return json.dumps(self.to_dict())

    @classmethod
    def from_json(cls, json_str: str) -> "SimStadtResults":
        return cls.from_dict(json.loads(json_str))

    @classmethod
    def create(cls, provider: str, **kwargs) -> "SimStadtResults":
        if provider not in cls._registry:
            raise ValueError(f"Unknown workflow provider: {provider}")
        return cls._registry[provider](**kwargs)


@dataclass(repr=False)
class HeatDemandResults(SimStadtResults):
    """Results for Heat Demand (and Heat+Cool Demand) workflows."""

    MAX_FLAT_ROOF_DIFFERENCE = 0.1  # [m]
    TOTAL_DEMAND = "Total Yearly Heating + DHW demand"

    supported_providers: ClassVar[list[str]] = [
        Providers.HEAT_DEMAND,
        Providers.HEAT_DEMAND_WITH_SHADOW,
        Providers.HEAT_DEMAND_WITH_REFURBISHMENT,
        Providers.HEAT_DEMAND_WITH_HISTORIC_REFURBISHMENT,
    ]
    csv_identifier = "DIN18599"

    def _parse_results(self) -> pd.DataFrame:
        csv_decimal = detect_decimal(self.csv_path, 5, "Latitude")
        df = pd.read_csv(self.csv_path, skiprows=list(range(19)) + [20], sep=";", decimal=csv_decimal)
        df["has_flat_roof"] = (
            df["Ridge/mean Height"] - df["Eaves/mean Height"] < self.MAX_FLAT_ROOF_DIFFERENCE
        ).fillna(False)
        df.attrs["Heating"] = "Yearly Heating demand" in df.columns
        df.attrs["Cooling"] = "Yearly Cooling demand" in df.columns
        return df

    @property
    def kpis(self) -> list[KPI]:
        df = self.dataframe
        custom_kpis = [KPI("Number of buildings", df.shape[0], precision=0)]
        sums = [("Heated area", "m²"), ("Footprint area", "m²")]

        if df.attrs["Heating"]:
            specific_heat_demand = df[self.TOTAL_DEMAND].sum() / df["Heated area"].sum()
            heated_buildings = int((df[self.TOTAL_DEMAND] > 10_000).sum())
            custom_kpis.extend(
                [
                    KPI("Number of heated buildings", heated_buildings, precision=0),
                    KPI("Specific Heating Demand", specific_heat_demand, "kWh / (m² · a)", precision=0),
                ]
            )
            sums.extend(
                [
                    ("Yearly Heating demand", "kWh / a"),
                    (self.TOTAL_DEMAND, "kWh / a"),
                ]
            )

        if df.attrs["Cooling"]:
            specific_cooling_demand = df["Yearly Cooling demand"].sum() / df["Heated area"].sum()
            custom_kpis.append(KPI("Specific Cooling Demand", specific_cooling_demand, "kWh / (m² · a)", precision=0))
            sums.append(("Yearly Cooling demand", "kWh / a"))

        averages = [("Mean Uvalue", "W / (m² · K)", 1), ("Year of construction", None), ("Storey number", None)]
        return custom_kpis + self._prepare_kpis(sums, averages)

    @property
    def diagrams(self) -> dict[str, Path]:
        heat_png = self.workflow_path / "heating.png"
        ax = self.monthly_df().plot.bar(
            rot=0,
            ylabel="[MWh]",
            width=0.8,
            color={"Monthly Heating Demand": "darkred", "Monthly Cooling Demand": "darkblue"},
        )
        plt.savefig(heat_png, bbox_inches="tight", dpi=300)
        plt.close(ax.figure)
        return {"Monthly demands": heat_png}

    def monthly_df(self) -> pd.DataFrame:
        months = ["Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"]
        df = self.dataframe
        monthly_df = pd.DataFrame([], index=pd.Index(months))
        for mode in ["Heating", "Cooling"]:
            search_str = f"{mode} demand"
            found_cols = [
                col for col in df.columns if search_str in col and "Yearly" not in col and "Specific" not in col
            ]
            if found_cols:
                monthly_df[f"Monthly {mode} Demand"] = (df[found_cols].sum() / 1000).values
        return monthly_df


@dataclass(repr=False)
class PhotovoltaicResults(SimStadtResults):
    """Results for Photovoltaic (PV) workflows."""

    supported_providers: ClassVar[list[str]] = [
        Providers.PHOTOVOLTAIC,
        Providers.PHOTOVOLTAIC_FINANCE,
        Providers.PHOTOVOLTAIC_WITH_SHADOW,
    ]
    csv_identifier = "_pv_potential"

    def _count_header_lines(self) -> int:
        count = 0
        with open(self.csv_path) as csv:
            for line in csv:
                if line.startswith("Building ID"):
                    return count
                count += 1
        raise ValueError(f"Header not found in {self.csv_path}")

    def _parse_results(self) -> pd.DataFrame:
        csv_decimal = detect_decimal(self.csv_path, 4, "Latitude")
        header_length = self._count_header_lines()
        df = pd.read_csv(
            self.csv_path,
            skiprows=list(range(header_length)) + [header_length + 1],
            sep=";",
            decimal=csv_decimal,
        )
        df = df.rename(columns={"Area": "Roof area for PV"})
        return df.dropna(axis=1, how="all")

    @property
    def kpis(self) -> list[KPI]:
        sums = [("Roof area for PV", "m²"), ("PV potential nominal power", "kWp"), ("PV potential yield", "MWh / a")]
        averages = [("Irradiance in module plane", "W / m²"), ("PV specific yield", "kWh / (kWp · a)")]
        return self._prepare_kpis(sums, averages)


@dataclass(repr=False)
class LoadProfileResults(SimStadtResults):
    """Results for LoadProfile workflows.

    The dataframe has one column per building (building GML ID as column name)
    and 8760 rows of hourly energy demand [kWh/h].
    """

    supported_providers: ClassVar[list[str]] = [Providers.LOAD_PROFILE]
    csv_identifier = "_load_profile_Hourly"

    def _parse_results(self) -> pd.DataFrame:
        csv_decimal = detect_decimal(self.csv_path, 6, "Area")
        df = pd.read_csv(self.csv_path, skiprows=range(1, 12), sep=";", decimal=csv_decimal)
        # First column is the timestamp; remaining columns are per-building loads
        return df.drop(columns=[df.columns[0]])

    @property
    def kpis(self) -> list[KPI]:
        # TODO: Add sum, average. Add timestep too?
        # TODO: Add people? Add total heated area?
        df = self.dataframe
        total = df.sum().sum()
        return [
            KPI("Number of buildings", df.shape[1], precision=0),
            KPI("Total load", total, "kWh/a", precision=1),
            KPI("Average load", total / 8760, "kWh/h", precision=1),
        ]


@dataclass(repr=False)
class SolarPotentialResults(SimStadtResults):
    """Results for SolarPotential workflows.

    The output is a .prn weather file with 8760 hourly rows and three columns:
    GHI (W/m²), DHI (W/m²), Ta (°C).
    """

    supported_providers: ClassVar[list[str]] = [Providers.SOLAR_POTENTIAL]
    csv_identifier = "_solar_potential"  # no CSV output; csv_path will raise if called

    def _parse_results(self) -> pd.DataFrame:
        prn_paths = self.get_all_by_extension(".prn")
        if len(prn_paths) != 1:
            raise ValueError(f"Expected exactly one .prn file, found {len(prn_paths)}")
        return pd.read_csv(prn_paths[0], sep=r"\s+", names=["GHI", "DHI", "Ta"], header=None)

    @property
    def kpis(self) -> list[KPI]:
        df = self.dataframe
        return [
            #TODO: Add averages too?
            KPI("Annual GHI", df["GHI"].sum() / 1000, "kWh/m²", precision=1),
            KPI("Annual DHI", df["DHI"].sum() / 1000, "kWh/m²", precision=1),
            KPI("Mean temperature", float(df["Ta"].mean()), "°C", precision=1),
        ]


@dataclass(repr=False)
class GreenWaterResults(SimStadtResults):
    """Results for GreenWater workflows (hourly time series with tree species breakdown)."""

    csv_identifier = "_greenwater"
    supported_providers: ClassVar[list[str]] = [Providers.GREEN_WATER]

    def _parse_results(self) -> pd.DataFrame:
        csv_decimal = detect_decimal(self.csv_path, 5, "tree area")
        df = pd.read_csv(
            self.csv_path,
            sep=";",
            header=[0, 1],
            skiprows=list(range(8)) + list(range(10, 13)),
            decimal=csv_decimal,
        )

        # Flatten multi-level columns: species names are forward-filled across sub-columns
        cols = df.columns.tolist()
        new_cols = []
        current_species = "All"
        species = []

        for col in cols:
            level0, level1 = col
            if pd.notna(level0) and not level0.startswith("Unnamed") and not level0.startswith("#"):
                current_species = level0
                species.append(current_species)
            new_cols.append(level1 if current_species == "All" else f"{current_species} - {level1}")

        df.columns = new_cols

        # Parse header metadata (tree counts, areas per species)
        with open(self.csv_path, encoding="utf-8") as csv:
            line = ""
            for _ in range(5):
                line = next(csv)
            tree_count = int(line.split(";")[1])
            line = next(csv)
            tree_area_header = re.split("[;,]", line.strip())
            total_tree_area = float(tree_area_header[1].replace(",", "."))
            if tree_count == 0:
                raise ValueError(f"CityGML {self.citygml} does not have any tree!")
            for _ in range(5):
                line = next(csv)
            species_count = [int(c) for c in re.split(";+", line.strip())[1:-1]]
            line = next(csv)
            species_area = [float(c.replace(",", ".")) for c in re.split(";+", line.strip())[1:-1]]

        species.pop(0)
        df.attrs = {
            "species": {
                name: {"count": count, "area": area}
                for name, count, area in zip(species, species_count, species_area)
            },
            "tree count": tree_count,
            "tree area": total_tree_area,
        }
        return df

    def df_with_time(self) -> pd.DataFrame:
        """Return the dataframe with a DatetimeIndex (year 2005, hourly)."""
        df = self.dataframe
        times = pd.date_range(start="2005-01-01 00:00", freq="1h", periods=8760)
        df = df.set_index(times)
        return df.drop(columns=["#"])

    @property
    def kpis(self) -> list[KPI]:
        df = self.dataframe
        rain = df["Rain [mm]"].sum()
        etc = df["Average ETc [mm]"].sum()
        cooling = df["Evaporative cooling [kWh/h]"].sum() / 1000
        irrigation_mm = df["Average irrigation [mm]"].sum()
        tree_area = df.attrs["tree area"]
        irrigation_m3 = irrigation_mm * tree_area / 1000
        return [
            KPI("Number of trees", df.attrs["tree count"], precision=0),
            KPI("Projected tree area", tree_area, "m²", precision=1),
            KPI("Rain", rain, "mm / a", precision=0),
            KPI("Average ETc", etc, "mm / a", precision=0),
            KPI("Irrigation", irrigation_m3, "m³ / a", precision=0),
            KPI("Evaporative cooling", cooling, "MWh / a", precision=0),
        ]

    @property
    def more_info(self) -> dict[str, list[KPI]]:
        df = self.dataframe
        all_info = {}
        for name, info in df.attrs["species"].items():
            count = info["count"]
            area = info["area"]
            irrigation_mm = df[f"{name} - Average irrigation [mm]"].sum()
            average_area = area / count
            all_info[f"_{name}_"] = [
                KPI("Number of trees", count, precision=0),
                KPI("Average tree area", average_area, "m²", precision=1),
                KPI("Average ETc", df[f"{name} - Average ETc [mm]"].sum(), "mm / a", precision=0),
                KPI("Irrigation", irrigation_mm, "mm / a", precision=0),
                KPI("Average irrigation", average_area * irrigation_mm / 1000, "m³ / a", precision=0),
            ]
        return all_info

    @property
    def diagrams(self) -> dict[str, Path]:
        rain_irrigation_png = self.workflow_path / "rain_and_irrigation.png"
        df_months = self.df_with_time()[["Rain [mm]", "Average irrigation [mm]"]].resample("ME").sum()
        fig, ax = plt.subplots()
        df_months = df_months.rename(
            columns={"Rain [mm]": "Niederschlag", "Average irrigation [mm]": "Künstliche Bewässerung"}
        )
        df_months.plot(kind="bar", stacked=True, ax=ax, ylabel="[mm]", figsize=(13, 6))
        ax.set_xticklabels([x.strftime("%b") for x in df_months.index], rotation=0)
        plt.savefig(rain_irrigation_png, bbox_inches="tight", dpi=300)
        plt.close(fig)
        return {"Rain and Irrigation": rain_irrigation_png}


def create_simstadt_results(description: str, output_files: list[Path]) -> SimStadtResults:
    """Factory: detect workflow type from params.xml and return the correct subclass."""
    provider = get_workflow_provider(output_files)
    return SimStadtResults.create(provider, description=description, output_files=sorted(output_files))
