from __future__ import annotations

import json
from collections.abc import Iterator
from enum import Enum, auto
from typing import Annotated, Any, ClassVar, Optional

import pandas as pd
from pydantic import BaseModel, constr, validator
from typing_extensions import TypeAlias

from cognite.powerops.prerun_transformations.transformations import Transformation as TransformationV2

ExternalId: TypeAlias = Annotated[str, constr(min_length=1, max_length=255)]


class Auction(str, Enum):
    week = "week"
    weekend = "weekend"


class RetrievalType(Enum):
    RANGE = auto()
    START = auto()
    END = auto()


class AggregationMethod(Enum):
    sum = auto()
    mean = auto()
    std = auto()
    sem = auto()
    max = auto()
    min = auto()
    median = auto()
    first = auto()
    last = auto()


class TransformationType(Enum):
    GATE_SCHEDULE_FLAG_VALUE_MAPPING = auto()
    GENERATOR_PRODUCTION_SCHEDULE_FLAG_VALUE_MAPPING = auto()
    PLANT_PRODUCTION_SCHEDULE_FLAG_VALUE_MAPPING = auto()
    GATE_OPENING_METER_TO_PERCENT = auto()
    TO_BOOL = auto()
    TO_INT = auto()
    ZERO_IF_NOT_ONE = auto()
    MULTIPLY = auto()
    MULTIPLY_FROM_OFFSET = auto()
    ADD = auto()
    ADD_FROM_OFFSET = auto()
    DYNAMIC_ADD_FROM_OFFSET = auto()
    RESERVOIR_LEVEL_TO_VOLUME = auto()
    STATIC = auto()
    DYNAMIC_STATIC = auto()
    ONE_IF_TWO = auto()
    ADD_WATER_IN_TRANSIT = auto()


ATTRIBUTE_DEFAULT_AGGREGATION: dict[str, AggregationMethod | None] = {
    "gate.max_flow": AggregationMethod.min,
    "gate.min_flow": AggregationMethod.max,
    "gate.schedule_flag": AggregationMethod.first,
    "gate.schedule_m3s": AggregationMethod.first,
    "generator.committed_flag": AggregationMethod.first,
    "generator.committed_in": None,
    "generator.initial_state": None,
    "generator.maintenance_flag": AggregationMethod.max,
    "generator.max_p_constr": AggregationMethod.min,
    "generator.min_p_constr": AggregationMethod.max,
    "generator.production_schedule_flag": AggregationMethod.first,
    "market.buy_price": AggregationMethod.mean,
    "market.sale_price": AggregationMethod.mean,
    "plant.feeding_fee": AggregationMethod.mean,
    "plant.maintenance_flag": AggregationMethod.max,
    "plant.max_p_constr": AggregationMethod.min,
    "plant.max_q_constr": AggregationMethod.min,
    "plant.min_p_constr": AggregationMethod.max,
    "plant.mip_flag": AggregationMethod.first,
    "plant.production_schedule": AggregationMethod.first,
    "plant.production_schedule_flag": AggregationMethod.first,
    "reserve_group.rr_up_obligation": AggregationMethod.max,
    "reserve_group.rr_down_obligation": AggregationMethod.max,
    "reservoir.inflow": AggregationMethod.mean,
    "reservoir.max_vol_constr": AggregationMethod.min,
    "reservoir.min_vol_constr": AggregationMethod.max,
    "reservoir.start_head": None,
    "reservoir.tactical_cost_max": AggregationMethod.mean,
    "reservoir.tactical_cost_min": AggregationMethod.mean,
    "reservoir.tactical_limit_max": AggregationMethod.mean,
    "reservoir.tactical_limit_min": AggregationMethod.mean,
}


class Transformation(BaseModel):  # type: ignore[no-redef]  # mypy says kwargs is redefined on this line ¯\_(ツ)_/¯
    transformation: TransformationType
    kwargs: Optional[dict] = None

    @validator("transformation", pre=True)
    def to_type(cls, value):
        return TransformationType[value.upper()] if isinstance(value, str) else value

    def to_dict(self) -> dict:
        return {"transformation": self.transformation.name, "kwargs": self.kwargs}


