from decimal import Decimal
import os
from typing import Optional
import json
import uuid
import httpx
from fastapi import HTTPException, status
from sqlalchemy import Delete, desc, Select, select, func
from sqlalchemy.orm import selectinload

from src.aeros_equipment.service import save_default_equipment
from src.aeros_simulation.service import save_default_simulation_node
from src.auth.service import CurrentUser
from src.config import WINDOWS_AEROS_BASE_URL, CLAMAV_HOST, CLAMAV_PORT
from src.aeros_utils import aeros_post, aeros_file_upload
from src.database.core import DbSession
from src.database.service import search_filter_sort_paginate
from src.utils import sanitize_filename
import clamd
import io

from .model import (
    AerosProject, 
    SimulationMasterType, 
    SimulationModel, 
    SimulationModelAsset, 
    SimulationModelAssetDistribution, 
    SimulationModelAssetOverhaul, 
    SimulationModelAssetPm, 
    SimulationModelAssetRepairDuration, 
    SimulationModelSchema, 
    SimulationResult, 
    SimulationResultAsset, 
    SimulationResultEvent, 
    SimulationResultSchema
)
from .schema import AerosProjectInput, OverhaulScheduleCreate, OverhaulScheduleUpdate
import asyncio
ALLOWED_EXTENSIONS = {".aro"}
MAX_FILE_SIZE = 100 * 1024 * 1024  # 100MB
# client = httpx.AsyncClient(timeout=300.0) # Managed in src.aeros_utils
# We still use a local client for WINDOWS_AEROS_BASE_URL if it's not managed by aeros_utils
client = httpx.AsyncClient(timeout=300.0)


async def import_aro_project(*, db_session: DbSession, aeros_project_in: AerosProjectInput):
    # windows_aeros_base_url = WINDOWS_AEROS_BASE_URL


    file = aeros_project_in.aro_file
    
    # Sanitize and validate filename
    try:
        clean_filename = sanitize_filename(file.filename)
    except ValueError as e:
        raise HTTPException(
            status_code=400,
            detail=f"Invalid filename: {str(e)}"
        )

    # Check if mime type is application/octet-stream 
    if file.content_type != "application/octet-stream":
        raise HTTPException(
            status_code=400,
            detail="Invalid file type. Allowed: application/octet-stream"
        )


    # Get filename
    filename_without_ext = os.path.splitext(clean_filename)[0]

    # Get file extension
    file_ext = os.path.splitext(clean_filename)[1].lower()

    # Validate file extension
    if file_ext not in ALLOWED_EXTENSIONS:
        raise HTTPException(
            status_code=400,
            detail=f"File type not allowed. Allowed: {ALLOWED_EXTENSIONS}"
        )

    print("read file")


    # Read and check file size
    content = await file.read()
    if len(content) > MAX_FILE_SIZE:
        raise HTTPException(
            status_code=400,
            detail="File too large. Max size: 100Mb"
        )

    # ClamAV Scan
    try:
        cd = clamd.ClamdNetworkSocket(CLAMAV_HOST, CLAMAV_PORT)
        scan_result = cd.instream(io.BytesIO(content))
        # Result format: {'stream': ('FOUND', 'Eicar-Test-Signature')} or {'stream': ('OK', None)}
        if scan_result and scan_result.get('stream') and scan_result['stream'][0] == 'FOUND':
            raise HTTPException(
                status_code=400,
                detail=f"Virus detected: {scan_result['stream'][1]}"
            )
    except clamd.ConnectionError:
        pass
    # except HTTPException:
    #     raise
    # except Exception as e:
    #     print(f"ClamAV error: {e}")
    #     raise HTTPException(
    #         status_code=500,
    #         detail=f"Antivirus check failed: {str(e)}"
    #     )


    # Project name hardcode
    # project_name = "trialapi"

    # Project name
    project_name = filename_without_ext


    ## save File to windows app
    # Output is string of file path, examole
    # # Example response "C/dsad/dsad.aro"
    try:
        # Reset file position since we already read it for size check
        # await file.seek(0)

        # Prepare file for upload
        # files = {
        #     "file": (clean_filename, content, file.content_type or "application/octet-stream")
        # }

        response = await aeros_file_upload(
            "/upload",
            content,
            "file",
            clean_filename,
            base_url=WINDOWS_AEROS_BASE_URL,
        )

        response.raise_for_status()

        # print("fetch")
        # response = await client.post(
        #     f"{WINDOWS_AEROS_BASE_URL}/upload-file",
        #     files=files
        # )
        # response.raise_for_status()

        # Get the file path from the response
        upload_result = response.json()
        aro_path = upload_result.get("full_path")
        filename = upload_result.get("stored_filename").replace(".aro", "")

        if not aro_path:
            raise HTTPException(
                status_code=500,
                detail="Failed to get file path from upload response"
            )

    except httpx.HTTPStatusError as e:
        raise HTTPException(
            status_code=e.response.status_code,
            detail=f"Upload failed: {e.response.text}"
        )
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e)
        )

    await asyncio.sleep(2)

    # aro_path = r"C:/Users/user/Documents/Aeros/sample_project.aro"

    aeros_project = AerosProject(project_name=filename, aro_file_path=aro_path)

    # find aeros record first, if not found, then create a new one
    stmt = select(AerosProject).order_by(desc(AerosProject.created_at)).limit(1)
    result = await db_session.execute(stmt)
    latest_project = result.scalar_one_or_none()

    # If aeros record found, then update it
    if latest_project:
        latest_project.project_name = filename
        latest_project.aro_file_path = aro_path
    else: # else create new aeros record
        db_session.add(aeros_project)

    await db_session.commit()

    # aro_json = json.dumps(aro_path)

    # Update path to AEROS APP
    # Example BODy "C/dsad/dsad.aro"
    try:
        response = await aeros_post(
            "/api/Project/ImportAROFile",
            json=aro_path,
            headers={"Content-Type": "application/json"},
        )
        
        response.raise_for_status() 
        response_json = response.json()
        
        
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e)
        )

    await _initialize_default_project_data(
            db_session=db_session,
            project_name=filename
    )

