rbd-app/temporal/activity.py

import asyncio
from datetime import datetime
from email.message import EmailMessage
import json
from temporalio import activity
import os
import requests
import redis.asyncio as redis
import smtplib
from uuid import UUID

REDIS_HOST = os.getenv("REDIS_HOST", "192.168.1.82")
REDIS_PORT = int(os.getenv("REDIS_PORT", 6379))
ERROR_NOTIFICATION_TIMEOUT_SECONDS = int(os.getenv("ERROR_NOTIFICATION_TIMEOUT_SECONDS", "10"))


redis_client = redis.Redis(
    host=REDIS_HOST,   # or docker service name
    port=REDIS_PORT,
    db=0,
    decode_responses=True
)


def _shorten(value, limit=2000):
    text = str(value)
    return text if len(text) <= limit else f"{text[:limit]}... [truncated]"


def _send_error_notification(subject: str, message: str, context: dict) -> dict:
    payload = {
        "service": os.getenv("SERVICE_NAME", "be-rbd"),
        "subject": subject,
        "message": _shorten(message),
        "context": context,
        "timestamp": datetime.utcnow().isoformat() + "Z",
    }
    results = {}

    webhook_url = os.getenv("ERROR_NOTIFICATION_WEBHOOK_URL")
    if webhook_url:
        try:
            response = requests.post(webhook_url, json=payload, timeout=ERROR_NOTIFICATION_TIMEOUT_SECONDS)
            response.raise_for_status()
            results["webhook"] = "sent"
        except Exception as e:
            results["webhook"] = f"failed: {_shorten(e, 300)}"

    email_to = os.getenv("ERROR_NOTIFICATION_EMAIL_TO")
    smtp_host = os.getenv("SMTP_HOST")
    if email_to and smtp_host:
        try:
            msg = EmailMessage()
            msg["Subject"] = subject
            msg["From"] = os.getenv("SMTP_FROM") or os.getenv("SMTP_USERNAME") or "noreply@localhost"
            msg["To"] = email_to
            msg.set_content(json.dumps(payload, indent=2, default=str))

            port = int(os.getenv("SMTP_PORT", "587"))
            with smtplib.SMTP(smtp_host, port, timeout=ERROR_NOTIFICATION_TIMEOUT_SECONDS) as smtp:
                if os.getenv("SMTP_USE_TLS", "true").lower() == "true":
                    smtp.starttls()
                username = os.getenv("SMTP_USERNAME")
                password = os.getenv("SMTP_PASSWORD")
                if username and password:
                    smtp.login(username, password)
                smtp.send_message(msg)
            results["email"] = "sent"
        except Exception as e:
            results["email"] = f"failed: {_shorten(e, 300)}"

    if not results:
        results["notification"] = "skipped: configure ERROR_NOTIFICATION_WEBHOOK_URL or SMTP_HOST + ERROR_NOTIFICATION_EMAIL_TO"

    return {"status": True, "channels": results, "payload": payload}


@activity.defn
async def update_equipment_for_simulation_activity(params: dict):
    # ✅ Import inside the activity function
    from src.aeros_equipment.service import update_equipment_for_simulation
    from src.database.core import async_session, async_aeros_session
    
    
    async with async_session() as db_session, async_aeros_session() as aeros_db_session:
        return await update_equipment_for_simulation(
            db_session=db_session,
            aeros_db_session=aeros_db_session,
            project_name=params["projectName"],
            overhaul_duration=params["OverhaulDuration"],
            overhaul_interval=params["OverhaulInterval"],
            offset=params["OffSet"],
            schematic_name=params["SchematicName"],
            custom_input=params["CustomInput"],
            simulation_duration=params["SimDuration"]
        )


@activity.defn
async def get_simulation_engine_activity(params: dict) -> str:
    from src.aeros_simulation.ramp_engine import selected_simulation_engine

    return selected_simulation_engine(params)


@activity.defn
async def collect_rbd_reliability_for_ramp_activity(params: dict):
    from src.aeros_equipment.service import collect_rbd_reliability_for_ramp
    from src.database.core import async_session

    async with async_session() as db_session:
        return await collect_rbd_reliability_for_ramp(
            db_session=db_session,
            schematic_name=params["SchematicName"],
            custom_input=params.get("CustomInput"),
        )


@activity.defn
async def execute_simulation_activity(params: dict):
    from src.database.core import async_session
    from src.aeros_simulation.simulation_save_service import execute_simulation
    async with async_session() as db_session:
        return await execute_simulation(
            db_session=db_session,
            simulation_id=params["sim_data"]["HubCnnId"],
            sim_data=params["sim_data"],
            is_saved=True,
            eq_update=params["eq_update"]
        )