class TimeSeriesMappingEntryV2(BaseModel):
    object_type: str
    object_name: str
    attribute_name: str
    time_series_external_id: Optional[str] = None
    transformations: Optional[list[TransformationV2]] = None
    retrieve: Optional[RetrievalType] = None
    aggregation: Optional[AggregationMethod] = None

    @validator("aggregation", pre=True)
    def to_enum(cls, value):
        return AggregationMethod[value] if isinstance(value, str) else value

    @validator("aggregation", always=True)
    def set_default(cls, value, values):
        if value is not None:
            return value
        # TODO: do we want default `None` here or raise error?
        return ATTRIBUTE_DEFAULT_AGGREGATION.get(f"{values.get('object_type')}.{values.get('attribute_name')}")

    @validator("retrieve", pre=True)
    def to_retrival_enum(cls, value):
        return RetrievalType[value] if isinstance(value, str) else value

    @property
    def shop_model_path(self) -> str:
        return f"{self.object_type}.{self.object_name}.{self.attribute_name}"

    def _transformations_to_strings(self, max_cols: int) -> list[str]:
        if self.transformations:
            # ensure_ascii=False to allow æøåÆØÅ
            transformation_string = json.dumps([t.to_dict() for t in self.transformations], ensure_ascii=False)
        else:
            transformation_string = json.dumps([])

        max_chars = 255
        return [transformation_string[i * max_chars : (i + 1) * max_chars] for i in range(max_cols)]

    def to_sequence_row(self, max_transformation_cols: int) -> list[str | float]:
        return [
            self.shop_model_path,
            self.time_series_external_id or float("nan"),
            *self._transformations_to_strings(max_cols=max_transformation_cols),
            self.retrieve.name if self.retrieve else float("nan"),
            self.aggregation.name if self.aggregation else float("nan"),
        ]


class TimeSeriesMappingEntry(BaseModel):
    object_type: str
    object_name: str
    attribute_name: str
    time_series_external_id: Optional[str] = None
    transformations: Optional[list[Transformation]] = None
    retrieve: Optional[RetrievalType] = None
    aggregation: Optional[AggregationMethod] = None

    @validator("aggregation", pre=True)
    def to_enum(cls, value):
        return AggregationMethod[value] if isinstance(value, str) else value

    @validator("aggregation", always=True)
    def set_default(cls, value, values):
        if value is not None:
            return value
        # TODO: do we want default `None` here or raise error?
        return ATTRIBUTE_DEFAULT_AGGREGATION.get(f"{values.get('object_type')}.{values.get('attribute_name')}")

    @validator("retrieve", pre=True)
    def to_retrival_enum(cls, value):
        return RetrievalType[value] if isinstance(value, str) else value

    @property
    def shop_model_path(self) -> str:
        return f"{self.object_type}.{self.object_name}.{self.attribute_name}"

    def _transformations_to_strings(self, max_cols: int) -> list[str]:
        if self.transformations:
            # ensure_ascii=False to allow æøåÆØÅ
            transformation_string = json.dumps([t.to_dict() for t in self.transformations], ensure_ascii=False)
        else:
            transformation_string = json.dumps([])

        max_chars = 255
        return [transformation_string[i * max_chars : (i + 1) * max_chars] for i in range(max_cols)]

    def to_sequence_row(self, max_transformation_cols: int) -> list[str | float]:
        return [
            self.shop_model_path,
            self.time_series_external_id or float("nan"),
            *self._transformations_to_strings(max_cols=max_transformation_cols),
            self.retrieve.name if self.retrieve else float("nan"),
            self.aggregation.name if self.aggregation else float("nan"),
        ]


