from __future__ import annotations

import asyncio
import json
import logging
import re
import time as _time
from pathlib import Path
from typing import Any, Awaitable, Callable

import anthropic

from epochler.agents import searcher
from epochler.agents.router import Tier, classify as router_classify
from epochler.app_settings import get_app_settings, get_workspace_root
from epochler.config import get_settings
from epochler.contract.models import (
    Conversation,
    ConversationState,
    DecisionStatus,
    HardwareProfile,
    LockState,
    PhaseExecutionStatus,
)
from epochler.storage.store import workspace_path as store_workspace_path, workspaces_dir

logger = logging.getLogger(__name__)

_SYSTEM_PROMPT = (Path(__file__).parent / "prompts" / "planner.md").read_text(encoding="utf-8")

_TOOLS = [
    {
        "name": "inspect_path",
        "description": "Read a file or list directory contents to inspect datasets, reference repos, notebooks, etc.",
        "strict": True,
        "input_schema": {
            "type": "object",
            "properties": {
                "path": {"type": "string"},
                "max_lines": {"type": "integer"},
            },
            "required": ["path", "max_lines"],
            "additionalProperties": False,
        },
    },
    {
        "name": "web_search",
        "description": "Search the web for current ML knowledge when needed.",
        "strict": True,
        "input_schema": {
            "type": "object",
            "properties": {
                "query": {"type": "string"},
            },
            "required": ["query"],
            "additionalProperties": False,
        },
    },
]


_DATA_EXTENSIONS = frozenset({
    ".csv", ".tsv", ".json", ".jsonl", ".parquet", ".h5", ".hdf5",
    ".npy", ".npz", ".mat", ".pkl", ".pickle", ".feather", ".arrow",
    ".xlsx", ".xls", ".txt", ".dat", ".gz", ".zip", ".tar",
})


def _scan_workspace(ws_root: str | Path | None, max_entries: int = 80) -> str:
    """Return a concise tree of the workspace root (top 2 levels + data dirs)."""
    if not ws_root:
        return ""
    root = Path(ws_root)
    if not root.is_dir():
        return ""
    lines: list[str] = []
    try:
        for item in sorted(root.iterdir()):
            if item.name.startswith("."):
                continue
            if item.is_dir():
                lines.append(f"{item.name}/")
                try:
                    children = sorted(item.iterdir())[:20]
                    for child in children:
                        if child.name.startswith("."):
                            continue
                        suffix = "/" if child.is_dir() else ""
                        lines.append(f"  {child.name}{suffix}")
                except PermissionError:
                    pass
            else:
                size_kb = item.stat().st_size / 1024
                ext_note = " [DATA]" if item.suffix.lower() in _DATA_EXTENSIONS else ""
                lines.append(f"{item.name}  ({size_kb:.0f} KB){ext_note}")
            if len(lines) >= max_entries:
                lines.append("... (truncated)")
                break
    except PermissionError:
        return ""
    return "\n".join(lines)


