from __future__ import annotations

import asyncio
import json
import logging
import re
from datetime import datetime, timezone
from typing import Any, Awaitable, Callable

from google import genai

from epochler.agents import planner
from epochler.config import get_settings
from epochler.contract.models import (
    AgentActivity,
    ActivityType,
    ChatMessage,
    Conversation,
    ConversationState,
    DecisionRecord,
    DecisionStatus,
    DecisionType,
    ExperimentJournal,
    HardwareProfile,
)
from epochler.app_settings import get_app_settings, save_app_settings
from epochler.contract.engine import invalidate_downstream_phases
from epochler.decisioning import (
    apply_decision_action,
    apply_decision_proposal,
    refresh_conversation_state,
)
from epochler.runner.runner import ExperimentRunner
from epochler.storage import store

_IMPORTANT_DECISION_TYPES = frozenset({
    DecisionType.DATASET,
    DecisionType.PREPROCESSING,
    DecisionType.MODEL,
    DecisionType.EVALUATION,
    DecisionType.TARGET,
})

logger = logging.getLogger(__name__)


def _parse_dict_field(value: Any) -> dict[str, Any]:
    """Parse a field that may be a dict or a JSON string (from structured output)."""
    if isinstance(value, dict):
        return value
    if isinstance(value, str):
        if not value.strip():
            return {}
        try:
            parsed = json.loads(value)
            if isinstance(parsed, dict):
                return parsed
        except (json.JSONDecodeError, ValueError):
            pass
        return {"_raw": value}
    return {}

_TITLE_PROMPT = (
    "Generate a short, descriptive title (max 6 words) for this conversation. "
    "Reply with ONLY the title text, nothing else. No quotes."
)

_AUTO_ALL_PATTERNS = re.compile(
    r"(decide\s+everything|auto[- ]?decide\s+all|fully\s+autonomous|make\s+all(?:\s+the)?\s+decisions(?:\s+yourself)?(?:\s+from\s+now\s+on)?|you\s+decide\s+everything|handle\s+all(?:\s+the)?\s+decisions(?:\s+yourself)?)",
    re.IGNORECASE,
)
_GUIDED_PATTERNS = re.compile(
    r"(let\s+me\s+(?:decide|approve)\s+everything|i\s+(?:want\s+to\s+)?(?:decide|approve|review)\s+(?:everything|all)|guided\s+mode|ask\s+me\s+everything)",
    re.IGNORECASE,
)
_AUTO_MINOR_PATTERNS = re.compile(
    r"(semi[- ]?autonom|auto[- ]?decide\s+minor|only\s+ask\s+(?:for\s+)?(?:important|key|major))",
    re.IGNORECASE,
)


def _detect_autonomy_override(text: str) -> str | None:
    """Detect if a user message contains an autonomy level change request."""
    if _GUIDED_PATTERNS.search(text):
        return "guided"
    if _AUTO_ALL_PATTERNS.search(text):
        return "auto_all"
    if _AUTO_MINOR_PATTERNS.search(text):
        return "auto_minor"
    return None


def _sanitize_assistant_text(text: str) -> str:
    """Strip internal planner schema jargon from user-facing chat text."""
    cleaned = text
    cleaned = re.sub(
        r"(?:with|using)\s+`?auto_advance`?\s*:\s*true",
        "and approve them automatically",
        cleaned,
        flags=re.IGNORECASE,
    )
    cleaned = re.sub(
        r"`?auto_advance`?\s*:\s*true",
        "automatic approval",
        cleaned,
        flags=re.IGNORECASE,
    )
    cleaned = re.sub(
        r"`?(required_user_input|decision_ids|phase_impacts)`?",
        "internal planner fields",
        cleaned,
        flags=re.IGNORECASE,
    )
    cleaned = re.sub(r"\n{3,}", "\n\n", cleaned)
    return cleaned.strip()

# Type for the broadcast callback
BroadcastFn = Callable[[dict[str, Any]], Awaitable[None]]