async def fetch_aro_record(*, db_session: DbSession):
    stmt = select(AerosProject).order_by(desc(AerosProject.updated_at)).limit(1)
    result = await db_session.execute(stmt)
    found_record = result.scalar_one_or_none()

    return found_record

async def get_project(*, db_session: DbSession):
    stmt = select(AerosProject).order_by(desc(AerosProject.updated_at)).limit(1)
    result = await db_session.execute(stmt)
    found_record = result.scalar_one_or_none()

    return found_record

async def _initialize_default_project_data(
    *,
    db_session: DbSession,
    project_name: str
) -> None:
    """
    Initialize default equipment and simulation nodes for a project.

    Args:
        db_session: Database session
        project_name: Name of the project to initialize
    """
    try:
        # Save default equipment
        await save_default_equipment(
            db_session=db_session,
            project_name=project_name
        )

        # # Save default simulation node
        # await save_default_simulation_node(
        #     db_session=db_session,
        #     project_name=project_name
        # )

        await db_session.commit()

    except Exception as e:
        await db_session.rollback()
        raise e


async def reset_project(*, db_session: DbSession):
    project = await fetch_aro_record(db_session=db_session)

    try:
        response = await aeros_post(
            "/api/Project/ImportAROFile",
            data=f'"{project.aro_file_path}"',
            headers={"Content-Type": "application/json"},
        )
        response.raise_for_status()
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e)
        )

    return project.aro_file_path


async def create(*, db_session: DbSession, overhaul_job_in: OverhaulScheduleCreate):
    # Placeholder for creation logic
    print(f"create overhaul job called: {overhaul_job_in}")
    pass

