from __future__ import annotations

from datetime import timedelta
from temporalio import workflow

with workflow.unsafe.imports_passed_through():
    from ops_agent.temporal.activities import (
        run_ops_diagnose_activity,
        run_ops_remediate_activity
    )

@workflow.defn
class OpsRemediateWorkflow:
    """Orchestrates OpsAgent diagnosis, waits for human operator approval signal, and executes remediation."""

    def __init__(self) -> None:
        self._approved = False
        self._rejected = False

    @workflow.signal
    def approve(self) -> None:
        self._approved = True

    @workflow.signal
    def reject(self) -> None:
        self._rejected = True

    @workflow.run
    async def run(self, target_type: str, target_id: str) -> dict:
        # Step 1: Diagnose issues
        params = {"target_type": target_type, "target_id": target_id}
        diagnose_res = await workflow.execute_activity(
            run_ops_diagnose_activity,
            params,
            start_to_close_timeout=timedelta(minutes=5),
        )
        
        proposed_action = diagnose_res.get("proposed_action", "none")
        
        # Step 2: Check if remediation is proposed. If 'none', exit.
        if proposed_action == "none":
            return {
                "status": "completed",
                "diagnosis": diagnose_res["diagnosis"],
                "action_proposed": "none",
                "action_executed": "none",
                "result": "No problems detected. No action needed."
            }
            
        # Step 3: Approval Gate (wait for approval signal or reject signal or timeout)
        # Timeout after 24 hours (or shorter in test/production depending on requirements)
        approval_timeout = timedelta(hours=2)
        
        workflow.logger.info(
            f"Action proposed: {proposed_action}. Waiting for approval signal on workflow ID: {workflow.info().workflow_id}"
        )
        
        # Wait for approval condition
        try:
            is_condition_met = await workflow.wait_condition(
                lambda: self._approved or self._rejected,
                timeout=approval_timeout
            )
        except Exception:
            # wait_condition times out
            is_condition_met = False
            
        if not is_condition_met:
            return {
                "status": "timed_out",
                "diagnosis": diagnose_res["diagnosis"],
                "action_proposed": proposed_action,
                "action_executed": "none",
                "result": "Timed out waiting for administrator approval."
            }
            
        if self._rejected:
            return {
                "status": "rejected",
                "diagnosis": diagnose_res["diagnosis"],
                "action_proposed": proposed_action,
                "action_executed": "none",
                "result": "Remediation action rejected by administrator."
            }
            
        # Step 4: Execute remediation
        remed_params = {**diagnose_res, "is_approved": True}
        exec_res = await workflow.execute_activity(
            run_ops_remediate_activity,
            remed_params,
            start_to_close_timeout=timedelta(minutes=5),
        )
        
        return {
            "status": "remediated",
            "diagnosis": diagnose_res["diagnosis"],
            "action_proposed": proposed_action,
            "action_executed": proposed_action,
            "result": exec_res["action_result"]
        }
