from enum import Enum
from typing import List, Optional, Union
from sqlalchemy import UUID, Column, Float, ForeignKey, Integer, String
from sqlalchemy.orm import relationship
from src.database.core import Base, DbSession
from src.models import DefaultMixin, IdentityMixin, TimeStampMixin, UUIDMixin


class OverhaulReferenceType(str, Enum):
    SCOPE = "SCOPE"
    ASSET = "ASSET"


class CalculationParam(Base, DefaultMixin, IdentityMixin):
    __tablename__ = "oh_ms_calculation_param"

    avg_failure_cost = Column(Float, nullable=False)
    overhaul_cost = Column(Float, nullable=False)

    # Relationships
    calculation_data = relationship(
        "CalculationData", back_populates="parameter")
    results = relationship("CalculationResult", back_populates="parameter")

    # @classmethod
    # async def create_with_references(
    #     cls,
    #     db: DbSession,
    #     avg_failure_cost: float,
    #     overhaul_cost: float,
    #     created_by: str,
    #     # list of {"reference_type": OverhaulReferenceType, "reference_id": str}
    #     references: List[dict]
    # ):
    #     # Create parameter
    #     param = cls(
    #         avg_failure_cost=avg_failure_cost,
    #         overhaul_cost=overhaul_cost,
    #         created_by=created_by
    #     )
    #     db.add(param)
    #     await db.flush()  # Flush to get the param.id

    #     # Create reference links
    #     for ref in references:
    #         reference_link = ReferenceLink(
    #             parameter_id=param.id,
    #             overhaul_reference_type=ref["reference_type"],
    #             reference_id=ref["reference_id"]
    #         )
    #         db.add(reference_link)

    #     await db.commit()
    #     await db.refresh(param)
    #     return param


class CalculationData(Base, DefaultMixin, IdentityMixin):
    __tablename__ = "oh_tr_calculation_data"

    parameter_id = Column(UUID(as_uuid=True), ForeignKey(
        'oh_ms_calculation_param.id'), nullable=True)
    overhaul_reference_type = Column(String(10), nullable=False)
    reference_id = Column(String, nullable=False)
    optimum_oh_day = Column(Integer, nullable=True)

    parameter = relationship(
        "CalculationParam", back_populates="calculation_data")
    scope = relationship("Scope", foreign_keys=[
                         reference_id], primaryjoin="and_(CalculationData.overhaul_reference_type=='SCOPE', " "CalculationData.reference_id==Scope.id)")
    equipment = relationship("ScopeEquipment", foreign_keys=[
                             reference_id], primaryjoin="and_(CalculationData.overhaul_reference_type=='ASSET', " "CalculationData.reference_id==ScopeEquipment.assetnum)")

    @classmethod
    async def create_with_param(
        cls,
        overhaul_reference_type: str,
        reference_id: Union[str, UUID],
        db: DbSession,
        avg_failure_cost: Optional[float],
        overhaul_cost: Optional[float],
        created_by: str,
        params_id: Optional[UUID]
    ):
        if not params_id:
            # Create Params
            params = CalculationParam(
                avg_failure_cost=avg_failure_cost,
                overhaul_cost=overhaul_cost,
                created_by=created_by
            )

            db.add(params)
            await db.flush()
            params_id = params.id

        calculation_data = cls(
            overhaul_reference_type=overhaul_reference_type,
            reference_id=str(reference_id),
            created_by=created_by,
            parameter_id=params_id
        )

        db.add(calculation_data)

        await db.commit()
        await db.refresh(calculation_data)

        return calculation_data


class CalculationResult(Base, DefaultMixin):

    __tablename__ = "oh_tr_calculation_result"

    parameter_id = Column(UUID(as_uuid=True), ForeignKey(
        'oh_ms_calculation_param.id'), nullable=False)
    calculation_data_id = Column(UUID(as_uuid=True), ForeignKey(
        'oh_tr_calculation_data.id'), nullable=False)
    day = Column(Integer, nullable=False)
    corrective_cost = Column(Float, nullable=False)
    overhaul_cost = Column(Float, nullable=False)
    num_failures = Column(Integer, nullable=False)

    parameter = relationship("CalculationParam", back_populates="results")
    reference_link = relationship("CalculationData")