def _build_system(conversation: Conversation, hardware: HardwareProfile | None) -> str:
    parts = [_SYSTEM_PROMPT]
    parts.append("\n\n## Current Spec\n")
    parts.append(f"```json\n{conversation.contract.model_dump_json(indent=2)}\n```")
    parts.append("\n\n## Current Decisions\n")
    parts.append(f"```json\n{json.dumps([d.model_dump(mode='json') for d in conversation.decisions[-10:]], indent=2)}\n```")

    pending_input = [
        d for d in conversation.decisions
        if d.status.value == "needs_input"
    ]
    if pending_input:
        titles = ", ".join(f'"{d.title}"' for d in pending_input)
        parts.append(
            f"\n\n## !! PENDING USER INPUT !!\n"
            f"The following decisions are waiting for user input: {titles}.\n"
            f"FOCUS on these first. If the latest user message answers any of them, "
            f"incorporate the answer and re-emit the decision as proposed (without required_user_input). "
            f"Do NOT create new unrelated decisions until these are resolved."
        )

    contract = conversation.contract
    dataset_frozen = bool(contract.dataset and contract.dataset.lock_state == LockState.FROZEN)
    preprocessing_frozen = bool(contract.preprocessing and contract.preprocessing.lock_state == LockState.FROZEN)
    model_frozen = bool(contract.model_spec and contract.model_spec.lock_state == LockState.FROZEN)
    eval_frozen = bool(contract.eval_spec and contract.eval_spec.lock_state == LockState.FROZEN)

    if not dataset_frozen:
        parts.append(
            "\n\n## !! PHASE ORDERING: DATASET FIRST !!\n"
            "CRITICAL: The dataset must be proposed and approved before any other "
            "component decisions (preprocessing, model, evaluation). "
            "Do NOT propose non-dataset decisions until a dataset decision is approved. "
            "If you have enough information to identify the dataset, propose it NOW."
        )
    elif not preprocessing_frozen:
        parts.append(
            "\n\n## Phase Ordering\n"
            "Dataset is locked. You should now propose preprocessing decisions.\n"
            "IMPORTANT: Propose the COMPLETE preprocessing plan in ONE decision. "
            "Include ALL steps (filtering, normalization, feature extraction, etc.) "
            "in the `after.transforms` array. Do NOT split into multiple decisions."
        )
    elif not model_frozen:
        parts.append(
            "\n\n## Phase Ordering\n"
            "Dataset and preprocessing are locked. Propose model architecture next."
        )
    elif not eval_frozen:
        parts.append(
            "\n\n## Phase Ordering\n"
            "Dataset, preprocessing, and model are locked. Propose evaluation spec next."
        )
    elif contract.target_threshold is None:
        parts.append(
            "\n\n## !! TARGET THRESHOLD REQUIRED !!\n"
            "Dataset, preprocessing, model, and evaluation are all locked. "
            "CRITICAL: You MUST propose a `target` decision NOW to set the success "
            "threshold and stop condition. The pipeline CANNOT run until `target_threshold` "
            "is set. Do NOT propose a `run_plan` until the target is approved.\n"
            "Example: `{\"type\": \"target\", \"title\": \"Target Threshold\", "
            "\"after\": \"{\\\"target_threshold\\\": 0.95, \\\"threshold_direction\\\": "
            "\\\"higher_is_better\\\"}\"}`"
        )
    else:
        has_approved_run_plan = any(
            d.type.value == "run_plan" and d.status.value == "approved"
            for d in conversation.decisions
        )
        if has_approved_run_plan:
            parts.append(
                "\n\n## !! RUN PLAN ALREADY APPROVED !!\n"
                "A run plan has already been approved. Do NOT propose another `run_plan`. "
                "The system will automatically start execution. If the user asks a question, "
                "respond conversationally without proposing new decisions."
            )
        else:
            parts.append(
                "\n\n## Phase Ordering\n"
                "All spec components are locked and the target is set. "
                "You may now propose a `run_plan` decision to start execution."
            )

    if conversation.state == ConversationState.COMPLETED:
        parts.append(
            "\n\n## Pipeline Completed\n"
            "The pipeline has finished successfully. Your role now is to help the user understand "
            "and act on the results. You should:\n"
            "- Summarize key metrics and outcomes when asked.\n"
            "- Generate comprehensive reports covering the full pipeline, experiments, and findings.\n"
            "- Discuss specific artifacts (charts, confusion matrices, feature importances) by name.\n"
            "- Suggest concrete next steps (iterate with different models, tune hyperparameters, "
            "try more data, etc.) and propose `post_run_recommendation` decisions if the user "
            "wants to iterate.\n"
            "- Use `inspect_path` to look at run outputs if you need specifics.\n"
            "CRITICAL: Do NOT refuse requests for reports or summaries. This is a core part of your job."
        )
    elif conversation.state == ConversationState.REVIEWING_RESULTS:
        parts.append(
            "\n\n## Reviewing Results\n"
            "The pipeline run has produced results that are ready for review. Help the user by:\n"
            "- Summarizing what happened during the run and the key metrics.\n"
            "- Explaining any charts or artifacts that were generated.\n"
            "- Recommending whether to iterate, adjust, or finalize.\n"
            "- Generating reports or summaries when asked.\n"
            "Use `inspect_path` on the run outputs to gather specifics."
        )
    elif conversation.state == ConversationState.FAILED:
        failed_phases = [
            name for name, phase in contract.phases.items()
            if phase.status == PhaseExecutionStatus.FAILED
        ]
        last_errors = [
            m.content for m in conversation.messages[-5:]
            if m.role == "assistant" and "failed" in m.content.lower()
        ]
        error_hint = last_errors[-1][:500] if last_errors else "Check activity log for details."
        parts.append(
            f"\n\n## !! PIPELINE FAILED: RECOVERY NEEDED !!\n"
            f"The pipeline is in a FAILED state. Failed phase(s): {', '.join(failed_phases) or 'unknown'}.\n"
            f"Context: {error_hint}\n\n"
            f"Your primary task is to help the user recover. Suggest concrete, actionable changes:\n"
            f"- If dataset acquisition failed, propose a new `dataset` decision with a different "
            f"`source_type` or `source_identifier`.\n"
            f"- If preprocessing or training failed, propose revised decisions for those phases.\n"
            f"- Always explain what went wrong in plain language before proposing a fix.\n"
            f"- Do NOT just say 'let me investigate'. Use your tools to find the actual error and "
            f"propose a specific fix."
        )
    elif conversation.state == ConversationState.STOPPED:
        parts.append(
            "\n\n## Pipeline Stopped\n"
            "The user manually stopped the pipeline. If they want to resume or restart, "
            "help them adjust decisions as needed and confirm they are ready to go."
        )

    autonomy = conversation.autonomy_level or get_app_settings().get("autonomy_level", "guided")
    if autonomy == "auto_all":
        parts.append(
            "\n\n## User Autonomy Preference\n"
            "The user wants fully autonomous decisions. Mark confident decisions for automatic "
            "approval. In your chat message, speak plainly to the user and never mention internal "
            "schema field names. Only ask the user for input when you truly cannot determine the answer."
        )
    elif autonomy == "auto_minor":
        parts.append(
            "\n\n## User Autonomy Preference\n"
            "The user wants semi-autonomous mode: auto-decide minor details (preprocessing, "
            "evaluation specifics) but ask for key choices (dataset, model architecture, target). "
            "Mark confident minor decisions for automatic approval, and never mention internal "
            "schema field names in your chat message."
        )

    if hardware:
        parts.append("\n\n## Hardware Profile\n")
        parts.append(f"```json\n{hardware.model_dump_json(indent=2)}\n```")
    project_ws = conversation.workspace_root or None
    ws_root = project_ws or get_workspace_root()
    ws_dir = workspaces_dir(project_ws)
    parts.append("\n\n## Workspace\n")
    if ws_root:
        parts.append(f"User project root (READ-ONLY): {ws_root}")
    parts.append(f"Execution workspace (code runs here): {ws_dir}")
    parts.append(
        "\nIMPORTANT: Never edit files in the user's project. Only read them. "
        "All generated scripts run in the execution workspace."
    )

    ws_listing = _scan_workspace(ws_root)
    if ws_listing:
        parts.append("\n\n## Live Workspace Contents\n")
        parts.append("Current files/folders at the project root (auto-scanned):\n")
        parts.append(f"```\n{ws_listing}\n```")
        parts.append(
            "\nIf you see data files here that match the pipeline's needs, use them "
            "directly instead of acquiring new data."
        )

    from datetime import date as _date
    parts.append(f"\n\n## Today's Date\n{_date.today().strftime('%A, %B %d, %Y')}")

    return "\n".join(parts)