class Orchestrator:
    """Central coordinator that routes user messages and manages agent interactions."""

    def __init__(self, broadcast: BroadcastFn) -> None:
        self._raw_broadcast = broadcast
        self._hardware: HardwareProfile | None = None
        self._running_tasks: dict[str, asyncio.Task] = {}
        self._chat_tasks: dict[str, asyncio.Task] = {}
        self._runners: dict[str, ExperimentRunner] = {}
        self._paused_convs: set[str] = set()
        self._stream_buffers: dict[str, str] = {}
        self._active_convs: set[str] = set()
        self._last_status: dict[str, str] = {}

    def get_stream_buffer(self, conv_id: str) -> str:
        return self._stream_buffers.get(conv_id, "")

    def track_conversation(self, conv_id: str) -> None:
        self._active_convs.add(conv_id)

    def untrack_conversation(self, conv_id: str) -> None:
        self._active_convs.discard(conv_id)
        self._last_status.pop(conv_id, None)

    async def start_heartbeat(self) -> None:
        """Periodically recompute and broadcast agent status for active conversations."""
        while True:
            await asyncio.sleep(2)
            for conv_id in list(self._active_convs):
                try:
                    conv = store.load_conversation(conv_id)
                    if not conv:
                        self._active_convs.discard(conv_id)
                        self._last_status.pop(conv_id, None)
                        continue
                    processing = self.is_processing(conv_id)
                    refresh_conversation_state(conv, is_processing=processing)
                    new_status = conv.readiness.agent_status
                    old_status = self._last_status.get(conv_id)
                    if new_status != old_status:
                        self._last_status[conv_id] = new_status
                        await self._broadcast({
                            "type": "status_update",
                            "conversation_id": conv_id,
                            "agent_status": conv.readiness.agent_status,
                            "agent_status_detail": conv.readiness.agent_status_detail,
                            "pending_decision_count": conv.readiness.pending_decision_count,
                        })
                except Exception:
                    logger.debug("Heartbeat error for %s", conv_id, exc_info=True)

    async def _broadcast(self, message: dict[str, Any]) -> None:
        """Fire-and-forget broadcast -- never raises, never blocks processing."""
        try:
            await self._raw_broadcast(message)
        except Exception:
            logger.debug("Broadcast failed (no connected clients?), continuing")

    def is_processing(self, conv_id: str) -> bool:
        task = self._chat_tasks.get(conv_id)
        return task is not None and not task.done()

    def set_hardware(self, hw: HardwareProfile) -> None:
        self._hardware = hw

    async def _emit_activity(
        self,
        contract_id: str,
        activity_type: ActivityType,
        summary: str,
        detail: str = "",
        experiment_id: str | None = None,
    ) -> None:
        activity = AgentActivity(
            activity_type=activity_type,
            summary=summary,
            detail=detail,
            experiment_id=experiment_id,
        )
        if contract_id:
            store.append_activity(contract_id, activity)
        await self._broadcast({
            "type": "activity",
            "data": activity.model_dump(mode="json"),
        })

    async def _emit_workspace(self, conversation: Conversation) -> None:
        processing = self.is_processing(conversation.id)
        refresh_conversation_state(conversation, is_processing=processing)
        self._last_status[conversation.id] = conversation.readiness.agent_status
        await self._broadcast({
            "type": "workspace_updated",
            "conversation_id": conversation.id,
            "state": conversation.state.value,
            "readiness": conversation.readiness.model_dump(mode="json"),
            "decisions": [d.model_dump(mode="json") for d in conversation.decisions],
            "contract": conversation.contract.model_dump(mode="json"),
        })

    @staticmethod
    def _decision_from_payload(payload: dict[str, Any]) -> DecisionRecord:
        decision_type = DecisionType(payload.get("type", DecisionType.GENERAL.value))
        return DecisionRecord(
            type=decision_type,
            title=str(payload.get("title", "")).strip(),
            status=DecisionStatus.NEEDS_INPUT
            if payload.get("required_user_input")
            else DecisionStatus.PROPOSED,
            summary=str(payload.get("summary", "")).strip(),
            rationale=str(payload.get("rationale", "")).strip(),
            evidence=[str(item) for item in payload.get("evidence", [])],
            before=_parse_dict_field(payload.get("before", {})),
            after=_parse_dict_field(payload.get("after", {})),
            blocking_reasons=[str(item) for item in payload.get("blocking_reasons", [])],
            downstream_impact=[str(item) for item in payload.get("downstream_impact", [])],
            required_user_input=[str(item) for item in payload.get("required_user_input", [])],
            phase_impacts=[str(item) for item in payload.get("phase_impacts", [])],
        )

    def _build_post_run_decision(self, conversation: Conversation, journal: dict[str, Any]) -> DecisionRecord:
        best_metric = journal.get("best_metric_value")
        best_id = journal.get("best_experiment_id")
        metric_name = conversation.contract.eval_spec.primary_metric if conversation.contract.eval_spec else "metric"
        target = conversation.contract.target_threshold
        met_target = bool(
            target is not None
            and best_metric is not None
            and (
                best_metric >= target
                if conversation.contract.threshold_direction == "higher_is_better"
                else best_metric <= target
            )
        )
        recommendation = (
            f"Target met on `{metric_name}` with value `{best_metric}`."
            if met_target
            else f"Best `{metric_name}` is `{best_metric}`. Continue iterating or revise the spec."
        )
        return DecisionRecord(
            type=DecisionType.POST_RUN_RECOMMENDATION,
            title="Post-run recommendation",
            status=DecisionStatus.PROPOSED,
            summary=recommendation,
            rationale="Every run should produce an explicit next-step recommendation instead of leaving the user at a dead end.",
            evidence=[
                f"best_experiment_id={best_id}",
                f"best_metric_value={best_metric}",
                f"target_threshold={target}",
            ],
            after={
                "best_experiment_id": best_id,
                "best_metric_value": best_metric,
                "target_met": met_target,
                "recommended_next_action": "finalize" if met_target else "iterate",
            },
            downstream_impact=[
                "Analysis is updated with the latest result.",
                "The readiness engine can now offer finalize or revise actions.",
            ],
            phase_impacts=["analysis"],
        )

    @staticmethod
    def _build_run_summary_message(
        conv: Conversation,
        journal: dict[str, Any],
        decision: DecisionRecord,
    ) -> ChatMessage:
        best_id = journal.get("best_experiment_id", "?")
        best_metric = journal.get("best_metric_value")
        total = len(journal.get("experiments", []))
        metric_name = conv.contract.eval_spec.primary_metric if conv.contract.eval_spec else "metric"
        target = conv.contract.target_threshold
        met = decision.after.get("target_met", False) if isinstance(decision.after, dict) else False

        parts = [f"## Pipeline Complete\n"]
        parts.append(f"Ran **{total} experiment(s)**. Best result: **{best_id}**\n")
        if best_metric is not None:
            parts.append(f"- **{metric_name}**: `{best_metric}`")
        if target is not None:
            parts.append(f"- **Target**: `{target}` — {'met' if met else 'not met'}")
        parts.append("")
        if met:
            parts.append("The target threshold has been reached. You can **approve** the recommendation below and click **Finalize** to wrap up, or keep iterating.")
        else:
            parts.append("The target was not reached. You can approve to finalize with current results, or reject to keep iterating.")

        return ChatMessage(
            role="assistant",
            content="\n".join(parts),
            metadata={"decision_ids": [decision.id]},
        )

    @staticmethod
    def _has_open_decisions(conversation: Conversation) -> bool:
        return any(
            decision.status in {
                DecisionStatus.PROPOSED,
                DecisionStatus.NEEDS_INPUT,
                DecisionStatus.BLOCKED,
            }
            for decision in conversation.decisions
        )

    async def _apply_planner_response(
        self,
        conv: Conversation,
        response: dict[str, Any],
        *,
        emit_stream: bool,
    ) -> tuple[ChatMessage, bool]:
        """Apply the planner response to the conversation.

        Returns (assistant_message, auto_advance).
        """
        decision_ids: list[str] = []
        decision_errors: list[str] = []
        raw_decisions = response.get("decisions", [])
        for payload in raw_decisions:
            try:
                decision = self._decision_from_payload(payload)
                apply_decision_proposal(conv, decision)
                decision_ids.append(decision.id)
            except Exception as exc:
                logger.exception("Invalid decision payload: %s", payload)
                decision_errors.append(f"Failed to apply '{payload.get('title', 'unknown')}': {exc}")

        planner_auto = bool(response.get("auto_advance", False))
        autonomy = conv.autonomy_level or get_app_settings().get("autonomy_level", "guided")

        has_important = any(
            DecisionType(p.get("type", "general")) in _IMPORTANT_DECISION_TYPES
            for p in raw_decisions
        )

        if autonomy == "auto_all":
            auto_advance = bool(raw_decisions)
        elif autonomy == "auto_minor":
            auto_advance = bool(raw_decisions) and not has_important
        else:
            auto_advance = planner_auto and bool(raw_decisions) and not has_important

        assistant_text = _sanitize_assistant_text(
            str(response.get("assistant_message", "")).strip() or "I updated the workspace."
        )
        if decision_errors:
            assistant_text += "\n\n---\n**Note:** " + "; ".join(decision_errors)
        meta: dict[str, Any] = {}
        if decision_ids:
            meta["decision_ids"] = decision_ids
        if auto_advance:
            meta["auto_advance"] = True
        assistant_msg = ChatMessage(
            role="assistant",
            content=assistant_text,
            metadata=meta,
        )
        conv.messages.append(assistant_msg)
        refresh_conversation_state(conv, is_processing=self.is_processing(conv.id))
        store.save_conversation(conv)

        if emit_stream:
            self._stream_buffers[conv.id] = assistant_text
            await self._broadcast({
                "type": "stream_start",
                "conversation_id": conv.id,
            })
            await self._broadcast({
                "type": "stream_chunk",
                "conversation_id": conv.id,
                "chunk": assistant_text,
            })
            self._stream_buffers.pop(conv.id, None)
            await self._broadcast({
                "type": "stream_end",
                "conversation_id": conv.id,
                "message": assistant_msg.model_dump(mode="json"),
                "state": conv.state.value,
                "decisions": [d.model_dump(mode="json") for d in conv.decisions],
                "readiness": conv.readiness.model_dump(mode="json") if conv.readiness else None,
            })
        return assistant_msg, auto_advance

    @staticmethod
    def _build_run_context(journal: ExperimentJournal) -> str:
        """Build a concise run-status summary for planner context."""
        exps = journal.experiments
        kept = sum(1 for e in exps if e.status.value == "keep")
        discarded = sum(1 for e in exps if e.status.value == "discard")
        crashed = sum(1 for e in exps if e.status.value == "crash")
        lines = [
            f"## Current Run Status",
            f"Experiments: {len(exps)} total ({kept} kept, {discarded} discarded, {crashed} crashed)",
        ]
        if journal.best_metric_value is not None:
            lines.append(f"Best metric: {journal.best_metric_value} (experiment {journal.best_experiment_id})")
        if journal.phase_history:
            last_phase = journal.phase_history[-1].get("phase", "unknown")
            lines.append(f"Latest phase: {last_phase}")
        recent = exps[-3:] if exps else []
        if recent:
            lines.append("Recent experiments:")
            for e in recent:
                metrics_str = ", ".join(f"{k}={v}" for k, v in (e.metrics or {}).items() if not k.startswith("_"))
                lines.append(f"  #{e.run_number} ({e.status.value}): {e.hypothesis[:80]} — {metrics_str}")
        if journal.agent_commentary:
            lines.append(f"Commentary: {journal.agent_commentary}")
        return "\n".join(lines)

    async def _trigger_planner(self, conv: Conversation) -> None:
        """Stream a planner response back to the frontend, with auto-advance loop."""
        conv_id = conv.id
        planner_conv = conv
        if conv.state.value in ("running", "paused", "reviewing_results"):
            contract_id = conv.contract.id if conv.contract else ""
            if contract_id:
                journal = store.load_journal(contract_id)
                if journal and journal.experiments:
                    ctx = self._build_run_context(journal)
                    from epochler.contract.models import ChatMessage as _CM
                    planner_conv = conv.model_copy(deep=True)
                    planner_conv.messages.append(_CM(role="system", content=ctx))

        self._stream_buffers[conv_id] = ""
        await self._broadcast({
            "type": "stream_start",
            "conversation_id": conv_id,
        })

        async def _stream_status(status: str) -> None:
            self._stream_buffers[conv_id] = self._stream_buffers.get(conv_id, "") + status
            await self._broadcast({
                "type": "stream_chunk",
                "conversation_id": conv_id,
                "chunk": status,
            })

        try:
            response = await planner.respond(planner_conv, self._hardware, on_status=_stream_status)
        except Exception as exc:
            logger.exception("Planner error during auto-trigger")
            error_text = f"I hit an error: {exc}. Could you try again or rephrase?"
            error_msg = ChatMessage(role="assistant", content=error_text)
            real_conv = store.load_conversation(conv_id) or conv
            real_conv.messages.append(error_msg)
            store.save_conversation(real_conv)
            self._stream_buffers.pop(conv_id, None)
            await self._broadcast({
                "type": "stream_end",
                "conversation_id": conv_id,
                "message": error_msg.model_dump(mode="json"),
            })
            return

        conv = store.load_conversation(conv_id) or conv

        num_decisions = len(response.get("decisions", []))
        if num_decisions > 0:
            await _stream_status(f"[decide] Applying {num_decisions} decision(s)...\n")

        assistant_msg, auto_advance = await self._apply_planner_response(conv, response, emit_stream=False)

        contract_id = conv.contract.id if conv.contract else ""
        summary = f"Proposed {num_decisions} decision(s)" if num_decisions else "Response generated"
        await self._emit_activity(contract_id, ActivityType.THINKING, summary)

        self._stream_buffers.pop(conv.id, None)
        await self._broadcast({
            "type": "stream_end",
            "conversation_id": conv.id,
            "message": assistant_msg.model_dump(mode="json"),
            "state": conv.state.value,
            "decisions": [d.model_dump(mode="json") for d in conv.decisions],
            "readiness": conv.readiness.model_dump(mode="json") if conv.readiness else None,
        })
        await self._emit_workspace(conv)

        asyncio.create_task(self._maybe_update_subtitle(conv))

        if auto_advance and num_decisions > 0:
            logger.info("Trigger planner: auto_advance=True, entering _auto_advance for %s", conv.id)
            await self._auto_advance(conv)

    async def _maybe_update_subtitle(self, conv: Conversation) -> None:
        """Generate/refresh the conversation subtitle after the 1st exchange, then every ~4 messages."""
        msg_count = len(conv.messages)
        if msg_count < 2:
            return
        if msg_count > 2 and (msg_count - 2) % 4 != 0:
            return

        settings = get_settings()
        if not settings.gemini_api_key:
            return

        snippet_parts: list[str] = []
        for m in conv.messages[-8:]:
            role = m.role
            text = m.content[:200]
            snippet_parts.append(f"{role}: {text}")
        snippet = "\n".join(snippet_parts)

        try:
            def _generate() -> str:
                client = genai.Client(api_key=settings.gemini_api_key)
                resp = client.models.generate_content(
                    model=settings.gemini_lite_model,
                    contents=f"{_TITLE_PROMPT}\n\n---\n\n{snippet}",
                    config=genai.types.GenerateContentConfig(temperature=0.3),
                )
                return resp.text.strip().strip('"').strip("'")

            subtitle = await asyncio.to_thread(_generate)
        except Exception:
            logger.debug("Subtitle generation failed", exc_info=True)
            return

        if not subtitle or len(subtitle) > 80:
            return

        conv_fresh = store.load_conversation(conv.id)
        if not conv_fresh:
            return
        conv_fresh.subtitle = subtitle
        store.save_conversation(conv_fresh)

        await self._broadcast({
            "type": "subtitle_updated",
            "conversation_id": conv.id,
            "subtitle": subtitle,
        })

    async def _maybe_inject_phase_prompt(self, conv: Conversation) -> None:
        """If readiness allows running a phase, inject an actionable message into the chat.

        Guards: skip if conversation is in a terminal/post-run state, or if
        the same phase_action was already the most recent assistant message.
        """
        terminal_states = {
            ConversationState.FAILED,
            ConversationState.COMPLETED,
            ConversationState.REVIEWING_RESULTS,
            ConversationState.STOPPED,
        }
        if conv.state in terminal_states:
            return

        r = conv.readiness
        if not r:
            return

        if r.can_run_preprocessing:
            text = "The approved spec is ready for preprocessing. You can start the preprocessing pipeline when ready."
            meta = {"phase_action": "start_preprocessing"}
        elif r.can_run_training:
            text = "All spec components are approved. You can start training and evaluation when ready."
            meta = {"phase_action": "start_training"}
        else:
            return

        action_key = meta["phase_action"]
        recent_assistant = [
            m for m in conv.messages[-10:]
            if m.role == "assistant" and m.metadata
        ]
        if any(m.metadata.get("phase_action") == action_key for m in recent_assistant):
            return

        msg = ChatMessage(role="assistant", content=text, metadata=meta)
        conv.messages.append(msg)
        await self._broadcast({
            "type": "message_saved",
            "conversation_id": conv.id,
            "message": msg.model_dump(mode="json"),
        })

    async def _auto_advance(self, conv: Conversation) -> None:
        """Auto-approve proposed decisions and re-trigger planner (silently) until input is needed."""
        MAX_AUTO_STEPS = 8
        seen_type_sets: list[frozenset[str]] = []

        for step in range(MAX_AUTO_STEPS):
            conv = store.load_conversation(conv.id) or conv

            proposed = [
                d for d in conv.decisions
                if d.status == DecisionStatus.PROPOSED
            ]
            if not proposed:
                break

            current_types = frozenset(d.type.value for d in proposed)
            seen_type_sets.append(current_types)

            if len(seen_type_sets) >= 3 and current_types in seen_type_sets[:-1]:
                logger.warning(
                    "Auto-advance: loop detected (seen %s before in history), stopping",
                    current_types,
                )
                for d in proposed:
                    try:
                        apply_decision_action(conv, d.id, "approve")
                    except Exception:
                        logger.exception("Auto-advance: failed to approve %s", d.id)
                store.save_conversation(conv)
                await self._emit_workspace(conv)
                break

            titles = ", ".join(d.title or d.type.value for d in proposed)
            logger.info("Auto-advance step %d: approving %s", step, titles)

            for d in proposed:
                try:
                    apply_decision_action(conv, d.id, "approve")
                except Exception:
                    logger.exception("Auto-advance: failed to approve %s", d.id)

            refresh_conversation_state(conv, is_processing=False)

            if conv.readiness.can_run_preprocessing or conv.readiness.can_run_training:
                store.save_conversation(conv)
                await self._emit_workspace(conv)
                logger.info("Auto-advance: readiness met, auto-starting run for %s", conv.id)
                result = await self.handle_run_control(conv.id, "start")
                if "error" in result:
                    logger.error("Auto-advance: auto-start failed: %s", result["error"])
                break

            store.save_conversation(conv)
            await self._emit_workspace(conv)

            if not conv.readiness.blocking_reasons:
                break

            try:
                response = await planner.respond(conv, self._hardware)
            except Exception as exc:
                logger.exception("Planner error during auto-advance")
                error_msg = ChatMessage(role="assistant", content=f"I hit an error: {exc}. Could you try again or rephrase?")
                conv.messages.append(error_msg)
                store.save_conversation(conv)
                await self._broadcast({
                    "type": "message_saved",
                    "conversation_id": conv.id,
                    "message": error_msg.model_dump(mode="json"),
                })
                break

            conv = store.load_conversation(conv.id) or conv
            assistant_msg, _ = await self._apply_planner_response(conv, response, emit_stream=True)
            await self._emit_workspace(conv)

    async def start_discovery(self, conv_id: str) -> None:
        """Autonomous workspace exploration triggered by the Start Discovery button."""
        conv = store.load_conversation(conv_id)
        if not conv:
            await self._broadcast({
                "type": "error",
                "detail": f"Conversation {conv_id} not found",
            })
            return

        discovery_msg = ChatMessage(
            role="user",
            content="Explore my workspace and tell me what you find. Then start proposing decisions for an ML pipeline.",
            metadata={"discovery": True},
        )
        conv.messages.append(discovery_msg)
        store.save_conversation(conv)

        await self._broadcast({
            "type": "message_saved",
            "conversation_id": conv_id,
            "message": discovery_msg.model_dump(mode="json"),
        })

        try:
            await self._trigger_planner(conv)
        except asyncio.CancelledError:
            logger.info("Discovery cancelled for %s", conv_id)
            self._chat_tasks.pop(conv_id, None)
            await self._broadcast({
                "type": "stream_end",
                "conversation_id": conv_id,
                "message": None,
            })
            return
        except Exception as e:
            logger.exception("Discovery planner error")
            self._chat_tasks.pop(conv_id, None)
            await self._broadcast({
                "type": "stream_end",
                "conversation_id": conv_id,
                "message": ChatMessage(role="assistant", content=f"Discovery error: {e}").model_dump(mode="json"),
                "state": ConversationState.FAILED.value,
            })
            return

        self._chat_tasks.pop(conv_id, None)

    def start_discovery_task(self, conv_id: str) -> None:
        existing = self._chat_tasks.get(conv_id)
        if existing and not existing.done():
            existing.cancel()
        task = asyncio.create_task(self.start_discovery(conv_id))
        self._chat_tasks[conv_id] = task

    async def handle_user_message(self, conv_id: str, content: str) -> None:
        conv = store.load_conversation(conv_id)
        if not conv:
            await self._broadcast({
                "type": "error",
                "detail": f"Conversation {conv_id} not found",
            })
            return

        user_msg = ChatMessage(role="user", content=content)
        conv.messages.append(user_msg)
        store.save_conversation(conv)

        await self._broadcast({
            "type": "message_saved",
            "conversation_id": conv_id,
            "message": user_msg.model_dump(mode="json"),
        })

        new_autonomy = _detect_autonomy_override(content)
        if new_autonomy:
            conv.autonomy_level = new_autonomy
            store.save_conversation(conv)
            label = {"auto_all": "fully autonomous", "auto_minor": "semi-autonomous", "guided": "guided"}.get(new_autonomy, new_autonomy)
            await self._broadcast({
                "type": "settings_updated",
                "conversation_id": conv.id,
                "autonomy_level": new_autonomy,
                "detail": f"Autonomy set to {label}.",
            })

        try:
            await self._trigger_planner(conv)
        except asyncio.CancelledError:
            logger.info("Chat generation cancelled for %s", conv_id)
            self._chat_tasks.pop(conv_id, None)
            self._stream_buffers.pop(conv_id, None)
            await self._broadcast({
                "type": "stream_end",
                "conversation_id": conv_id,
                "message": None,
            })
            return
        except Exception as e:  # pragma: no cover - defensive path
            logger.exception("Planner agent error")
            self._chat_tasks.pop(conv_id, None)
            self._stream_buffers.pop(conv_id, None)
            await self._broadcast({
                "type": "stream_end",
                "conversation_id": conv_id,
                "message": ChatMessage(role="assistant", content=f"Agent error: {e}").model_dump(mode="json"),
                "state": ConversationState.FAILED.value,
            })
            return

        self._chat_tasks.pop(conv_id, None)

    def stop_generation(self, conv_id: str) -> bool:
        task = self._chat_tasks.pop(conv_id, None)
        if task and not task.done():
            task.cancel()
            return True
        return False

    def start_chat_task(self, conv_id: str, content: str) -> None:
        existing = self._chat_tasks.get(conv_id)
        if existing and not existing.done():
            existing.cancel()
        task = asyncio.create_task(self.handle_user_message(conv_id, content))
        self._chat_tasks[conv_id] = task

    async def handle_decision_action(
        self,
        conv_id: str,
        decision_id: str,
        action: str,
        note: str = "",
    ) -> dict[str, Any]:
        conv = store.load_conversation(conv_id)
        if not conv:
            return {"error": "Conversation not found"}
        try:
            decision = apply_decision_action(conv, decision_id, action, note)
        except ValueError as exc:
            return {"error": str(exc)}

        if action == "provide_input" and note.strip():
            user_msg = ChatMessage(
                role="user",
                content=note.strip(),
                metadata={"decision_id": decision_id},
            )
            conv.messages.append(user_msg)
            store.save_conversation(conv)
            await self._emit_workspace(conv)
            await self._broadcast({
                "type": "message_saved",
                "conversation_id": conv_id,
                "message": user_msg.model_dump(mode="json"),
            })
            await self._trigger_planner(conv)
            return {
                "ok": True,
                "decision": decision.model_dump(mode="json"),
                "state": conv.state.value,
                "readiness": conv.readiness.model_dump(mode="json"),
            }

        if decision.type == DecisionType.POST_RUN_RECOMMENDATION:
            store.save_conversation(conv)
            await self._emit_workspace(conv)
            return {
                "ok": True,
                "decision": decision.model_dump(mode="json"),
                "state": conv.state.value,
                "readiness": conv.readiness.model_dump(mode="json") if conv.readiness else {},
            }

        auto_followup = False
        if action == "approve":
            refresh_conversation_state(conv, is_processing=self.is_processing(conv.id))
            remaining_open = [
                d for d in conv.decisions
                if d.status in (DecisionStatus.PROPOSED, DecisionStatus.NEEDS_INPUT)
            ]
            auto_followup = bool(
                conv.readiness.blocking_reasons and not remaining_open
            )

        if auto_followup:
            store.save_conversation(conv)
            await self._emit_workspace(conv)
            await self._trigger_planner(conv)
        else:
            if action == "approve" and conv.readiness:
                await self._maybe_inject_phase_prompt(conv)
            store.save_conversation(conv)
            await self._emit_workspace(conv)

        return {
            "ok": True,
            "decision": decision.model_dump(mode="json"),
            "state": conv.state.value,
            "readiness": conv.readiness.model_dump(mode="json") if conv.readiness else {},
        }

    async def handle_run_control(
        self, conv_id: str, action: str
    ) -> dict[str, Any]:
        conv = store.load_conversation(conv_id)
        if not conv:
            return {"error": "Conversation not found"}

        if action == "restart":
            if conv.state == ConversationState.RUNNING:
                existing_runner = self._runners.pop(conv_id, None)
                if existing_runner:
                    existing_runner.cancel()
                existing_task = self._running_tasks.pop(conv_id, None)
                if existing_task and not existing_task.done():
                    existing_task.cancel()
                    try:
                        await existing_task
                    except (asyncio.CancelledError, Exception):
                        pass

            invalidate_downstream_phases(conv.contract, conv.contract.phase_order[0], "Full pipeline restart requested")
            store.save_contract(conv.contract)
            store.clear_journal(conv.contract.id)
            conv.state = ConversationState.DRAFT
            logger.info("Restart: invalidated all phases and cleared journal for %s", conv_id)
            action = "start"

        if action == "start" and conv.state == ConversationState.COMPLETED:
            return {"error": "Cannot start from completed state. Create a new conversation."}

        if action == "start":
            if conv.state in {ConversationState.FAILED, ConversationState.STOPPED}:
                conv.state = ConversationState.DRAFT
            existing_task = self._running_tasks.get(conv_id)
            if existing_task and not existing_task.done():
                old_runner = self._runners.pop(conv_id, None)
                if old_runner:
                    old_runner.cancel()
                existing_task.cancel()
                try:
                    await existing_task
                except (asyncio.CancelledError, Exception):
                    pass
                self._running_tasks.pop(conv_id, None)

            refresh_conversation_state(conv)
            if not (conv.readiness.can_run_preprocessing or conv.readiness.can_run_training):
                return {"error": "The workspace is not ready to run yet."}
            conv.state = ConversationState.RUNNING
            store.save_conversation(conv)
            self._paused_convs.discard(conv_id)

            async def _check_paused():
                return conv_id in self._paused_convs

            runner = ExperimentRunner(
                conv.contract, self._broadcast, _check_paused,
                workspace_root=conv.workspace_root,
            )
            self._runners[conv_id] = runner

            async def _run_wrapper():
                try:
                    journal = await runner.run()
                    jdata = journal.model_dump(mode="json")
                    ran_experiments = len(journal.experiments) > 0
                    has_failed_phase = any(
                        entry.get("status") == "failed"
                        for entry in journal.phase_history
                    )
                    early_failure = not ran_experiments and has_failed_phase

                    c = store.load_conversation(conv_id)
                    if c:
                        c.contract = runner.contract

                        if early_failure:
                            failed_phases = [
                                entry.get("phase", "unknown")
                                for entry in journal.phase_history
                                if entry.get("status") == "failed"
                            ]
                            c.state = ConversationState.FAILED
                            error_detail = ", ".join(failed_phases)
                            fail_msg = ChatMessage(
                                role="assistant",
                                content=(
                                    f"## Pipeline Failed\n\n"
                                    f"The pipeline could not proceed — "
                                    f"**{error_detail}** phase(s) failed before any experiments ran.\n\n"
                                    f"Check the activity log for details. You may need to adjust "
                                    f"the dataset spec or retry."
                                ),
                            )
                            c.messages.append(fail_msg)
                            refresh_conversation_state(c, is_processing=False)
                            store.save_contract(c.contract)
                            store.save_conversation(c)
                            await self._emit_workspace(c)
                            await self._broadcast({
                                "type": "message_saved",
                                "conversation_id": conv_id,
                                "message": fail_msg.model_dump(mode="json"),
                            })
                            await self._broadcast({
                                "type": "run_error",
                                "conversation_id": conv_id,
                                "error": f"Pipeline failed during {error_detail}",
                            })
                        else:
                            c.state = ConversationState.REVIEWING_RESULTS
                            post_decision = self._build_post_run_decision(c, jdata)
                            apply_decision_proposal(c, post_decision)

                            summary_msg = self._build_run_summary_message(c, jdata, post_decision)
                            c.messages.append(summary_msg)

                            refresh_conversation_state(c, is_processing=self.is_processing(conv_id))
                            store.save_contract(c.contract)
                            store.save_conversation(c)
                            await self._emit_workspace(c)
                            await self._broadcast({
                                "type": "message_saved",
                                "conversation_id": conv_id,
                                "message": summary_msg.model_dump(mode="json"),
                            })
                            await self._broadcast({
                                "type": "run_completed",
                                "conversation_id": conv_id,
                                "journal": jdata,
                                "state": c.state.value,
                                "contract": c.contract.model_dump(mode="json"),
                                "decisions": [d.model_dump(mode="json") for d in c.decisions],
                                "readiness": c.readiness.model_dump(mode="json") if c.readiness else None,
                            })
                except asyncio.CancelledError:
                    # Persist the runner's live contract so a restart skips completed phases
                    try:
                        c = store.load_conversation(conv_id)
                        if c:
                            c.contract = runner.contract
                            c.state = ConversationState.STOPPED
                            store.save_contract(c.contract)
                            store.save_conversation(c)
                    except Exception:
                        logger.exception("Failed to persist contract on cancellation for %s", conv_id)
                    raise
                except Exception as e:
                    logger.exception("Runner failed for %s", conv_id)
                    failed = store.load_conversation(conv_id)
                    if failed:
                        failed.state = ConversationState.FAILED
                        refresh_conversation_state(failed, is_processing=False)
                        store.save_conversation(failed)
                        await self._emit_workspace(failed)
                    await self._broadcast({
                        "type": "run_error",
                        "conversation_id": conv_id,
                        "error": str(e),
                    })
                finally:
                    self._runners.pop(conv_id, None)
                    self._running_tasks.pop(conv_id, None)

            task = asyncio.create_task(_run_wrapper())
            self._running_tasks[conv_id] = task
            await self._broadcast({
                "type": "run_started",
                "conversation_id": conv_id,
                "contract_id": conv.contract.id,
            })
            return {"ok": True, "state": conv.state.value}

        elif action == "pause":
            if conv.state != ConversationState.RUNNING:
                return {"error": f"Cannot pause from state {conv.state.value}"}
            conv.state = ConversationState.PAUSED
            self._paused_convs.add(conv_id)
            store.save_conversation(conv)
            await self._broadcast({
                "type": "run_paused",
                "conversation_id": conv_id,
            })
            await self._emit_workspace(conv)
            return {"ok": True, "state": conv.state.value}

        elif action == "resume":
            if conv.state != ConversationState.PAUSED:
                return {"error": f"Cannot resume from state {conv.state.value}"}
            conv.state = ConversationState.RUNNING
            self._paused_convs.discard(conv_id)
            store.save_conversation(conv)
            await self._broadcast({
                "type": "run_resumed",
                "conversation_id": conv_id,
            })
            await self._emit_workspace(conv)
            return {"ok": True, "state": conv.state.value}

        elif action in {"stop", "finalize"}:
            if action == "finalize":
                if conv.state == ConversationState.COMPLETED:
                    return {"error": "Pipeline is already finalized."}
                if not conv.readiness or not conv.readiness.can_finalize:
                    return {"error": "Nothing is ready to finalize right now."}

            runner = self._runners.pop(conv_id, None)
            if runner:
                try:
                    conv.contract = runner.contract
                    store.save_contract(conv.contract)
                except Exception:
                    logger.exception("Failed to persist runner contract on stop for %s", conv_id)
                runner.cancel()

            if action == "finalize":
                for d in conv.decisions:
                    if d.type == DecisionType.POST_RUN_RECOMMENDATION and d.status == DecisionStatus.PROPOSED:
                        d.status = DecisionStatus.COMPLETED
                        d.updated_at = datetime.now(timezone.utc)

            conv.state = ConversationState.STOPPED if action == "stop" else ConversationState.COMPLETED
            refresh_conversation_state(conv, is_processing=False)
            store.save_conversation(conv)
            self._paused_convs.discard(conv_id)
            task = self._running_tasks.pop(conv_id, None)
            if task and not task.done():
                task.cancel()
            await self._broadcast({
                "type": "run_stopped" if action == "stop" else "run_finalized",
                "conversation_id": conv_id,
            })
            await self._emit_workspace(conv)
            return {"ok": True, "state": conv.state.value}

        return {"error": f"Unknown action: {action}"}
