from datetime import timedelta
from temporalio import workflow
from temporalio.client import Client
from temporalio.common import RetryPolicy
import os
with workflow.unsafe.imports_passed_through():
    from temporal.activity import calculate_plant_eaf_activity, collect_rbd_reliability_for_ramp_activity, execute_ramp_simulation_activity, execute_simulation_activity, get_simulation_engine_activity, update_contribution_bulk_mappings_activity, update_equipment_for_simulation_activity,call_callback_ahm, acquire_lock, release_lock,execute_update_node, handle_workflow_failure


TEMPORAL_URL = os.environ.get("TEMPORAL_URL", "http://192.168.1.86:7233")
_RETRY_POLICY = RetryPolicy(
    initial_interval=timedelta(seconds=10),
    backoff_coefficient=2.0,
    maximum_interval=timedelta(minutes=2),
    maximum_attempts=3,
)
_NO_RETRY_POLICY = RetryPolicy(maximum_attempts=1)

def _workflow_error_message(error: Exception) -> str:
    messages = []
    current = error
    while current:
        message = str(current)
        if message and message not in messages:
            messages.append(message)
        current = getattr(current, "__cause__", None)
    return " | caused by: ".join(messages) or error.__class__.__name__

@workflow.defn
class SimulationWorkflow:
    def __init__(self):
        self.status = "INITIALIZED"
        self.isDone = False

    @workflow.query
    def get_status(self) -> str:
        return self.status
    
    @workflow.query
    def is_done(self) -> bool:
        return self.isDone
    
    @workflow.run
    async def run(self, sim_data: dict) -> str:
        self.status = "QUEUE"
        lock_acquired = False
        workflow_failed = False

        try:
            await workflow.execute_activity(
                acquire_lock,
                sim_data["HubCnnId"],
                start_to_close_timeout=timedelta(hours=2, minutes=5),
                retry_policy=_NO_RETRY_POLICY,
            )
            lock_acquired = True
            
            self.status = "RUNNING"

            engine = await workflow.execute_activity(
                get_simulation_engine_activity,
                {**sim_data},
                start_to_close_timeout=timedelta(seconds=30),
                retry_policy=_NO_RETRY_POLICY,
            )

            if engine == "aeros":
                self.status = "UPDATING_AEROS_EQUIPMENT"
                results = await workflow.execute_activity(
                    update_equipment_for_simulation_activity,
                    {**sim_data},
                    start_to_close_timeout=timedelta(minutes=30),
                    retry_policy=_RETRY_POLICY,
                )

                await workflow.execute_activity(
                    execute_update_node,
                    {"projectName": sim_data['projectName'], "equipment_nodes": results['equipment_nodes']},
                    start_to_close_timeout=timedelta(minutes=30),
                    retry_policy=_RETRY_POLICY,
                )

                self.status = "RUNNING_AEROS_SIMULATION"
                await workflow.execute_activity(
                    execute_simulation_activity,
                    {"sim_data": sim_data, "eq_update": results['reliability']},
                    start_to_close_timeout=timedelta(hours=2),
                    retry_policy=_NO_RETRY_POLICY,
                )

                await workflow.execute_activity(
                    calculate_plant_eaf_activity,
                    sim_data,
                    start_to_close_timeout=timedelta(minutes=30),
                    retry_policy=_RETRY_POLICY,
                )
            elif engine == "ramp":
                self.status = "BUILDING_RAMP_DISTRIBUTION_OVERRIDES"
                results = await workflow.execute_activity(
                    collect_rbd_reliability_for_ramp_activity,
                    {**sim_data},
                    start_to_close_timeout=timedelta(minutes=30),
                    retry_policy=_RETRY_POLICY,
                )

                self.status = "RUNNING_RAMP_SIMULATION"
                await workflow.execute_activity(
                    execute_ramp_simulation_activity,
                    {"sim_data": sim_data, "eq_update": results},
                    start_to_close_timeout=timedelta(hours=2),
                    retry_policy=_NO_RETRY_POLICY,
                )
            else:
                raise ValueError(f"Unsupported simulation engine: {engine}")
            
            await workflow.execute_activity(
                update_contribution_bulk_mappings_activity,
                sim_data["HubCnnId"],
                start_to_close_timeout=timedelta(minutes=30),
                retry_policy=_RETRY_POLICY,
            )

            if sim_data.get("AhmJobId"):
                await workflow.execute_activity(
                    call_callback_ahm,
                    {
                        "simulation_id": sim_data["HubCnnId"],
                        "job_id": sim_data["AhmJobId"]
                    },
                    start_to_close_timeout=timedelta(minutes=5),
                    retry_policy=_RETRY_POLICY,
                )

            self.status = "COMPLETED"
            self.isDone = True
            
            if sim_data.get("CallbackWorkflowId"):
                handle = workflow.get_external_workflow_handle(
                    sim_data["CallbackWorkflowId"]
                )
                await handle.signal("notify_done")

            return sim_data["HubCnnId"]

        except Exception as e:
            failed_stage = self.status
            self.status = "FAILED"
            workflow_failed = True
            await workflow.execute_activity(
                handle_workflow_failure,
                {
                    "simulation_id": sim_data["HubCnnId"],
                    "stage": failed_stage,
                    "error": _workflow_error_message(e),
                    "ahm_job_id": sim_data.get("AhmJobId"),
                },
                start_to_close_timeout=timedelta(seconds=30),
                retry_policy=_NO_RETRY_POLICY,
            )
            raise
        finally:
            if lock_acquired:
                try:
                    await workflow.execute_activity(
                        release_lock,
                        sim_data["HubCnnId"],
                        start_to_close_timeout=timedelta(minutes=10),
                        retry_policy=_RETRY_POLICY,
                    )
                except Exception as release_error:
                    await workflow.execute_activity(
                        handle_workflow_failure,
                        {
                            "simulation_id": sim_data["HubCnnId"],
                            "stage": "release_lock",
                            "error": f"Failed to release simulation lock: {release_error}",
                            "ahm_job_id": sim_data.get("AhmJobId"),
                        },
                        start_to_close_timeout=timedelta(seconds=30),
                        retry_policy=_NO_RETRY_POLICY,
                    )
                    if not workflow_failed:
                        raise