_RECENT_FULL = 12
_CONTEXT_WINDOW = 30
_TRUNCATE_LEN = 300


def _conversation_to_messages(conv: Conversation) -> list[dict[str, Any]]:
    """Build the message list for the LLM.

    Last 12 messages in full; messages 13-30 truncated to 300 chars for context.
    Merges consecutive same-role messages (Anthropic requires alternating roles).
    Ensures the list always ends with a user message.
    """
    raw: list[dict[str, Any]] = []
    recent = conv.messages[-_CONTEXT_WINDOW:]
    full_start = max(0, len(recent) - _RECENT_FULL)
    for i, message in enumerate(recent):
        if message.role not in {"user", "assistant"}:
            continue
        content = message.content
        if i < full_start and len(content) > _TRUNCATE_LEN:
            content = content[:_TRUNCATE_LEN] + "..."
        raw.append({"role": message.role, "content": content})

    msgs: list[dict[str, Any]] = []
    for entry in raw:
        if msgs and msgs[-1]["role"] == entry["role"]:
            msgs[-1]["content"] += "\n\n" + entry["content"]
        else:
            msgs.append(entry)

    if msgs and msgs[-1]["role"] == "assistant":
        msgs.append({"role": "user", "content": "Continue."})
    return msgs


_BINARY_EXTENSIONS = frozenset({
    ".mat", ".pkl", ".pickle", ".npy", ".npz", ".h5", ".hdf5",
    ".pt", ".pth", ".bin", ".onnx", ".pb", ".tflite",
    ".zip", ".tar", ".gz", ".bz2", ".xz", ".7z",
    ".png", ".jpg", ".jpeg", ".gif", ".bmp", ".tiff",
    ".mp3", ".wav", ".mp4", ".avi", ".mov",
    ".exe", ".dll", ".so", ".dylib", ".o",
    ".parquet", ".arrow", ".feather",
})
_INSPECT_CHAR_CAP = 8000