async def update(*, db_session: DbSession, overhaul_schedule_id: str, overhaul_job_in: OverhaulScheduleUpdate):
    # Placeholder for update logic
    print(f"update overhaul job called for id {overhaul_schedule_id}: {overhaul_job_in}")
    pass

async def delete(*, db_session: DbSession, overhaul_schedule_id: str):
    # Placeholder for delete logic
    print(f"delete overhaul job called for id {overhaul_schedule_id}")
    pass


# helper: get or create type
async def get_or_create_type(db: DbSession, name: str):
    result = await db.execute(
        select(SimulationMasterType).where(SimulationMasterType.name == name)
    )
    obj = result.scalar_one_or_none()

    if obj:
        return obj

    obj = SimulationMasterType(name=name)
    db.add(obj)
    await db.flush()  # biar dapet id

    return obj

# helper: convert to float
def to_float(value):
    if isinstance(value, Decimal):
        return float(value)
    return value


# store simulation model
async def import_simulation_model_service(
    db: DbSession,
    name: str,
    data: dict
):

    asset_count = 0
    schema_count = 0

    # =====================================
    # CREATE ROOT MODEL
    # =====================================
    simulation_model = SimulationModel(
        name=name,

        # raw original json
        # payload=data,

        # full gui state
        gui_payload=data.get("gui_state")
    )

    db.add(simulation_model)

    await db.flush()

    # =====================================
    # MEMORY MAPS
    # =====================================
    asset_map = {}
    schema_map = {}

    # =====================================
    # IMPORT ASSETS
    # =====================================
    for asset in data.get("assets", []):

        asset_uuid = uuid.uuid4()

        db_asset = SimulationModelAsset(

            id=asset_uuid,

            simulation_model_id=simulation_model.id,

            assetId=asset.get("id"),

            name=asset.get("name"),

            max_capacity=asset.get(
                "max_capacity"
            ),

            design_capacity=asset.get(
                "design_capacity"
            ),

            buffer_duration=asset.get(
                "buffer_duration"
            ),

            transition_delay=asset.get(
                "transition_delay"
            ),

            initial_age=asset.get(
                "initial_age"
            ),
        )

        db.add(db_asset)

        asset_map[
            db_asset.assetId
        ] = db_asset

        # =================================
        # NORMALIZE DISTRIBUTION
        # =================================
        model = asset.get("model", {})

        dist = asset.get("distribution")

        if not dist:
            dist = model.get(
                "distribution"
            )

        if not dist and model:
            dist = model

        # =================================
        # DISTRIBUTION
        # =================================
        if dist and dist.get("type"):

            dist_type = (
                await get_or_create_type(
                    db,
                    dist.get("type")
                )
            )

            db.add(
                SimulationModelAssetDistribution(

                    id=uuid.uuid4(),

                    simulation_model_assets_id=(
                        asset_uuid
                    ),

                    type_id=dist_type.id,

                    eta=dist.get("eta"),

                    beta=dist.get("beta"),

                    mean=dist.get("mean"),

                    std_dev=dist.get(
                        "std_dev"
                    ),

                    shape=dist.get("shape"),

                    scale=dist.get("scale"),
                )
            )

        # =================================
        # REPAIR
        # =================================
        repair = asset.get(
            "repair_duration"
        )

        if repair and repair.get("type"):

            repair_type = (
                await get_or_create_type(
                    db,
                    repair.get("type")
                )
            )

            db.add(
                SimulationModelAssetRepairDuration(

                    id=uuid.uuid4(),

                    simulation_model_assets_id=(
                        asset_uuid
                    ),

                    type_id=repair_type.id,

                    mean=repair.get("mean"),

                    std_dev=repair.get(
                        "std_dev"
                    ),
                )
            )

        # =================================
        # OVERHAUL
        # =================================
        overhaul = asset.get(
            "overhaul"
        )

        if overhaul:

            db.add(
                SimulationModelAssetOverhaul(

                    id=uuid.uuid4(),

                    simulation_model_assets_id=(
                        asset_uuid
                    ),

                    interval=overhaul.get(
                        "interval"
                    ),

                    duration=overhaul.get(
                        "duration"
                    ),
                )
            )

        # =================================
        # PM
        # =================================
        pm = asset.get("pm")

        if pm:

            db.add(
                SimulationModelAssetPm(

                    id=uuid.uuid4(),

                    simulation_model_assets_id=(
                        asset_uuid
                    ),

                    interval=pm.get(
                        "interval"
                    ),

                    duration=pm.get(
                        "duration"
                    ),

                    restoration_factor=(
                        pm.get(
                            "restoration_factor"
                        )
                    ),
                )
            )

        asset_count += 1

    # =====================================
    # PREPARE SCHEMAS
    # =====================================
    schemas = data.get(
        "schemas",
        {}
    )

    rbd_root = data.get(
        "rbd_root"
    )

    # inject root
    if rbd_root:

        root_id = rbd_root.get(
            "id",
            "rbd_root"
        )

        if root_id not in schemas:
            schemas[root_id] = rbd_root

    # =====================================
    # CREATE SCHEMAS
    # =====================================
    for schema_key, schema_value in schemas.items():

        structure = schema_value.get(
            "structure",
            {}
        )

        type_name = structure.get(
            "type"
        )

        name = schema_value.get(
            "name"
        )

        type_obj = (
            await get_or_create_type(
                db,
                type_name
            )
        )

        db_schema = SimulationModelSchema(

            id=uuid.uuid4(),

            simulation_model_id=(
                simulation_model.id
            ),

            schemaId=schema_key,

            name=name,

            type_id=type_obj.id,
        )

        db.add(db_schema)

        schema_map[
            schema_key
        ] = db_schema

        schema_count += 1

    # =====================================
    # LINK TREE
    # =====================================
    for schema_key, schema_value in schemas.items():

        parent = schema_map.get(
            schema_key
        )

        if not parent:
            continue

        children = (
            schema_value
            .get("structure", {})
            .get("children", [])
        )

        for child in children:

            child_type = child.get(
                "type"
            )

            ref_id = child.get(
                "ref_id"
            )

            # =============================
            # SCHEMA CHILD
            # =============================
            if child_type == "Schema":

                child_schema = (
                    schema_map.get(ref_id)
                )

                if child_schema:

                    child_schema.simulation_model_schemas_id = (
                        parent.id
                    )

            # =============================
            # ASSET CHILD
            # =============================
            elif child_type == "Asset":

                asset_obj = asset_map.get(
                    ref_id
                )

                if asset_obj:

                    asset_obj.simulation_model_schemas_id = (
                        parent.id
                    )

    # =====================================
    # COMMIT
    # =====================================
    await db.commit()

    return {

        "simulation_model_id": (
            simulation_model.id
        ),

        "schema_count": schema_count,

        "asset_count": asset_count,
    }


