"""Orchestrator — drives the FSM through stages sequentially."""

from __future__ import annotations

import logging
import time
import uuid
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from dpe.adapters.base import ModelAdapter
from dpe.adapters.registry import build_adapters, is_adapter_premium
from dpe.adapters.simulation import SimulationAdapter
from dpe.artifacts.registry import ArtifactRegistry
from dpe.core_bridge import emit_core_artifacts
from dpe.engine.budget import BudgetController, BudgetExceededError
from dpe.engine.prompts import STAGE_PROMPTS
from dpe.engine.states import STAGE_ORDER, Stage
from dpe.models.adapter import ModelRequest
from dpe.models.config import DPEConfig
from dpe.models.manifest import BudgetActuals, RunManifest, StageRecord
from dpe.validation import validate_prd_has_acceptance_criteria

logger = logging.getLogger(__name__)

# Graceful TelemetryEmitter import — telemetry must never crash the orchestrator
try:
    from abi_control_core.sdk.telemetry_emitter import TelemetryEmitter

    _HAS_TELEMETRY = True
except ImportError:
    _HAS_TELEMETRY = False
    TelemetryEmitter = None  # type: ignore[assignment,misc]

# Map DPE stage roles to TelemetryEmitter VALID_ROLES
_ROLE_MAP: dict[str, str] = {
    "normalizer": "extractor",
    "ideator_a": "planner",
    "ideator_b": "planner",
    "critic_a": "reviewer",
    "critic_b": "reviewer",
    "synthesizer": "summarizer",
    "prd_author": "builder",
    "prd_critic": "reviewer",
    "editor": "builder",
}

# Map adapter names to TelemetryEmitter VALID_PROVIDERS
_PROVIDER_MAP: dict[str, str] = {
    "claude_code_cli": "anthropic",
    "codex_cli": "openai",
    "local": "ollama",
    "simulation": "mock",
}


_SMOKE_RUN_ID = "000000000000"
_SMOKE_EPOCH = datetime(2026, 1, 1, 0, 0, 0, tzinfo=UTC)


