from __future__ import annotations

import asyncio
import inspect
import logging
import time

from app.modules.agent.engine.orchestrator.execution_context import ExecutionContext
from app.modules.agent.engine.orchestrator.models import PlanStatus, PlanStep, StepResult, StepStatus
from app.modules.agent.engine.orchestrator.quality_gates import QualityGateRunner
from app.modules.agent.engine.orchestrator.step_registry import StepRegistry

LOGGER = logging.getLogger(__name__)


class ExecutionEngine:
    def __init__(self, step_registry: StepRegistry, gates: QualityGateRunner) -> None:
        self._steps = step_registry
        self._gates = gates

    async def run(self, ctx: ExecutionContext) -> list[StepResult]:
        ctx.plan.status = PlanStatus.RUNNING
        step_results: list[StepResult] = []

        for step in ctx.plan.steps:
            dep_issue = self._dependency_issue(step, step_results)
            if dep_issue:
                result = StepResult(
                        step_id=step.step_id,
                        status=StepStatus.SKIPPED,
                        warnings=[dep_issue],
                    )
                step_results.append(result)
                self._log_step_result(ctx, step, result)
                continue

            result = await self._run_with_retry(step, ctx)
            step_results.append(result)
            self._log_step_result(ctx, step, result)
            if result.status in {StepStatus.FAILED, StepStatus.RETRY_EXHAUSTED} and step.on_failure == "fail":
                ctx.plan.status = PlanStatus.FAILED
                return step_results

        passed, global_messages = self._gates.check_global(ctx.plan.global_gates, ctx)
        if not passed:
            step_results.append(
                StepResult(
                    step_id="global_gates",
                    status=StepStatus.FAILED,
                    warnings=global_messages,
                )
            )
            ctx.plan.status = PlanStatus.FAILED
            return step_results

        if any(item.status in {StepStatus.FAILED, StepStatus.RETRY_EXHAUSTED} for item in step_results):
            ctx.plan.status = PlanStatus.FAILED
        elif any(item.status == StepStatus.SKIPPED for item in step_results):
            ctx.plan.status = PlanStatus.PARTIAL
        else:
            ctx.plan.status = PlanStatus.COMPLETED
        return step_results

    async def _run_with_retry(self, step: PlanStep, ctx: ExecutionContext) -> StepResult:
        max_attempts = max(1, int(step.retry.max_attempts or 1))
        attempt = 0
        last_error: Exception | None = None

        while attempt < max_attempts:
            attempt += 1
            started_at = time.monotonic()
            LOGGER.warning(
                "orchestrator step start: task_id=%s step_id=%s action_id=%s executor=%s attempt=%s graph_id=%s",
                ctx.task.task_id,
                step.step_id,
                step.action_id,
                step.executor,
                attempt,
                step.graph_id or "",
            )
            await self._emit_progress(ctx, f"orchestrator.step.{step.step_id}", step.title)

            try:
                artifact_ids = await self._steps.execute(step, ctx)
                passed, gate_messages = self._gates.check_step(step, ctx)
                if not passed:
                    raise RuntimeError(";".join(gate_messages) or "step_quality_gate_failed")

                elapsed = int((time.monotonic() - started_at) * 1000)
                return StepResult(
                    step_id=step.step_id,
                    status=StepStatus.SUCCESS,
                    produced_artifact_ids=artifact_ids,
                    warnings=gate_messages,
                    duration_ms=elapsed,
                )
            except Exception as exc:
                last_error = exc
                if attempt < max_attempts and step.retry.backoff_sec > 0:
                    await asyncio.sleep(step.retry.backoff_sec)

        elapsed = int((time.monotonic() - started_at) * 1000)
        return StepResult(
            step_id=step.step_id,
            status=StepStatus.RETRY_EXHAUSTED if max_attempts > 1 else StepStatus.FAILED,
            error_code="step_execution_failed",
            error_message=str(last_error) if last_error else "step_execution_failed",
            duration_ms=elapsed,
        )

    def _dependency_issue(self, step: PlanStep, results: list[StepResult]) -> str | None:
        if not step.depends_on:
            return None
        by_step = {item.step_id: item for item in results}
        for dep in step.depends_on:
            dep_result = by_step.get(dep)
            if dep_result is None:
                return f"dependency_not_executed:{dep}"
            if dep_result.status != StepStatus.SUCCESS:
                return f"dependency_not_success:{dep}:{dep_result.status.value}"
        return None

    async def _emit_progress(self, ctx: ExecutionContext, stage: str, message: str) -> None:
        if ctx.progress_cb is None:
            return
        result = ctx.progress_cb(stage, message, "task_progress", {"layer": "orchestrator"})
        if inspect.isawaitable(result):
            await result

    def _log_step_result(self, ctx: ExecutionContext, step: PlanStep, result: StepResult) -> None:
        artifact_keys = []
        for artifact_id in result.produced_artifact_ids:
            item = next((artifact for artifact in ctx.artifacts.all_items() if artifact.artifact_id == artifact_id), None)
            if item is not None:
                artifact_keys.append(item.key)
        LOGGER.warning(
            "orchestrator step result: task_id=%s step_id=%s action_id=%s status=%s duration_ms=%s artifact_keys=%s warnings=%s error=%s",
            ctx.task.task_id,
            step.step_id,
            step.action_id,
            result.status.value,
            result.duration_ms,
            artifact_keys,
            result.warnings,
            result.error_message or "",
        )