# =========================================================
# GET SIMULATION MODEL (USE THIS)
# =========================================================
async def get_simulation_model_by_id_service(
    db: DbSession,
    simulation_model_id: str
):

    # ======================================
    # LOAD MODEL
    # ======================================
    stmt = (
        select(SimulationModel)
        .where(
            SimulationModel.id == simulation_model_id
        )
        .options(

            # ==================================
            # SCHEMAS
            # ==================================
            selectinload(
                SimulationModel.schemas
            ).selectinload(
                SimulationModelSchema.type
            ),

            selectinload(
                SimulationModel.schemas
            ).selectinload(
                SimulationModelSchema.assets
            ),

            # ==================================
            # ASSETS
            # ==================================
            selectinload(
                SimulationModel.assets
            )
            .selectinload(
                SimulationModelAsset.asset_distribution
            )
            .selectinload(
                SimulationModelAssetDistribution.type
            ),

            selectinload(
                SimulationModel.assets
            )
            .selectinload(
                SimulationModelAsset.asset_repair_duration
            )
            .selectinload(
                SimulationModelAssetRepairDuration.type
            ),

            selectinload(
                SimulationModel.assets
            )
            .selectinload(
                SimulationModelAsset.asset_overhaul
            ),

            selectinload(
                SimulationModel.assets
            )
            .selectinload(
                SimulationModelAsset.asset_pm
            ),
        )
    )

    result = await db.execute(stmt)

    simulation_model = (
        result.scalar_one_or_none()
    )

    if not simulation_model:
        return None

    # ======================================
    # HELPERS
    # ======================================
    def clean_nulls(data):

        if isinstance(data, dict):

            return {
                k: clean_nulls(v)
                for k, v in data.items()
                if v is not None
            }

        if isinstance(data, list):

            return [
                clean_nulls(v)
                for v in data
            ]

        return data

    # ======================================
    # SCHEMA MAP
    # ======================================
    schema_map = {}

    for schema in simulation_model.schemas:

        schema_map[str(schema.id)] = schema

    # ======================================
    # CHILDREN MAP
    # ======================================
    children_map = {}

    for schema in simulation_model.schemas:

        parent_id = (
            str(
                schema.simulation_model_schemas_id
            )
            if schema.simulation_model_schemas_id
            else None
        )

        children_map.setdefault(
            parent_id,
            []
        ).append(schema)

    # ======================================
    # SERIALIZE ASSET
    # ======================================
    def serialize_asset(asset):

        distribution = None

        if asset.asset_distribution:

            distribution = {

                "type": (
                    asset.asset_distribution.type.name
                    if asset.asset_distribution.type
                    else None
                ),

                "shape": to_float(
                    asset.asset_distribution.shape
                ),

                "scale": to_float(
                    asset.asset_distribution.scale
                ),

                "mean": to_float(
                    asset.asset_distribution.mean
                ),

                "std_dev": to_float(
                    asset.asset_distribution.std_dev
                ),

                "eta": to_float(
                    asset.asset_distribution.eta
                ),

                "beta": to_float(
                    asset.asset_distribution.beta
                ),
            }

        repair_duration = None

        if asset.asset_repair_duration:

            repair_duration = {

                "type": (
                    asset.asset_repair_duration.type.name
                    if asset.asset_repair_duration.type
                    else None
                ),

                "mean": to_float(
                    asset.asset_repair_duration.mean
                ),

                "std_dev": to_float(
                    asset.asset_repair_duration.std_dev
                ),
            }

        overhaul = None

        if asset.asset_overhaul:

            overhaul = {

                "interval": to_float(
                    asset.asset_overhaul.interval
                ),

                "duration": to_float(
                    asset.asset_overhaul.duration
                ),
            }

        pm = None

        if asset.asset_pm:

            pm = {

                "interval": to_float(
                    asset.asset_pm.interval
                ),

                "duration": to_float(
                    asset.asset_pm.duration
                ),

                "restoration_factor": to_float(
                    asset.asset_pm.restoration_factor
                ),
            }

        asset_json = {

            "id": asset.assetId,

            "name": asset.name,

            "max_capacity": to_float(
                asset.max_capacity
            ),

            "design_capacity": to_float(
                asset.design_capacity
            ),

            "buffer_duration": to_float(
                asset.buffer_duration
            ),

            "transition_delay": to_float(
                asset.transition_delay
            ),

            "initial_age": to_float(
                asset.initial_age
            ),

            # normalized
            "distribution": distribution,

            "repair_duration": repair_duration,

            "overhaul": overhaul,

            "pm": pm,
        }

        return clean_nulls(asset_json)

    # ======================================
    # BUILD SCHEMA JSON
    # ======================================
    schemas_json = {}

    root_schema = None

    for schema in simulation_model.schemas:

        if (
            not schema.simulation_model_schemas_id
        ):
            root_schema = schema

        children_json = []

        # ==================================
        # CHILD SCHEMAS
        # ==================================
        child_schemas = children_map.get(
            str(schema.id),
            []
        )

        for child_schema in child_schemas:

            children_json.append({

                "type": "Schema",

                "ref_id": (
                    child_schema.schemaId
                )
            })

        # ==================================
        # CHILD ASSETS
        # ==================================
        for asset in schema.assets:

            children_json.append({

                "type": "Asset",

                "ref_id": asset.assetId
            })

        schemas_json[
            schema.schemaId
        ] = {

            "name": schema.name,

            "structure": {

                "type": (
                    schema.type.name
                    if schema.type
                    else None
                ),

                "children": children_json
            }
        }

    # ======================================
    # ROOT
    # ======================================
    rbd_root = None

    if root_schema:

        rbd_root = schemas_json.pop(
            root_schema.schemaId
        )

    # ======================================
    # ASSETS
    # ======================================
    assets_json = []

    for asset in simulation_model.assets:

        assets_json.append(
            serialize_asset(asset)
        )

    # ======================================
    # FINAL JSON
    # ======================================
    final_json = {

        "id": str(
            simulation_model.id
        ),

        "name": simulation_model.name,

        "rbd_root": rbd_root,

        "schemas": schemas_json,

        "assets": assets_json,

        # restore gui state
        "gui_state": (
            simulation_model.gui_payload
        ),
    }

    return clean_nulls(
        final_json
    )