def _inspect_path(path: str, max_lines: int = 200, *, allowed_roots: list[Path] | None = None) -> str:
    if not allowed_roots:
        return "Access denied: no workspace root configured. File inspection is disabled."
    try:
        raw = Path(path).expanduser()
        if not raw.is_absolute():
            p = (allowed_roots[0] / raw).resolve()
        else:
            p = raw.resolve()
    except (TypeError, ValueError) as exc:
        return f"Invalid path: {exc}"
    is_allowed = False
    for root in allowed_roots:
        root_resolved = root.resolve()
        if p == root_resolved or root_resolved in p.parents:
            is_allowed = True
            break
    if not is_allowed:
        roots_str = ", ".join(str(r) for r in allowed_roots)
        return f"Access denied: path '{path}' is outside the allowed workspace(s) ({roots_str}). Only paths within the project workspace can be inspected."
    if not p.exists():
        return f"Path does not exist: {path}"
    if p.is_dir():
        entries = []
        try:
            for item in sorted(p.iterdir()):
                try:
                    kind = "dir" if item.is_dir() else "file"
                    size = item.stat().st_size if item.is_file() else 0
                    entries.append(f"{kind}\t{size}\t{item.name}")
                except OSError:
                    entries.append(f"error\t0\t{item.name}")
        except PermissionError:
            return f"Permission denied reading directory: {path}"
        except OSError as exc:
            return f"Error reading directory: {exc}"
        return "\n".join(entries[:100])
    if p.suffix.lower() in _BINARY_EXTENSIONS:
        try:
            size_mb = p.stat().st_size / (1024 * 1024)
        except OSError:
            size_mb = 0.0
        return f"Binary file: {p.name} ({size_mb:.1f} MB, format: {p.suffix}). Use a specialized loader to read this."
    try:
        text = p.read_text(encoding="utf-8", errors="replace")
    except Exception as exc:
        return f"Could not read {p.name}: {exc}"
    lines = text.splitlines()
    result = "\n".join(lines[:max_lines])
    if len(result) > _INSPECT_CHAR_CAP:
        result = result[:_INSPECT_CHAR_CAP] + f"\n... (truncated, file has {len(lines)} lines total)"
    return result


_SEARCH_CHAR_CAP = 4000


