from enum import Enum
from pydantic import BaseModel, ConfigDict
from typing import Literal, Union, Optional, List, Tuple, TYPE_CHECKING

from prodsys.models.performance_indicators import (
    KPIEnum,
    KPILevelEnum,
    KPI_UNION,
)


class Event(BaseModel):
    """
    Class that represents an event in the event log.

    Args:
        time (float): Time of the event.
        resource (str): Resource of the event.
        state (str): State of the event.
        state_type (str): State type of the event, representing a state.StateTypeEnum.
        activity (str): Activity of the event, representing a state.StateEnum.
        product (Optional[str], optional): Product of the event. Defaults to None.
        expected_end_time (Optional[float], optional): Expected end time of the event. Defaults to None.
        target_location (Optional[str], optional): Target location of the event. Defaults to None.
    """

    time: float
    resource: str
    state: str
    state_type: str
    activity: str
    product: Optional[str] = None
    expected_end_time: Optional[float] = None
    target_location: Optional[str] = None

    model_config = ConfigDict(json_schema_extra={
        "examples": [
            {
                "time": 12.0,
                "resource": "R1",
                "state": "P1",
                "state_type": "Production",
                "activity": "start state",
                "product": "Product_1_12",
                "expected_end_time": 24.3,
                "target_location": None,
            }, 
            {
                "time": 24.3,
                "resource": "R1",
                "state": "P1",
                "state_type": "Production",
                "activity": "end state",
                "product": "Product_1_12",
                "expected_end_time": None,
                "target_location": "L1",
            },
        ]
    })


class Performance(BaseModel):
    """
    Class that represents the performance of a simulation run.

    Args:
        event_log (List[Event]): Event log of the simulation run.
        kpis (List[KPI_UNION]): List of KPIs of the simulation run.
    """

    model_config = ConfigDict(json_schema_extra={
        "examples": [
            {
                "event_log": Event.model_config["json_schema_extra"]["examples"],
                "kpis": [
                    {
                        "name": "throughput",
                        "target": "max",
                        "weight": 1,
                        "value": 4.32,
                        "context": ["system", "product_type"],
                        "product_type": "ProductType_1",
                    },
                    {
                        "name": "WIP",
                        "target": "min",
                        "weight": 1,
                        "value": 121,
                        "context": ["system", "product_type"],
                        "product_type": "ProductType_1",
                    },
                ],
            }
        ]
    })

    event_log: List[Event]
    kpis: List[KPI_UNION]

    def get_kpi_for_context(self, context: Tuple[KPILevelEnum, ...]) -> List[KPI_UNION]:
        return [kpi for kpi in self.kpis if context == kpi.context]

    def get_kpi_for_name(self, name: KPIEnum) -> List[KPI_UNION]:
        return [kpi for kpi in self.kpis if name == kpi.name]

    def get_kpi_for_context_and_name(
        self, context: Tuple[KPILevelEnum, ...], name: KPIEnum
    ) -> List[KPI_UNION]:
        return [kpi for kpi in self.kpis if context == kpi.context and name == kpi.name]