# store simulation result into DB (JSON: input) 
async def import_simulation_result_service(db: DbSession, data: dict):

    try:
        # =========================================
        # CREATE MAIN RESULT
        # =========================================
        result_obj = SimulationResult(
            production_availability=data.get("production_availability"),
            simulation_duration=data.get("simulation_duration"),
            simulation_replications=data.get("simulation_replications"),
            system_availability=data.get("system_availability"),
            system_downtime=data.get("system_downtime"),
            system_failures=data.get("system_failures"),
            system_uptime=data.get("system_uptime"),
            total_demand=data.get("total_demand"),
            total_productions=data.get("total_productions"),
            system_down_causes=data.get("system_down_causes"),
        )

        db.add(result_obj)

        await db.flush()

        # =========================================
        # PRELOAD MODEL ASSETS
        # assetId -> model asset
        # =========================================
        asset_model_result = await db.execute(
            select(SimulationModelAsset)
        )

        asset_models = asset_model_result.scalars().all()

        asset_model_map = {
            asset.assetId: asset
            for asset in asset_models
        }

        # =========================================
        # PRELOAD MODEL SCHEMAS
        # schemaId -> model schema
        # =========================================
        schema_model_result = await db.execute(
            select(SimulationModelSchema)
        )

        schema_models = schema_model_result.scalars().all()

        schema_model_map = {
            schema.schemaId: schema
            for schema in schema_models
        }

        # =========================================
        # STORE ASSET RESULTS
        # =========================================
        asset_result_map = {}

        for asset in data.get("asset_results", []):

            model_asset = asset_model_map.get(
                asset.get("id")
            )

            asset_result_obj = SimulationResultAsset(
                simulation_results_id=result_obj.id,

                simulation_model_assets_id=(
                    model_asset.id
                    if model_asset
                    else None
                ),

                assetId=asset.get("id"),
                name=asset.get("name"),

                availability=asset.get("availability"),
                contribution_factor=asset.get("contribution_factor"),
                downtime=asset.get("downtime"),
                uptime=asset.get("uptime"),
                failures=asset.get("failures"),

                total_downtime_concurrent_with_system=(
                    asset.get(
                        "total_downtime_concurrent_with_system"
                    )
                ),
            )

            db.add(asset_result_obj)

            await db.flush()

            asset_result_map[
                asset.get("id")
            ] = asset_result_obj

            # =====================================
            # EVENTS
            # =====================================
            for event in asset.get("events", []):

                type_obj = None

                if event.get("event_type"):
                    type_obj = await get_or_create_type(
                        db,
                        event.get("event_type")
                    )

                db.add(
                    SimulationResultEvent(
                        simulation_result_id=result_obj.id,

                        simulation_result_assets_id=(
                            asset_result_obj.id
                        ),

                        type_id=(
                            type_obj.id
                            if type_obj
                            else None
                        ),

                        cause=event.get("cause"),
                        state=event.get("state"),
                        load_factor=event.get("load_factor"),

                        start_time=event.get("start"),
                        end_time=event.get("end"),
                    )
                )

        # =========================================
        # STORE SCHEMA RESULTS
        # =========================================
        schema_result_map = {}

        for schema in data.get("schema_results", []):

            model_schema = schema_model_map.get(
                schema.get("id")
            )

            schema_result_obj = SimulationResultSchema(
                simulation_results_id=result_obj.id,

                simulation_model_schemas_id=(
                    model_schema.id
                    if model_schema
                    else None
                ),

                schemaId=schema.get("id"),
                name=schema.get("name"),

                availability=schema.get("availability"),
                downtime=schema.get("downtime"),
                uptime=schema.get("uptime"),
            )

            db.add(schema_result_obj)

            await db.flush()

            schema_result_map[
                schema.get("id")
            ] = schema_result_obj

            # =====================================
            # EVENTS
            # =====================================
            for event in schema.get("events", []):

                type_obj = None

                if event.get("event_type"):
                    type_obj = await get_or_create_type(
                        db,
                        event.get("event_type")
                    )

                db.add(
                    SimulationResultEvent(
                        simulation_result_id=result_obj.id,

                        simulation_result_schemas_id=(
                            schema_result_obj.id
                        ),

                        type_id=(
                            type_obj.id
                            if type_obj
                            else None
                        ),

                        cause=event.get("cause"),
                        state=event.get("state"),
                        load_factor=event.get("load_factor"),

                        start_time=event.get("start"),
                        end_time=event.get("end"),
                    )
                )

        # =========================================
        # STORE SYSTEM EVENTS
        # =========================================
        for event in data.get("system_events", []):

            type_obj = None

            if event.get("event_type"):
                type_obj = await get_or_create_type(
                    db,
                    event.get("event_type")
                )

            db.add(
                SimulationResultEvent(
                    simulation_result_id=result_obj.id,

                    type_id=(
                        type_obj.id
                        if type_obj
                        else None
                    ),

                    cause=event.get("cause"),
                    state=event.get("state"),
                    load_factor=event.get("load_factor"),

                    start_time=event.get("start"),
                    end_time=event.get("end"),
                )
            )

        # =========================================
        # SINGLE COMMIT
        # =========================================
        await db.commit()

        return {
            "result_id": str(result_obj.id),
            "asset_count": len(
                data.get("asset_results", [])
            ),
            "schema_count": len(
                data.get("schema_results", [])
            ),
            "system_event_count": len(
                data.get("system_events", [])
            ),
        }
        
    except Exception as e:

        # logger.exception(
        #     "Failed importing simulation result"
        # )

        await db.rollback()

        raise