async def _handle_tool_call(
    tool_name: str,
    tool_input: dict[str, Any],
    *,
    allowed_roots: list[Path] | None = None,
) -> str:
    try:
        if tool_name == "inspect_path":
            path = tool_input.get("path")
            if not path or not isinstance(path, str):
                return "Error: 'path' argument is required and must be a string."
            try:
                max_lines = int(tool_input.get("max_lines", 200))
            except (ValueError, TypeError):
                max_lines = 200
            return _inspect_path(path, max_lines, allowed_roots=allowed_roots)
        if tool_name == "web_search":
            query = tool_input.get("query")
            if not query or not isinstance(query, str):
                return "Error: 'query' argument is required and must be a string."
            result = await searcher.search(query)
            text = json.dumps(result, indent=2)
            if len(text) > _SEARCH_CHAR_CAP:
                text = text[:_SEARCH_CHAR_CAP] + "\n... (search results truncated)"
            return text
        return f"Unknown tool: {tool_name}"
    except Exception as exc:
        logger.exception("Tool call error: %s", tool_name)
        return f"Tool call failed: {exc}"


def _extract_json(payload: str) -> dict[str, Any]:
    payload = payload.strip()
    try:
        return json.loads(payload)
    except json.JSONDecodeError:
        match = re.search(r"\{.*\}", payload, flags=re.S)
        if not match:
            raise
        return json.loads(match.group(0))


def _error_response(reason: str) -> dict[str, Any]:
    """Transparent error -- no decisions, no amnesia. User stays in control."""
    return {
        "assistant_message": f"I hit an error: {reason}. Could you try again or rephrase?",
        "decisions": [],
    }


def _sync_gemini_respond(system_prompt: str, messages: list[dict[str, Any]]) -> str:
    """Call Gemini Flash synchronously (for asyncio.to_thread)."""
    from google import genai

    settings = get_settings()
    client = genai.Client(api_key=settings.gemini_api_key)
    parts: list[str] = [system_prompt, "\n\n---\n\nConversation:\n"]
    for msg in messages:
        parts.append(f"{msg['role']}: {msg['content']}\n")
    parts.append(
        "\nRespond with a single JSON object: "
        '{"assistant_message": "...", "decisions": [...]}\n'
        "decisions array can be empty if you have no proposals."
    )
    response = client.models.generate_content(
        model=settings.gemini_model,
        contents="".join(parts),
        config=genai.types.GenerateContentConfig(temperature=0.3),
    )
    return response.text.strip()


async def _respond_light(
    conversation: Conversation,
    hardware: HardwareProfile | None,
) -> dict[str, Any]:
    """Handle a light turn via Gemini Flash. Never returns decisions."""
    settings = get_settings()
    if not settings.gemini_api_key:
        return _error_response("Gemini API key not configured")

    system_prompt = _build_system(conversation, hardware)
    messages = _conversation_to_messages(conversation)
    if not messages:
        return {"assistant_message": "Describe your ML task or point me to a dataset.", "decisions": []}

    last_error = ""
    for attempt in range(2):
        try:
            raw = await asyncio.wait_for(
                asyncio.to_thread(_sync_gemini_respond, system_prompt, messages),
                timeout=60,
            )
            parsed = _extract_json(raw)
            parsed.setdefault("assistant_message", "")
            parsed.setdefault("decisions", [])
            if parsed["decisions"]:
                logger.info("LIGHT path returned decisions -- upgrading to HEAVY")
                return {"_upgrade_to_heavy": True}
            return parsed
        except Exception as exc:
            last_error = str(exc)
            logger.warning("Gemini Flash attempt %d failed: %s", attempt + 1, last_error)
            if attempt == 0:
                continue

    logger.error("Gemini Flash planner failed after retry: %s", last_error)
    return _error_response(f"Gemini Flash failed: {last_error}")


StatusCallback = Callable[[str], Awaitable[None]]