class TimeSeriesMappingV2(BaseModel):
    rows: list[TimeSeriesMappingEntryV2] = []
    columns: ClassVar[list[str]] = [
        "shop_model_path",
        "time_series_external_id",
        "transformations",
        "transformations1",
        "transformations2",
        "transformations3",
        "retrieve",
        "aggregation",
    ]

    @property
    def transformations_cols(self) -> list[str]:
        return [col for col in self.columns if col.startswith("transformations")]

    def __iter__(self) -> Iterator[TimeSeriesMappingEntry]:
        yield from self.rows

    def __len__(self) -> int:
        return len(self.rows)

    def __add__(self, other: TimeSeriesMapping) -> TimeSeriesMapping:
        return TimeSeriesMapping(rows=self.rows + other.rows)

    def append(self, element: TimeSeriesMappingEntry) -> None:
        self.rows.append(element)

    def extend(self, other: TimeSeriesMapping) -> None:
        self.rows.extend(other.rows)

    @property
    def column_definitions(self) -> list[dict]:
        return [{"valueType": "STRING", "externalId": col} for col in self.columns]

    def to_dataframe(self) -> pd.DataFrame:
        rows = [row.to_sequence_row(max_transformation_cols=len(self.transformations_cols)) for row in self.rows]
        return pd.DataFrame(data=rows, columns=self.columns)

    def dumps(self) -> dict[str, Any]:
        rows = []
        for row in self.rows:
            row_raw = row.dict(exclude={"aggregation", "retrieve", "transformations"})
            if row.aggregation:
                row_raw["aggregation"] = row.aggregation.name
            if row.retrieve:
                row_raw["retrieve"] = row.retrieve.name
            if row.transformations:
                row_raw["transformations"] = [
                    {
                        **transformation.dict(exclude={"transformation"}),
                        "transformation": transformation.transformation.name,
                    }
                    for transformation in row.transformations
                ]
            rows.append(row_raw)
        return {"rows": rows}

    @classmethod
    def load_from_dict(cls, config: dict) -> TimeSeriesMappingV2:
        time_series_mappings_v2 = {"rows": []}
        for mapping_entry in config["rows"]:
            if transformations := mapping_entry.get("transformations"):
                loaded_transformations = [TransformationV2.load(t) for t in transformations]
                del mapping_entry["transformations"]
                time_series_mappings_v2["rows"].append(
                    TimeSeriesMappingEntryV2(transformations=loaded_transformations, **mapping_entry)
                )
                continue
            time_series_mappings_v2["rows"].append(TimeSeriesMappingEntryV2(**mapping_entry))

        return cls(**time_series_mappings_v2)


class TimeSeriesMapping(BaseModel):
    rows: list[TimeSeriesMappingEntry] = []
    columns: ClassVar[list[str]] = [
        "shop_model_path",
        "time_series_external_id",
        "transformations",
        "transformations1",
        "transformations2",
        "transformations3",
        "retrieve",
        "aggregation",
    ]

    @property
    def transformations_cols(self) -> list[str]:
        return [col for col in self.columns if col.startswith("transformations")]

    def __iter__(self) -> Iterator[TimeSeriesMappingEntry]:
        yield from self.rows

    def __len__(self) -> int:
        return len(self.rows)

    def __add__(self, other: TimeSeriesMapping) -> TimeSeriesMapping:
        return TimeSeriesMapping(rows=self.rows + other.rows)

    def append(self, element: TimeSeriesMappingEntry) -> None:
        self.rows.append(element)

    def extend(self, other: TimeSeriesMapping) -> None:
        self.rows.extend(other.rows)

    @property
    def column_definitions(self) -> list[dict]:
        return [{"valueType": "STRING", "externalId": col} for col in self.columns]

    def to_dataframe(self) -> pd.DataFrame:
        rows = [row.to_sequence_row(max_transformation_cols=len(self.transformations_cols)) for row in self.rows]
        return pd.DataFrame(data=rows, columns=self.columns)

    def dumps(self) -> dict[str, Any]:
        rows = []
        for row in self.rows:
            row_raw = row.dict(exclude={"aggregation", "retrieve", "transformations"})
            if row.aggregation:
                row_raw["aggregation"] = row.aggregation.name
            if row.retrieve:
                row_raw["retrieve"] = row.retrieve.name
            if row.transformations:
                row_raw["transformations"] = [
                    {
                        **transformation.dict(exclude={"transformation"}),
                        "transformation": transformation.transformation.name,
                    }
                    for transformation in row.transformations
                ]
            rows.append(row_raw)
        return {"rows": rows}