@activity.defn
async def execute_ramp_simulation_activity(params: dict):
    from src.database.core import async_session
    from src.aeros_simulation.ramp_engine import build_ramp_distribution_overrides, execute_ramp_simulation

    async with async_session() as db_session:
        return await execute_ramp_simulation(
            db_session=db_session,
            simulation_id=params["sim_data"]["HubCnnId"],
            sim_data=params["sim_data"],
            eq_update=params.get("eq_update") or build_ramp_distribution_overrides(params["sim_data"]),
        )


@activity.defn
async def execute_update_node(params:dict):
    from src.database.core import async_session
    from src.aeros_equipment.service import update_node
    async with async_session() as db_session:
        return await update_node(
            db_session=db_session,
            project_name=params["projectName"],
            equipment_nodes=params["equipment_nodes"]
        )

@activity.defn
async def calculate_plant_eaf_activity(params: dict):
    from src.aeros_simulation.simulation_save_service import calculate_plant_eaf
    from src.database.core import async_session
    
    async with async_session() as db_session:
        return await calculate_plant_eaf(
            db_session=db_session,
            simulation_id=params["HubCnnId"],
            po_downtime=params["OverhaulDuration"],
            oh_interval=params["OverhaulInterval"],
            offset=params["OffSet"],
            manual_deratings=params.get("deratings", []),
            manual_mo=params.get("maintenance_outages", [])
        )
        

@activity.defn
async def update_contribution_bulk_mappings_activity(sim_id: str):
    from src.aeros_contribution.service import update_contribution_bulk_mappings
    from src.database.core import async_session
    
    async with async_session() as db_session:
        return await update_contribution_bulk_mappings(
            db_session=db_session,
            simulation_id=sim_id,
        )
        

@activity.defn
async def call_callback_ahm(params):
    from src.aeros_simulation.service import call_ahm_callback_service
    sim_id = params["simulation_id"]
    job_id = params["job_id"]
    
    try:
        await call_ahm_callback_service(sim_id, job_id)
        return True
    except Exception as e:
        activity.logger.error(f"Callback to AHM failed: {str(e)}")
        raise e


@activity.defn
async def handle_workflow_failure(params: dict) -> dict:
    import src.aeros_equipment.model  # noqa: F401
    from src.aeros_simulation.model import AerosSimulation
    from src.database.core import async_session

    simulation_id = params.get("simulation_id")
    error = params.get("error", "Unknown workflow error")
    workflow_stage = params.get("stage", "workflow")

    async with async_session() as db_session:
        simulation = await db_session.get(AerosSimulation, UUID(str(simulation_id)))
        if simulation:
            simulation.status = "failed"
            simulation.error = {
                "stage": workflow_stage,
                "message": _shorten(error),
            }
            simulation.completed_at = datetime.now()
            await db_session.commit()

    notification = _send_error_notification(
        subject="RBD simulation workflow failed",
        message=error,
        context={
            "simulation_id": simulation_id,
            "stage": workflow_stage,
            "ahm_job_id": params.get("ahm_job_id"),
        },
    )
    return {"status": "failed_recorded", "notification": notification}
    
    
LOCK_KEY = "lock:aeros_process"
LOCK_TIMEOUT = int(os.getenv("SIMULATION_LOCK_TIMEOUT_SECONDS", "7200"))  # 2 hours

@activity.defn
async def acquire_lock(job_id: str) -> dict:
    try:
        while True:
            acquired = await redis_client.set(
                LOCK_KEY,
                job_id,
                nx=True,
                ex=LOCK_TIMEOUT
            )

            if acquired:
                # Lock acquired
                await redis_client.hset(
                    f"process:{job_id}",
                    mapping={
                        "job_id": job_id,
                        "status": "Processing"
                    }
                )
                await redis_client.expire(f"process:{job_id}", LOCK_TIMEOUT)

                return {
                    "job_id": job_id,
                    "status": "Ready"
                }

            # Someone else holds the lock
            await redis_client.hset(
                f"process:{job_id}",
                mapping={
                    "job_id": job_id,
                    "status": "Queue"
                }
            )
            await redis_client.expire(f"process:{job_id}", 300)

            # Cancellation-friendly wait
            await asyncio.sleep(3)


    except Exception:
        await redis_client.hset(
            f"process:{job_id}",
            mapping={
                "job_id": job_id,
                "status": "Failed"
            }
        )
        await redis_client.expire(f"process:{job_id}", 300)
        raise

@activity.defn
async def release_lock(job_id: str):
    await redis_client.delete(LOCK_KEY)

    await redis_client.hset(
        f"process:{job_id}",
        mapping={
            "job_id": job_id,
            "status": "Done"
        }
    )
    await redis_client.expire(f"process:{job_id}", 300)