def _tool_status(tool_name: str, tool_input: dict[str, Any]) -> str:
    if tool_name == "inspect_path":
        p = tool_input.get("path", "...")
        short = p.rsplit("/", 2)
        label = "/".join(short[-2:]) if len(short) > 1 else p
        return f"[explore] Reading {label}"
    if tool_name == "web_search":
        q = tool_input.get("query", "...")
        return f"[search] Looking up: {q}"
    return f"[tool] {tool_name}"


def _tool_done_status(tool_name: str) -> str:
    if tool_name == "inspect_path":
        return "[explore] Done reading file"
    if tool_name == "web_search":
        return "[search] Got search results"
    return f"[tool] Finished {tool_name}"


_TEXT_STREAM_PREFIX = "\x02"

_MAX_TOOL_ITERATIONS = 12
_TOKEN_BUDGET = 150_000
_COMPACT_KEEP_RECENT = 4
_COMPACT_PREVIEW_CHARS = 200
_STREAM_TIMEOUT = 180

_PLANNER_SCHEMA = {
    "type": "object",
    "properties": {
        "assistant_message": {"type": "string"},
        "decisions": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "type": {"type": "string"},
                    "title": {"type": "string"},
                    "summary": {"type": "string"},
                    "rationale": {"type": "string"},
                    "evidence": {"type": "array", "items": {"type": "string"}},
                    "before": {"type": "string"},
                    "after": {"type": "string"},
                    "blocking_reasons": {"type": "array", "items": {"type": "string"}},
                    "downstream_impact": {"type": "array", "items": {"type": "string"}},
                    "required_user_input": {"type": "array", "items": {"type": "string"}},
                    "phase_impacts": {"type": "array", "items": {"type": "string"}},
                },
                "required": ["type", "title", "summary", "after"],
                "additionalProperties": False,
            },
        },
        "auto_advance": {"type": "boolean"},
    },
    "required": ["assistant_message", "decisions"],
    "additionalProperties": False,
}


def _estimate_tokens(system_prompt: str, messages: list[dict[str, Any]]) -> int:
    """Fast char-based token estimate (~4 chars per token)."""
    total = len(system_prompt) // 4
    for msg in messages:
        content = msg.get("content", "")
        if isinstance(content, str):
            total += len(content) // 4
        elif isinstance(content, list):
            for block in content:
                if isinstance(block, dict):
                    raw = block.get("content", "") or block.get("text", "") or ""
                    if isinstance(raw, str):
                        total += len(raw) // 4
                    inp = block.get("input")
                    if inp:
                        total += len(json.dumps(inp)) // 4
    return total


def _compact_old_tool_results(messages: list[dict[str, Any]]) -> None:
    """Shrink tool_result content for older messages, keeping recent ones full."""
    tool_result_indices: list[int] = []
    for i, msg in enumerate(messages):
        content = msg.get("content")
        if isinstance(content, list):
            for block in content:
                if isinstance(block, dict) and block.get("type") == "tool_result":
                    tool_result_indices.append(i)
                    break
    to_compact = tool_result_indices[:-_COMPACT_KEEP_RECENT] if len(tool_result_indices) > _COMPACT_KEEP_RECENT else []
    for idx in to_compact:
        content = messages[idx]["content"]
        if not isinstance(content, list):
            continue
        for block in content:
            if block.get("type") != "tool_result":
                continue
            full = block.get("content", "")
            if isinstance(full, str) and len(full) > _COMPACT_PREVIEW_CHARS * 2:
                preview = full[:_COMPACT_PREVIEW_CHARS]
                lines_count = full.count("\n") + 1
                block["content"] = f"{preview}\n... (compacted, ~{lines_count} lines originally)"


def _compute_allowed_roots(conversation: Conversation) -> list[Path]:
    """Build the list of directories the agent is allowed to inspect."""
    roots: list[Path] = []
    project_ws = conversation.workspace_root or None
    effective_root = Path(project_ws) if project_ws else get_workspace_root()
    if effective_root:
        roots.append(effective_root.resolve())
    contract_id = conversation.contract.id if conversation.contract else None
    if contract_id:
        exec_ws = store_workspace_path(contract_id, project_ws)
        exec_resolved = exec_ws.resolve()
        if not any(exec_resolved == r or r in exec_resolved.parents for r in roots):
            roots.append(exec_resolved)
    return roots