# =========================================================
# GET SIMULATION RESULT V1 (USE THIS)
# =========================================================
async def get_simulation_result_service(
    db: DbSession,
    simulation_result_id: str
):

    # =====================================================
    # LOAD RESULT
    # =====================================================
    stmt = (
        select(SimulationResult)
        .where(
            SimulationResult.id == simulation_result_id
        )
        .options(

            # =========================================
            # ASSETS
            # =========================================
            selectinload(
                SimulationResult.assets
            ),

            # =========================================
            # SCHEMAS
            # =========================================
            selectinload(
                SimulationResult.schemas
            ),

            # =========================================
            # EVENTS
            # =========================================
            selectinload(
                SimulationResult.events
            ).selectinload(
                SimulationResultEvent.type
            ),
        )
    )

    result = await db.execute(stmt)

    simulation_result = (
        result.scalar_one_or_none()
    )

    if not simulation_result:
        return None

    # =====================================================
    # EVENT SERIALIZER
    # =====================================================
    def serialize_event(event):

        return {
            "event_type": (
                event.type.name
                if event.type
                else None
            ),

            "cause": event.cause,
            "state": event.state,

            "start": to_float(
                event.start_time
            ),

            "end": to_float(
                event.end_time
            ),

            "load_factor": to_float(
                event.load_factor
            ),
        }

    # =====================================================
    # GROUP EVENTS
    # =====================================================
    asset_event_map = {}
    schema_event_map = {}
    system_events = []

    # stable ordering
    events_sorted = sorted(
        simulation_result.events,
        key=lambda x: (
            x.start_time or 0
        )
    )

    for event in events_sorted:

        event_json = serialize_event(
            event
        )

        # =========================================
        # ASSET EVENT
        # =========================================
        if event.simulation_result_assets_id:

            key = str(
                event.simulation_result_assets_id
            )

            asset_event_map.setdefault(
                key,
                []
            ).append(event_json)

        # =========================================
        # SCHEMA EVENT
        # =========================================
        elif event.simulation_result_schemas_id:

            key = str(
                event.simulation_result_schemas_id
            )

            schema_event_map.setdefault(
                key,
                []
            ).append(event_json)

        # =========================================
        # SYSTEM EVENT
        # =========================================
        else:

            system_events.append(
                event_json
            )

    # =====================================================
    # ASSET SERIALIZER
    # =====================================================
    def serialize_asset(asset):

        asset_json = {

            "id": asset.assetId,

            "name": asset.name,

            "availability": to_float(
                asset.availability
            ),

            "contribution_factor": to_float(
                asset.contribution_factor
            ),

            "downtime": to_float(
                asset.downtime
            ),

            "uptime": to_float(
                asset.uptime
            ),

            "failures": to_float(
                asset.failures
            ),

            "total_downtime_concurrent_with_system": (
                to_float(
                    asset.total_downtime_concurrent_with_system
                )
            ),

            "events": asset_event_map.get(
                str(asset.id),
                []
            ),
        }

        # remove nulls
        return {
            k: v
            for k, v in asset_json.items()
            if v is not None
        }

    # =====================================================
    # SCHEMA SERIALIZER
    # =====================================================
    def serialize_schema(schema):

        schema_json = {

            "id": schema.schemaId,

            "name": schema.name,

            "availability": to_float(
                schema.availability
            ),

            "downtime": to_float(
                schema.downtime
            ),

            "uptime": to_float(
                schema.uptime
            ),

            "events": schema_event_map.get(
                str(schema.id),
                []
            ),
        }

        return {
            k: v
            for k, v in schema_json.items()
            if v is not None
        }

    # =====================================================
    # BUILD FINAL JSON
    # =====================================================
    final_json = {
        
        "id": str(
            simulation_result.id
        ),

        "production_availability": to_float(
            simulation_result.production_availability
        ),

        "simulation_duration": to_float(
            simulation_result.simulation_duration
        ),

        "simulation_replications": to_float(
            simulation_result.simulation_replications
        ),

        "system_availability": to_float(
            simulation_result.system_availability
        ),

        "system_downtime": to_float(
            simulation_result.system_downtime
        ),

        "system_failures": to_float(
            simulation_result.system_failures
        ),

        "system_uptime": to_float(
            simulation_result.system_uptime
        ),

        "total_demand": to_float(
            simulation_result.total_demand
        ),

        "total_productions": to_float(
            simulation_result.total_productions
        ),

        # =========================================
        # JSONB
        # =========================================
        "system_down_causes": (
            simulation_result.system_down_causes
            or []
        ),

        # =========================================
        # RESULTS
        # =========================================
        "asset_results": [
            serialize_asset(asset)
            for asset in simulation_result.assets
        ],

        "schema_results": [
            serialize_schema(schema)
            for schema in simulation_result.schemas
        ],

        "system_events": system_events,
    }

    # remove null root keys
    final_json = {
        k: v
        for k, v in final_json.items()
        if v is not None
    }

    return final_json