class Orchestrator:
    """Top-level coordinator: config -> FSM -> adapters -> artifacts.

    Processes stages strictly in order. No recursion. No parallel premium calls.
    Every model call goes through BudgetController.
    """

    def __init__(
        self,
        config: DPEConfig,
        adapters: dict[str, ModelAdapter] | None = None,
    ) -> None:
        self._config = config
        self._budget = BudgetController(limits=config.budget)
        self._adapters = adapters if adapters is not None else build_adapters(config)
        self._simulation = SimulationAdapter()
        self._run_id = (
            _SMOKE_RUN_ID if config.smoke_mode else uuid.uuid4().hex[:12]
        )
        self._run_dir = Path(config.output_dir).resolve() / self._run_id
        self._registry = ArtifactRegistry(self._run_dir, clock=self._now)
        self._stage_outputs: dict[Stage, str] = {}
        self._stage_records: list[StageRecord] = []
        self._stop_reason = "completed"
        self._warnings: list[str] = []
        self._idea: str = ""

        # Initialize TelemetryEmitter if abi-core is available
        self._telemetry: Any = None
        if _HAS_TELEMETRY:
            try:
                self._telemetry = TelemetryEmitter(
                    output_dir=self._run_dir, run_id=self._run_id
                )
            except Exception:
                pass  # telemetry init must never crash

    def _now(self) -> datetime:
        """Return current time, or a fixed epoch in smoke mode."""
        if self._config.smoke_mode:
            return _SMOKE_EPOCH
        return datetime.now(UTC)

    @property
    def run_id(self) -> str:
        return self._run_id

    def run(self, idea: str) -> RunManifest:
        """Execute the full dialectical workflow. Returns the run manifest."""
        self._idea = idea
        logger.info("Starting run %s", self._run_id)

        try:
            for stage in STAGE_ORDER:
                if self._should_skip(stage):
                    self._record_skip(stage)
                    continue
                self._execute_stage(stage, idea)
        except BudgetExceededError as e:
            self._stop_reason = "budget_exceeded"
            self._warnings.append(f"Budget exceeded: {e.reason}")
            logger.warning("Run stopped: %s", e.reason)
        except Exception as e:
            self._stop_reason = f"error: {type(e).__name__}"
            self._warnings.append(str(e))
            logger.exception("Run failed")

        return self._finalize_manifest()

    def _should_skip(self, stage: Stage) -> bool:
        """Determine if a stage should be skipped based on config."""
        if stage in (Stage.CROSS_CRITIQUE_A, Stage.CROSS_CRITIQUE_B):
            return not self._config.enable_cross_critique

        # Gate ALL critique stages on can_critique()
        if stage == Stage.PRD_CRITIQUE_R1:
            return not self._budget.can_critique()

        if stage == Stage.PRD_REVISION_R1:
            return Stage.PRD_CRITIQUE_R1 not in self._stage_outputs

        if stage == Stage.PRD_CRITIQUE_R2:
            r1_done = Stage.PRD_REVISION_R1 in self._stage_outputs
            return not (r1_done and self._budget.can_critique())

        if stage == Stage.PRD_REVISION_R2:
            return Stage.PRD_CRITIQUE_R2 not in self._stage_outputs

        return False

    def _execute_stage(self, stage: Stage, idea: str) -> None:
        """Execute a single FSM stage."""
        started_at = self._now()
        adapter_name = self._resolve_adapter_name(stage)
        is_premium = is_adapter_premium(adapter_name, self._config)

        # Budget gate — must pass before every model call
        decision = self._budget.pre_call_check(
            is_premium=is_premium, stage=stage.value, adapter=adapter_name
        )

        if not decision.allowed:
            raise BudgetExceededError(decision.reason)

        # Resolve adapter AFTER budget check; degrade if needed
        if decision.degraded:
            adapter_name = "local"
            is_premium = False
            self._warnings.append(
                f"Degraded to local at stage {stage.value}"
            )

        adapter = self._resolve_adapter(adapter_name)

        # Build prompt
        prompt = self._build_prompt(stage, idea)

        # Execute model call
        request = ModelRequest(
            prompt=prompt,
            role=self._stage_role(stage),
            stage=stage.value,
        )

        t0 = time.monotonic()
        result = adapter.generate(request)
        duration_ms = int((time.monotonic() - t0) * 1000)

        # Emit telemetry (fire-and-forget, never crashes)
        self._emit_telemetry(
            adapter_name=adapter_name,
            stage_role=self._stage_role(stage),
            stage=stage.value,
            duration_ms=duration_ms,
            throttle_detected=result.throttle_detected,
        )

        # Record call in budget using resolved premium status
        self._budget.record_call(is_premium=is_premium)

        # Handle throttle detection from adapter
        if result.throttle_detected:
            self._warnings.append(
                f"Throttle detected at {stage.value} via {adapter_name}"
            )

        # Record critique rounds
        if stage in (Stage.PRD_CRITIQUE_R1, Stage.PRD_CRITIQUE_R2):
            self._budget.record_critique_round()

        # Write artifact before proceeding to next stage
        self._registry.write_stage_artifact(stage, result.text)
        self._stage_outputs[stage] = result.text

        # Validate acceptance criteria after FINALIZE stage
        if stage == Stage.FINALIZE:
            is_valid, error_msg = validate_prd_has_acceptance_criteria(result.text)
            if not is_valid:
                self._warnings.append(error_msg)
                logger.warning("PRD validation failed: %s", error_msg)

        finished_at = self._now()
        self._stage_records.append(
            StageRecord(
                stage=stage.value,
                adapter_used=result.adapter_name,
                started_at=started_at,
                finished_at=finished_at,
            )
        )
        logger.info(
            "Stage %s completed via %s", stage.value, result.adapter_name
        )

    def _record_skip(self, stage: Stage) -> None:
        """Record that a stage was skipped."""
        now = self._now()
        self._stage_records.append(
            StageRecord(
                stage=stage.value,
                adapter_used="none",
                started_at=now,
                finished_at=now,
                skipped=True,
            )
        )
        logger.info("Stage %s skipped", stage.value)

    def _resolve_adapter_name(self, stage: Stage) -> str:
        """Map a stage to its configured adapter name."""
        if self._config.simulation_mode:
            return "simulation"

        assignments = self._config.model_assignments
        mapping: dict[Stage, str] = {
            Stage.INPUT_NORMALIZE: assignments.synthesis,
            Stage.IDEATION_A: assignments.ideation_a,
            Stage.IDEATION_B: assignments.ideation_b,
            Stage.CROSS_CRITIQUE_A: assignments.critique,
            Stage.CROSS_CRITIQUE_B: assignments.critique,
            Stage.SYNTHESIS: assignments.synthesis,
            Stage.PRD_V1: assignments.prd_draft,
            Stage.PRD_CRITIQUE_R1: assignments.prd_critique,
            Stage.PRD_REVISION_R1: assignments.prd_revision,
            Stage.PRD_CRITIQUE_R2: assignments.prd_critique,
            Stage.PRD_REVISION_R2: assignments.prd_revision,
            Stage.FINALIZE: assignments.synthesis,
        }
        return mapping.get(stage, assignments.fallback)

    def _resolve_adapter(self, name: str) -> ModelAdapter:
        """Get the adapter instance for a given name."""
        if name == "simulation" or self._config.simulation_mode:
            return self._simulation
        if name in self._adapters:
            return self._adapters[name]
        # Fail loudly in non-simulation mode if adapter is missing
        msg = f"Adapter '{name}' not registered"
        available = list(self._adapters.keys())
        logger.error("%s (available: %s)", msg, available)
        raise RuntimeError(msg)

    def _build_prompt(self, stage: Stage, idea: str) -> str:
        """Build the prompt for a stage using templates and prior outputs."""
        template = STAGE_PROMPTS[stage]

        if stage == Stage.INPUT_NORMALIZE:
            return template.format(idea=idea)

        if stage == Stage.SYNTHESIS:
            return template.format(
                ideation_a=self._stage_outputs.get(
                    Stage.IDEATION_A, ""
                ),
                ideation_b=self._stage_outputs.get(
                    Stage.IDEATION_B, ""
                ),
                critique_a=self._stage_outputs.get(
                    Stage.CROSS_CRITIQUE_A, "(skipped)"
                ),
                critique_b=self._stage_outputs.get(
                    Stage.CROSS_CRITIQUE_B, "(skipped)"
                ),
            )

        if stage in (Stage.PRD_REVISION_R1, Stage.PRD_REVISION_R2):
            prd_key = (
                Stage.PRD_V1 if stage == Stage.PRD_REVISION_R1
                else Stage.PRD_REVISION_R1
            )
            critique_key = (
                Stage.PRD_CRITIQUE_R1 if stage == Stage.PRD_REVISION_R1
                else Stage.PRD_CRITIQUE_R2
            )
            return template.format(
                prd=self._stage_outputs.get(prd_key, ""),
                critique=self._stage_outputs.get(critique_key, ""),
            )

        # Default: use previous stage output as {input}
        prev_output = self._get_previous_output(stage)
        return template.format(input=prev_output)

    def _get_previous_output(self, stage: Stage) -> str:
        """Get the output of the most recent non-skipped prior stage."""
        idx = STAGE_ORDER.index(stage)
        for i in range(idx - 1, -1, -1):
            prev = STAGE_ORDER[i]
            if prev in self._stage_outputs:
                return self._stage_outputs[prev]
        return ""

    def _emit_telemetry(
        self,
        adapter_name: str,
        stage_role: str,
        stage: str,
        duration_ms: int,
        throttle_detected: bool = False,
    ) -> None:
        """Emit a telemetry record for a model call. Never raises."""
        if self._telemetry is None:
            return
        try:
            self._telemetry.emit(
                provider=_PROVIDER_MAP.get(adapter_name, "mock"),
                model=adapter_name,
                role=_ROLE_MAP.get(stage_role, "builder"),
                duration_ms=duration_ms,
                status="rate_limited" if throttle_detected else "success",
                task_id=stage,
                task_category="prd_generation",
                repo="abi-prd",
            )
        except Exception:
            pass  # telemetry must never crash the orchestrator

    def _stage_role(self, stage: Stage) -> str:
        """Map stage to a logical role name."""
        role_map: dict[Stage, str] = {
            Stage.INPUT_NORMALIZE: "normalizer",
            Stage.IDEATION_A: "ideator_a",
            Stage.IDEATION_B: "ideator_b",
            Stage.CROSS_CRITIQUE_A: "critic_a",
            Stage.CROSS_CRITIQUE_B: "critic_b",
            Stage.SYNTHESIS: "synthesizer",
            Stage.PRD_V1: "prd_author",
            Stage.PRD_CRITIQUE_R1: "prd_critic",
            Stage.PRD_REVISION_R1: "prd_author",
            Stage.PRD_CRITIQUE_R2: "prd_critic",
            Stage.PRD_REVISION_R2: "prd_author",
            Stage.FINALIZE: "editor",
        }
        return role_map.get(stage, "unknown")

    def _finalize_manifest(self) -> RunManifest:
        """Build, write, and return the run manifest. Always called, even on failure."""
        wall_time = (
            0.0 if self._config.smoke_mode
            else round(self._budget.wall_time_elapsed, 2)
        )
        manifest = RunManifest(
            run_id=self._run_id,
            timestamp=self._now(),
            config_hash=self._config.config_hash(),
            models_used={
                r.stage: r.adapter_used
                for r in self._stage_records if not r.skipped
            },
            budget_limits=self._config.budget.model_dump(),
            budget_actuals=BudgetActuals(
                total_calls_used=self._budget.total_calls_used,
                premium_calls_used=self._budget.premium_calls_used,
                wall_time_seconds=wall_time,
                critique_rounds_used=self._budget.critique_rounds_used,
            ),
            throttle_events=self._budget.throttle_events,
            stages=self._stage_records,
            artifacts=self._registry.artifacts,
            stop_reason=self._stop_reason,
            warnings=self._warnings,
        )

        self._registry.write_manifest(manifest.model_dump(mode="json"))
        self._registry.write_artifact_index()

        # Emit canonical abi-control-core artifacts (additive, non-fatal)
        try:
            emit_core_artifacts(
                manifest=manifest,
                run_dir=self._run_dir,
                config=self._config,
                stage_outputs=self._stage_outputs,
            )
        except Exception as e:
            logger.warning("Core artifact emission failed: %s", e)
            self._warnings.append(f"core_bridge: {e}")

        # Emit telemetry summary
        if self._telemetry is not None:
            try:
                logger.info("Telemetry: %s", self._telemetry.summary_line())
            except Exception:
                pass

        logger.info(
            "Run %s finished: stop_reason=%s, stages=%d, artifacts=%d",
            self._run_id,
            self._stop_reason,
            len(self._stage_records),
            len(self._registry.artifacts),
        )
        return manifest