async def _respond_heavy(
    conversation: Conversation,
    hardware: HardwareProfile | None,
    on_status: StatusCallback | None = None,
) -> dict[str, Any]:
    """Handle a heavy turn via Claude Sonnet with streaming and tool use."""
    settings = get_settings()
    if not settings.anthropic_api_key:
        return _error_response("Anthropic API key not configured")

    allowed_roots = _compute_allowed_roots(conversation)

    async with anthropic.AsyncAnthropic(
        api_key=settings.anthropic_api_key,
        timeout=anthropic.Timeout(timeout=120.0, connect=10.0),
        max_retries=0,
    ) as client:
        system_prompt = _build_system(conversation, hardware)
        base_messages = _conversation_to_messages(conversation)

        if not base_messages:
            return {"assistant_message": "Describe your ML task or point me to a dataset.", "decisions": []}

        model = settings.anthropic_planner_model
        last_error = ""
        for attempt in range(2):
            messages = [dict(m) for m in base_messages]
            try:
                return await _run_heavy_loop(
                    client, model, system_prompt, messages,
                    on_status=on_status, attempt=attempt,
                    allowed_roots=allowed_roots,
                )
            except Exception as exc:
                last_error = str(exc)
                logger.warning("Planner (Sonnet) attempt %d failed: %s", attempt + 1, last_error, exc_info=True)
                if attempt == 0:
                    if on_status:
                        await on_status("[think] Retrying after error...\n")
                    continue

        logger.error("Planner (Sonnet) failed after retry: %s", last_error)
        return _error_response(f"Planner failed: {last_error}")


async def _run_heavy_loop(
    client: anthropic.AsyncAnthropic,
    model: str,
    system_prompt: str,
    messages: list[dict[str, Any]],
    on_status: StatusCallback | None,
    attempt: int,
    allowed_roots: list[Path] | None = None,
) -> dict[str, Any]:
    """Inner tool-use loop for the heavy path. Raises on failure (caller retries)."""
    iteration = 0
    while True:
        iteration += 1
        assistant_content: list[dict[str, Any]] = []
        text_chunks: list[str] = []
        tool_uses: list[Any] = []
        current_tool_input = ""
        current_tool_id = ""
        current_tool_name = ""
        force_final = iteration > _MAX_TOOL_ITERATIONS

        if on_status and iteration == 1 and attempt == 0:
            await on_status("[think] Analyzing your request...\n")
        elif on_status and iteration > 1:
            await on_status("[think] Thinking further...\n")

        est = _estimate_tokens(system_prompt, messages)
        if est > _TOKEN_BUDGET:
            _compact_old_tool_results(messages)
            logger.info("Token budget exceeded (~%dk), compacted old tool results", est // 1000)

        _output_config = {
            "format": {
                "type": "json_schema",
                "schema": _PLANNER_SCHEMA,
            },
        }

        stream_kwargs: dict[str, Any] = {
            "model": model,
            "max_tokens": 4096,
            "system": system_prompt,
            "messages": messages,
            "output_config": _output_config,
        }
        if force_final:
            if on_status:
                await on_status("[think] Summarizing findings...\n")
        else:
            stream_kwargs["tools"] = _TOOLS

        is_text_response = False
        t0 = _time.monotonic()

        async def _consume_stream() -> None:
            nonlocal is_text_response, current_tool_id, current_tool_name, current_tool_input
            async with client.messages.stream(**stream_kwargs) as stream:
                async for event in stream:
                    if event.type == "content_block_start":
                        block = event.content_block
                        if block.type == "tool_use":
                            current_tool_id = block.id
                            current_tool_name = block.name
                            current_tool_input = ""
                        elif block.type == "text":
                            is_text_response = True
                    elif event.type == "content_block_delta":
                        delta = event.delta
                        if delta.type == "text_delta":
                            text_chunks.append(delta.text)
                            if is_text_response and on_status:
                                await on_status(f"{_TEXT_STREAM_PREFIX}{delta.text}")
                        elif delta.type == "input_json_delta":
                            current_tool_input += delta.partial_json
                    elif event.type == "content_block_stop":
                        if current_tool_name:
                            try:
                                tool_input_parsed = json.loads(current_tool_input) if current_tool_input else {}
                            except (json.JSONDecodeError, ValueError):
                                logger.warning("Failed to parse tool input JSON: %s", current_tool_input[:200])
                                tool_input_parsed = {}
                            tool_uses.append(type("ToolUse", (), {
                                "id": current_tool_id,
                                "name": current_tool_name,
                                "input": tool_input_parsed,
                            })())
                            assistant_content.append({
                                "type": "tool_use",
                                "id": current_tool_id,
                                "name": current_tool_name,
                                "input": tool_input_parsed,
                            })
                            if on_status:
                                await on_status(_tool_status(current_tool_name, tool_input_parsed) + "\n")
                            current_tool_name = ""
                            current_tool_id = ""
                            current_tool_input = ""
                    elif event.type == "message_stop":
                        pass

        await asyncio.wait_for(_consume_stream(), timeout=_STREAM_TIMEOUT)

        elapsed = _time.monotonic() - t0
        logger.info("Planner stream iteration %d completed in %.1fs (%d tool calls)", iteration, elapsed, len(tool_uses))

        full_text = "".join(text_chunks)
        if full_text.strip():
            assistant_content.insert(0, {"type": "text", "text": full_text})

        if not tool_uses:
            if on_status and not is_text_response:
                await on_status("[think] Preparing response...\n")
            try:
                parsed = json.loads(full_text)
            except json.JSONDecodeError:
                parsed = _extract_json(full_text)
            parsed.setdefault("assistant_message", "")
            parsed.setdefault("decisions", [])
            return parsed

        messages.append({"role": "assistant", "content": assistant_content})
        tool_result_blocks: list[dict[str, Any]] = []
        for tool_use in tool_uses:
            result = await _handle_tool_call(tool_use.name, tool_use.input, allowed_roots=allowed_roots)
            if on_status:
                await on_status(_tool_done_status(tool_use.name) + "\n")
            tool_result_blocks.append({
                "type": "tool_result",
                "tool_use_id": tool_use.id,
                "content": result,
            })
        messages.append({"role": "user", "content": tool_result_blocks})


def _should_force_heavy(conversation: Conversation) -> bool:
    """Determine if this turn must use the HEAVY (Opus + tools) path."""
    if len(conversation.messages) <= 4:
        return True
    if conversation.contract.dataset is None and len(conversation.messages) > 2:
        return True
    has_needs_input = any(
        d.status == DecisionStatus.NEEDS_INPUT
        for d in conversation.decisions
    )
    if has_needs_input:
        return True
    return False


async def respond(
    conversation: Conversation,
    hardware: HardwareProfile | None = None,
    on_status: StatusCallback | None = None,
) -> dict[str, Any]:
    messages = _conversation_to_messages(conversation)
    if not messages:
        return {"assistant_message": "Describe your ML task or point me to a dataset.", "decisions": []}

    if _should_force_heavy(conversation):
        logger.info("Forcing HEAVY tier (discovery/early/needs_input)")
        return await _respond_heavy(conversation, hardware, on_status=on_status)

    latest_user = next(
        (m["content"] for m in reversed(messages) if m["role"] == "user"), ""
    )

    tier = await router_classify(latest_user, messages)
    logger.info("Router tier for message: %s", tier.value)

    if tier == Tier.LIGHT:
        if on_status:
            await on_status("[think] Processing...\n")
        result = await _respond_light(conversation, hardware)
        if result.get("_upgrade_to_heavy"):
            logger.info("LIGHT returned decisions, upgrading to HEAVY")
            return await _respond_heavy(conversation, hardware, on_status=on_status)
        return result
    return await _respond_heavy(conversation, hardware, on_status=on_status)
