inference.py

"""Inference runner for the Flatmate RL environment.

Uses the same Docker-backed pattern as `sudoku_rl/inference.py`.
It can either:
- call a chat model for the next action, or
- run purely with the built-in heuristic policy.
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
import re
import textwrap
from dataclasses import dataclass
from types import SimpleNamespace
from typing import Any

from openai import OpenAI
from openenv.core.containers.runtime.providers import LocalDockerProvider

try:
    from flatmate_rl.env_config import load_repo_env
    from flatmate_rl import FlatmateRlAction, FlatmateRlEnv
    from flatmate_rl.server.heuristic_policy import autopolicy_next_request, expected_policy_action
    from flatmate_rl.server.scenarios import SCENARIOS
except ImportError:
    from env_config import load_repo_env
    from __init__ import FlatmateRlAction, FlatmateRlEnv
    from server.heuristic_policy import autopolicy_next_request, expected_policy_action
    from server.scenarios import SCENARIOS

load_repo_env()


IMAGE_NAME = os.getenv("IMAGE_NAME") or "flatmate_rl:latest"
API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY")

API_BASE_URL = os.getenv("API_BASE_URL") or "https://router.huggingface.co/v1"
MODEL_NAME = os.getenv("MODEL_NAME") or "Qmeta-llama/Llama-3.1-8B-Instruct"
#meta-llama/Llama-3.1-8B-Instruct
MAX_STEPS_ENV = os.getenv("MAX_STEPS")
TEMPERATURE = float(os.getenv("TEMPERATURE", "0.2"))
MAX_TOKENS = int(os.getenv("MAX_TOKENS", "256"))
MALFORMED_ACTION_PENALTY = -0.05


class ModelConfigurationError(RuntimeError):
    """Raised when configured model inference cannot be used."""


SYSTEM_PROMPT = textwrap.dedent(
    """
    You are a flatmate visit-scheduling broker agent.

    Return exactly one JSON object with this schema:
    {"action_type":"assistant_message","assistant_message":"..."}
    or
    {"action_type":"tool_call","tool_name":"...","tool_arguments":{...}}

    Example valid tool call:
    {"action_type":"tool_call","tool_name":"store_user_details","tool_arguments":{}}

    Rules:
    - Return JSON only.
    - Use only tools listed in available_tools.
    - Never put a tool name in action_type.
    - Follow the observation state exactly.
    - If a tool can perform the next required operation, call the tool immediately.
    - Do not send acknowledgement or progress messages such as "I will search now" when a tool call is needed.
    - Prefer safe, incremental progress toward storing user details, matching listings, and booking visits.
    - Use exact tool argument names from the prompt. Never invent aliases such as visit_time.
    - Treat negative reward, violations, and feedback_summary as corrective feedback for the next action.
    """
).strip()


TOOL_CONTRACT_PROMPT = textwrap.dedent(
    """
    Tool argument contract:
    - store_user_details: tool_arguments can be {} after required buyer fields are gathered.
    - search_posts: tool_arguments can be {}.
    - match_location_preference: {"post_ids":["post_id", ...]}.
    - get_commute_time: {"post_ids":["post_id", ...]}.
    - check_calendar_slots: {"post_ids":["post_id", ...]}.
    - shortlist: {"post_ids":["post_id", ...]}.
    - contact_poster: {"post_id":"post_id","time_text":"exact slot from check_calendar_slots"}. This shares the buyer profile with the seller/poster and asks them to confirm profile fit plus visit time.
    - book_viewing: {"post_id":"post_id","time_text":"same exact slot confirmed by buyer and poster"}.

    Booking workflow:
    1. Ask for missing buyer fields before store_user_details.
    2. Store buyer details, search posts, match location, get commute time, then check calendar slots.
    3. Ask the buyer to confirm one exact slot from check_calendar_slots.
    4. Call contact_poster with post_id and time_text for that same slot.
    5. Only after buyer_confirmed and poster_confirmed are true, call book_viewing with post_id and time_text.
    """
).strip()


def log_start(task_id: str, model: str, source: str) -> None:
    print(f"[START] scenario={task_id} model={model} source={source}", flush=True)


def log_end(task_id: str, success: bool, steps: int, total_reward: float, booked_visits: int, final_status: str) -> None:
    print(
        f"[END] scenario={task_id} success={str(success).lower()} steps={steps} "
        f"total_reward={total_reward:.2f} booked_visits={booked_visits} final_status={final_status}",
        flush=True,
    )


def format_action(action: FlatmateRlAction | dict[str, Any] | None) -> str:
    if action is None:
        return "None"
    if isinstance(action, dict):
        try:
            action = FlatmateRlAction.model_validate(action)
        except Exception:
            return json.dumps(action, ensure_ascii=False, sort_keys=True)
    if action.action_type == "tool_call":
        return json.dumps(
            {
                "action_type": action.action_type,
                "tool_name": action.tool_name,
                "tool_arguments": action.tool_arguments,
            },
            ensure_ascii=False,
            sort_keys=True,
        )
    return json.dumps(
        {
            "action_type": action.action_type,
            "assistant_message": action.assistant_message,
        },
        ensure_ascii=False,
        sort_keys=True,
    )


def format_action_with_reasoning(action: FlatmateRlAction | None, reasoning: dict[str, Any] | None) -> str:
    if action is None:
        return "None"
    payload = json.loads(format_action(action))
    if reasoning and "error" not in reasoning:
        payload["reasoning"] = reasoning.get("decision_summary") or reasoning.get("why_this_action_now") or reasoning
    return json.dumps(payload, ensure_ascii=False, sort_keys=True)


def actions_match(actual: FlatmateRlAction | None, expected: FlatmateRlAction | None) -> bool:
    if actual is None:
        return False
    if expected is None:
        return True
    if actual.action_type != expected.action_type:
        return False
    if actual.action_type == "assistant_message":
        return bool(actual.assistant_message.strip())
    return actual.tool_name == expected.tool_name


def scenario_check_snapshot(task_id: str, observation: Any) -> dict[str, Any]:
    scenario = SCENARIOS[task_id]
    required_tools = scenario["ground_truth"]["required_tool_calls"]
    used_tools = [item.get("tool", "") for item in observation.tool_trace]
    missing_required_tools = [tool for tool in required_tools if tool not in used_tools]
    required_info = scenario["ground_truth"]["required_info"]
    gathered_fields = set(observation.gathered_fields)
    missing_required_info = [field for field in required_info if field not in gathered_fields]
    return {
        "required_bookings": scenario["ground_truth"]["required_bookings"],
        "bookings_so_far": len(observation.booked_visits),
        "required_tools": required_tools,
        "used_tools": used_tools,
        "missing_required_tools": missing_required_tools,
        "required_info": required_info,
        "gathered_fields": list(observation.gathered_fields),
        "missing_required_info": missing_required_info,
        "violations": list(observation.violations),
        "done": bool(observation.done),
    }


def log_verbose_scenario(task_id: str) -> None:
    scenario = SCENARIOS[task_id]
    truth = scenario["ground_truth"]
    buyer = scenario["buyer_profile"]
    creation_config = scenario["scenario_creation_config"]
    print("[VERBOSE] scenario_definition", flush=True)
    print(
        json.dumps(
            {
                "task_id": scenario["task_id"],
                "label": scenario["label"],
                "description": scenario["description"],
                "task_post_ids": scenario["task_post_ids"],
                "buyer_profile": {
                    "budget_max": buyer["budget_max"],
                    "dietary": buyer["dietary"],
                    "areas": buyer["areas"],
                    "occupation": buyer["occupation"],
                    "visit_availability": buyer["visit_availability"],
                    "initial_disclosure_fields": buyer["initial_disclosure_fields"],
                },
                "expected_answers": creation_config.get("expected_answers", {}),
                "followup_seller_expected_answers": creation_config.get("followup_seller_expected_answers", {}),
                "ground_truth": truth,
            },
            ensure_ascii=False,
            indent=2,
            sort_keys=True,
        ),
        flush=True,
    )
    print("[VERBOSE] scenario_check_functions", flush=True)
    print(
        json.dumps(
            {
                "field_tracking": "server/episode.py:_required_fields, _remaining_fields, _buyer_response",
                "tool_gating": "server/episode.py:_execute_tool",
                "buyer_store_check": "server/episode.py:_tool_store_user_details",
                "search_check": "server/episode.py:_tool_search_posts",
                "slot_fetch_check": "server/episode.py:_tool_check_calendar_slots",
                "poster_confirmation_check": "server/episode.py:_tool_contact_poster",
                "booking_check": "server/episode.py:_tool_book_viewing",
                "completion_check": "server/episode.py:_tool_book_viewing and _maybe_finish_from_message",
                "reward_check": "server/episode.py:_handle_assistant_message and _handle_tool_call",
            },
            ensure_ascii=False,
            indent=2,
            sort_keys=True,
        ),
        flush=True,
    )


def log_verbose_step(
    *,
    task_id: str,
    step: int,
    raw_observation: Any,
    policy_observation: Any,
    expected_action: FlatmateRlAction | None,
    actual_action: FlatmateRlAction | None,
    model_raw_response: str | None,
    model_debug_explanation: dict[str, Any] | None,
) -> None:
    user_prompt = build_user_prompt(step=step, observation=policy_observation)
    print(f"[VERBOSE] step={step} pre_step_checks", flush=True)
    print(
        json.dumps(
            {
                "full_state_checks": scenario_check_snapshot(task_id, raw_observation),
                "strict_eval": {
                    "scenario_id_visible": bool(policy_observation.scenario_id),
                    "scenario_label_visible": bool(policy_observation.scenario_label),
                    "gathered_fields_visible": policy_observation.gathered_fields,
                    "remaining_required_fields_visible": policy_observation.remaining_required_fields,
                    "violations_visible": policy_observation.violations,
                    "tool_trace_visible": policy_observation.tool_trace,
                    "total_reward_visible": policy_observation.total_reward,
                    "last_tool_result_visible": policy_observation.last_tool_result,
                },
                "expected_action_from_full_state": format_action(expected_action),
                "actual_action_from_policy_input": format_action(actual_action),
                "action_match": actions_match(actual_action, expected_action),
                "broker_feedback_payload": policy_observation.model_dump(),
            },
            ensure_ascii=False,
            indent=2,
            sort_keys=True,
        ),
        flush=True,
    )
    print(f"[VERBOSE] step={step} llm_system_prompt", flush=True)
    print(SYSTEM_PROMPT, flush=True)
    print(f"[VERBOSE] step={step} llm_user_prompt", flush=True)
    print(user_prompt, flush=True)
    print(f"[VERBOSE] step={step} llm_raw_response", flush=True)
    print(model_raw_response if model_raw_response is not None else "null", flush=True)
    print(f"[VERBOSE] step={step} llm_decision_explanation", flush=True)
    print(
        json.dumps(model_debug_explanation or {"message": "No model explanation available for this step."}, ensure_ascii=False, indent=2, sort_keys=True),
        flush=True,
    )


def log_verbose_post_step(task_id: str, step: int, observation: Any) -> None:
    print(f"[VERBOSE] step={step} post_step_checks", flush=True)
    print(
        json.dumps(
            scenario_check_snapshot(task_id, observation),
            ensure_ascii=False,
            indent=2,
            sort_keys=True,
        ),
        flush=True,
    )


def extract_new_chat_entries(history: list[dict[str, Any]], start_idx: int) -> tuple[list[dict[str, Any]], int]:
    return history[start_idx:], len(history)


def _clean_text(value: str | None) -> str:
    return re.sub(r"\s+", " ", value or "").strip()


def _block(title: str, body: str) -> str:
    return f"{title}\n{body}"


def log_initial_conversation(observation: Any) -> tuple[int, int]:
    buyer_history = observation.buyer_conversation_history
    seller_history = observation.seller_conversation_history
    if buyer_history:
        print(_block("user initiated conversation", _clean_text(buyer_history[0].get("content", ""))), flush=True)
        print("", flush=True)
    if seller_history:
        print(_block("seller initiated conversation", _clean_text(seller_history[0].get("content", ""))), flush=True)
        print("", flush=True)
    return len(buyer_history), len(seller_history)


def log_step_report(
    *,
    step: int,
    action: FlatmateRlAction | None,
    expected_action: FlatmateRlAction | None,
    model_raw_response: str | None,
    model_debug_explanation: dict[str, Any] | None,
    buyer_entries: list[dict[str, Any]],
    seller_entries: list[dict[str, Any]],
    reward: float,
    total_reward: float,
    status: str,
    done: bool,
    tool_result: dict[str, Any],
    message: str,
    source: str,
    error: str | None,
) -> None:
    policy_status = "success" if actions_match(action, expected_action) else "violation"
    print(f"step {step}", flush=True)
    print("", flush=True)

    for entry in buyer_entries:
        label = "broker response :" if entry.get("role") == "assistant" else "buyer_chat :"
        print(_block(label, _clean_text(entry.get("content", ""))), flush=True)
        print("", flush=True)

    for entry in seller_entries:
        label = "broker response :" if entry.get("role") == "assistant" else "seller_chat :"
        print(_block(label, _clean_text(entry.get("content", ""))), flush=True)
        print("", flush=True)

    print(_block("expected_response :", format_action(expected_action)), flush=True)
    print("", flush=True)

    print(_block("llm_raw_response :", model_raw_response or "null"), flush=True)
    print("", flush=True)

    print(_block("llm_parsed_action :", format_action_with_reasoning(action, model_debug_explanation)), flush=True)
    print("", flush=True)

    print(_block("policy_check :", policy_status), flush=True)
    print("", flush=True)

    tool_used = action.tool_name if action is not None and action.action_type == "tool_call" else "none"
    print(_block("tool_used :", tool_used), flush=True)
    print("", flush=True)

    print(_block("tool_result :", json.dumps(tool_result, ensure_ascii=False, sort_keys=True)), flush=True)
    print("", flush=True)

    reward_text = (
        f"step_reward={reward:.2f}\n"
        f"total_reward={total_reward:.2f}\n"
        f"status={status}\n"
        f"done={str(done).lower()}\n"
        f"source={source}\n"
        f"message={_clean_text(message)}\n"
        f"error={error or 'null'}"
    )
    print(_block("reward :", reward_text), flush=True)
    print("\n" + ("-" * 60) + "\n", flush=True)


def build_user_prompt(step: int, observation: Any) -> str:
    last_tool_result = dict(observation.last_tool_result or {})
    last_tool_result.pop("stored_profile", None)
    return textwrap.dedent(
        f"""
        Step: {step}
        Phase: {observation.phase}
        Status: {observation.status}
        Feedback summary: {observation.feedback_summary}
        Environment message: {observation.message}
        Step reward: {observation.step_reward}
        Total reward: {observation.total_reward}
        Violations: {observation.violations}
        Remaining required fields: {observation.remaining_required_fields}
        Available tools: {observation.available_tools}
        Last tool result: {json.dumps(last_tool_result, ensure_ascii=False)}
        Prerequisites satisfied: {json.dumps(observation.prerequisites_satisfied, ensure_ascii=False)}
        Recent tool calls: {json.dumps(observation.recent_tool_calls, ensure_ascii=False)}
        Booked visits: {observation.booked_visits}

        {TOOL_CONTRACT_PROMPT}

        Buyer/Broker transcript:
        {json.dumps(observation.buyer_conversation_history[-8:], ensure_ascii=False)}

        Seller/Broker transcript:
        {json.dumps(observation.seller_conversation_history[-8:], ensure_ascii=False)}

        Return the next action as JSON only.
        """
    ).strip()


@dataclass
class ParsedAction:
    action: FlatmateRlAction | None
    error: str | None = None
    warning: str | None = None


def _schema_error_details(candidate: dict[str, Any]) -> str | None:
    action_type = candidate.get("action_type")
    if action_type not in {"assistant_message", "tool_call"}:
        return f"action_type must be 'assistant_message' or 'tool_call', got {action_type!r}"
    if action_type == "assistant_message" and not str(candidate.get("assistant_message", "")).strip():
        return "assistant_message is required when action_type is 'assistant_message'"
    if action_type == "tool_call" and not str(candidate.get("tool_name", "")).strip():
        return "tool_name is required when action_type is 'tool_call'"
    return None


def normalize_action_candidate(candidate: dict[str, Any]) -> tuple[dict[str, Any], str | None]:
    normalized = dict(candidate)
    warning = None
    action_type = str(normalized.get("action_type", "")).strip()
    tool_name = str(normalized.get("tool_name", "")).strip()

    if not action_type and tool_name:
        normalized["action_type"] = "tool_call"
        normalized.setdefault("tool_arguments", {})
        warning = "coerced missing action_type to tool_call because tool_name was present"
        return normalized, warning

    if action_type in {"assistant", "message"} and "assistant_message" in normalized:
        normalized["action_type"] = "assistant_message"
        normalized.pop("tool_name", None)
        normalized.pop("tool_arguments", None)
        warning = f"coerced action_type {action_type!r} to assistant_message"
        return normalized, warning

    if action_type and action_type not in {"assistant_message", "tool_call"} and not tool_name:
        normalized["action_type"] = "tool_call"
        normalized["tool_name"] = action_type
        normalized.setdefault("tool_arguments", {})
        warning = f"coerced invalid action_type {action_type!r} into tool_name"
        return normalized, warning

    return normalized, warning


def parse_action(text: str, *, strict: bool = True) -> ParsedAction:
    cleaned = (text or "").strip()
    if not cleaned:
        return ParsedAction(action=None, error="json_parse_failed: empty model response")

    try:
        parsed = json.loads(cleaned)
    except json.JSONDecodeError as exc:
        if strict:
            return ParsedAction(action=None, error=f"json_parse_failed: {exc.msg} at line {exc.lineno} column {exc.colno}")
        parsed = None

    if strict:
        if not isinstance(parsed, dict):
            return ParsedAction(action=None, error=f"schema_validation_failed: expected JSON object, got {type(parsed).__name__}")
        schema_detail = _schema_error_details(parsed)
        if schema_detail:
            return ParsedAction(action=None, error=f"schema_validation_failed: {schema_detail}")
        try:
            return ParsedAction(action=FlatmateRlAction.model_validate(parsed))
        except Exception as exc:
            return ParsedAction(action=None, error=f"schema_validation_failed: {exc}")

    candidates: list[dict[str, Any]] = []
    if isinstance(parsed, dict):
        candidates.append(parsed)

    for match in re.finditer(r"\{.*\}", cleaned, flags=re.DOTALL):
        try:
            extracted = json.loads(match.group(0))
        except json.JSONDecodeError:
            continue
        if isinstance(extracted, dict):
            candidates.append(extracted)

    last_error = "schema_validation_failed: no valid action object found"
    for candidate in candidates:
        normalized, warning = normalize_action_candidate(candidate)
        try:
            return ParsedAction(action=FlatmateRlAction.model_validate(normalized), warning=warning)
        except Exception as exc:
            last_error = f"schema_validation_failed: {exc}"
            continue
    return ParsedAction(action=None, error=last_error)


def malformed_action_observation(observation: Any, details: str) -> Any:
    payload = observation.model_dump()
    malformed_result = {
        "error": "schema_validation_failed",
        "details": details,
        "expected_schema": FlatmateRlAction.model_json_schema(),
    }
    payload["status"] = "tool_result"
    payload["last_tool_result"] = malformed_result
    payload["tool_results"] = list(payload.get("tool_results", [])) + [malformed_result]
    payload["step_reward"] = MALFORMED_ACTION_PENALTY
    payload["reward"] = MALFORMED_ACTION_PENALTY
    payload["total_reward"] = float(payload.get("total_reward", 0.0)) + MALFORMED_ACTION_PENALTY
    payload["message"] = "Malformed action output. Return a valid FlatmateRlAction JSON object."
    payload["feedback_summary"] = "Use action_type='assistant_message' or action_type='tool_call' with a non-empty tool_name."
    payload["done"] = False
    return type(observation).model_validate(payload)


def apply_total_reward_adjustment(observation: Any, adjustment: float) -> Any:
    if not adjustment:
        return observation
    payload = observation.model_dump()
    payload["total_reward"] = float(payload.get("total_reward", 0.0)) + adjustment
    return type(observation).model_validate(payload)


def sanitize_observation_for_policy(observation: Any, strict_eval: bool) -> Any:
    if not strict_eval:
        return observation
    payload = observation.model_dump()
    payload["scenario_id"] = ""
    payload["scenario_label"] = ""
    payload["difficulty"] = ""
    payload["gathered_fields"] = []
    payload["remaining_required_fields"] = []
    payload["violations"] = []
    payload["tool_trace"] = []
    payload["step_reward"] = 0.0
    payload["total_reward"] = 0.0
    payload["last_tool_result"] = {
        key: value
        for key, value in payload["last_tool_result"].items()
        if key != "stored_profile"
    }
    payload["tool_results"] = [
        {key: value for key, value in item.items() if key != "stored_profile"}
        for item in payload["tool_results"]
    ]
    return type(observation).model_validate(payload)


def missing_fields_from_feedback(observation: dict[str, Any]) -> list[str]:
    feedback = " ".join(
        [
            str(observation.get("feedback_summary", "")),
            str(observation.get("message", "")),
            str(observation.get("last_tool_result", {}).get("message", "")),
        ]
    ).lower()
    fields = []
    for field in ["diet", "visit_availability", "occupation", "budget", "areas", "listing_choices"]:
        phrases = {
            "visit_availability": ["visit_availability", "visit availability"],
            "listing_choices": ["listing_choices", "listing choices"],
        }.get(field, [field])
        if any(phrase in feedback for phrase in phrases):
            fields.append(field)
    return fields


def heuristic_action(task_id: str, observation: Any) -> FlatmateRlAction:
    payload = expected_policy_action(task_id, observation.model_dump())
    if payload is None:
        raise RuntimeError("Heuristic policy produced no action for a non-terminal observation.")
    return FlatmateRlAction.model_validate(payload)


def build_explanation_prompt(step: int, observation: Any, action: FlatmateRlAction, raw_response: str) -> str:
    return textwrap.dedent(
        f"""
        You are auditing a broker policy decision in a flatmate scheduling environment.

        Explain the chosen action briefly as structured JSON only.

        Observation:
        {build_user_prompt(step=step, observation=observation)}

        Raw model response:
        {raw_response}

        Parsed chosen action:
        {format_action(action)}

        Return JSON with keys:
        - decision_summary
        - action_type
        - chosen_tool_or_message
        - why_this_action_now
        - evidence_from_state
        - why_not_other_tools
        - risks_or_uncertainties
        """
    ).strip()


def get_model_explanation(
    client: OpenAI,
    step: int,
    observation: Any,
    action: FlatmateRlAction,
    raw_response: str,
) -> dict[str, Any]:
    try:
        completion = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {
                    "role": "system",
                    "content": "You are a debugging assistant. Explain the selected broker action briefly in JSON only. Do not reveal hidden reasoning. Report explicit decision factors from the visible state.",
                },
                {
                    "role": "user",
                    "content": build_explanation_prompt(
                        step=step,
                        observation=observation,
                        action=action,
                        raw_response=raw_response,
                    ),
                },
            ],
            temperature=0.0,
            max_tokens=400,
            stream=False,
        )
        text = (completion.choices[0].message.content or "").strip()
        parsed = json.loads(text)
        if isinstance(parsed, dict):
            return parsed
        return {"message": text}
    except Exception as exc:
        return {"error": str(exc)}


def get_model_action(
    client: OpenAI,
    task_id: str,
    step: int,
    observation: Any,
    explain: bool,
    strict_parsing: bool,
) -> tuple[FlatmateRlAction | None, str, str | None, str, dict[str, Any] | None]:
    user_prompt = build_user_prompt(step=step, observation=observation)
    try:
        completion = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt},
            ],
            temperature=TEMPERATURE,
            max_tokens=MAX_TOKENS,
            stream=False,
        )
        text = (completion.choices[0].message.content or "").strip()
        parsed = parse_action(text, strict=strict_parsing)
        if parsed.warning:
            print(f"[WARN] parser_coercion {parsed.warning}", flush=True)
        if parsed.action is not None:
            explanation = get_model_explanation(client, step, observation, parsed.action, text) if explain else None
            return parsed.action, "model", None, text, explanation
        if strict_parsing:
            explanation = {"error": parsed.error} if parsed.error else None
            return None, "model_parse_error", parsed.error or "invalid_model_output", text, explanation
        fallback_action = heuristic_action(task_id, observation)
        if parsed.error:
            print(f"[WARN] parser_fallback {parsed.error}", flush=True)
        explanation = {
            "message": "Primary model output could not be parsed, so heuristic fallback was used.",
            "raw_model_response": text,
            "parse_error": parsed.error,
            "fallback_action": json.loads(format_action(fallback_action)),
        }
        return fallback_action, "heuristic_fallback", f"unparseable_model_output={text!r}", text, explanation
    except Exception as exc:
        raise ModelConfigurationError(
            f"MODEL_NAME is invalid or unsupported for API_BASE_URL: {MODEL_NAME!r}. "
            "Use --heuristic-only to run the heuristic policy, or set MODEL_NAME to a supported model. "
            f"Provider error: {exc}"
        ) from exc


async def run_scenario(
    env: FlatmateRlEnv,
    task_id: str,
    client: OpenAI | None,
    max_steps: int | None,
    strict_eval: bool,
    verbose: bool,
    strict_parsing: bool,
) -> dict[str, Any]:
    result = await env.reset(scenario_id=task_id)
    observation = result.observation
    limit = max_steps if max_steps is not None else 24
    steps_taken = 0
    buyer_logged_count = 0
    seller_logged_count = 0
    local_reward_adjustment = 0.0

    log_start(task_id=task_id, model=MODEL_NAME if client else "heuristic", source="model" if client else "heuristic")
    if verbose:
        log_verbose_scenario(task_id)
    buyer_logged_count, seller_logged_count = log_initial_conversation(observation)

    for step in range(1, limit + 1):
        if result.done:
            break

        policy_observation = sanitize_observation_for_policy(observation, strict_eval=strict_eval)
        expected_action = None
        model_raw_response = None
        model_debug_explanation = None
        expected_payload = expected_policy_action(task_id, observation.model_dump())
        if expected_payload is not None:
            expected_action = FlatmateRlAction.model_validate(expected_payload)
        if client is None:
            # Heuristic is the ground-truth policy: give it the full unredacted observation
            # so tool_trace (and other state) is intact even when --strict-eval is active.
            action = heuristic_action(task_id, observation)
            source = "heuristic"
            error = None
        else:
            action, source, error, model_raw_response, model_debug_explanation = get_model_action(
                client=client,
                task_id=task_id,
                step=step,
                observation=policy_observation,
                explain=verbose,
                strict_parsing=strict_parsing,
            )
        if verbose:
            log_verbose_step(
                task_id=task_id,
                step=step,
                raw_observation=observation,
                policy_observation=policy_observation,
                expected_action=expected_action,
                actual_action=action,
                model_raw_response=model_raw_response,
                model_debug_explanation=model_debug_explanation,
            )

        if action is None:
            local_reward_adjustment += MALFORMED_ACTION_PENALTY
            observation = malformed_action_observation(observation, error or "invalid_model_output")
            result = SimpleNamespace(observation=observation, reward=MALFORMED_ACTION_PENALTY, done=False)
        else:
            result = await env.step(action)
            observation = result.observation
            observation = apply_total_reward_adjustment(observation, local_reward_adjustment)
            result = SimpleNamespace(observation=observation, reward=result.reward, done=result.done)
        steps_taken = step
        buyer_entries, buyer_logged_count = extract_new_chat_entries(
            observation.buyer_conversation_history,
            buyer_logged_count,
        )
        seller_entries, seller_logged_count = extract_new_chat_entries(
            observation.seller_conversation_history,
            seller_logged_count,
        )

        log_step_report(
            step=step,
            action=action,
            expected_action=expected_action,
            model_raw_response=model_raw_response,
            model_debug_explanation=model_debug_explanation,
            buyer_entries=buyer_entries,
            seller_entries=seller_entries,
            reward=float(result.reward or 0.0),
            total_reward=float(observation.total_reward),
            status=observation.status,
            done=result.done,
            tool_result=observation.last_tool_result,
            source=source,
            message=observation.message,
            error=error,
        )
        if verbose:
            log_verbose_post_step(task_id=task_id, step=step, observation=observation)

        if result.done:
            break

    success = bool(observation.booked_visits) and not observation.violations
    summary = {
        "scenario": task_id,
        "success": success,
        "steps": steps_taken,
        "total_reward": float(observation.total_reward),
        "booked_visits": observation.booked_visits,
        "violations": observation.violations,
        "status": observation.status,
    }
    log_end(
        task_id=task_id,
        success=success,
        steps=steps_taken,
        total_reward=float(observation.total_reward),
        booked_visits=len(observation.booked_visits),
        final_status=observation.status,
    )
    return summary


async def connect_env(image_name: str, startup_timeout_s: float) -> FlatmateRlEnv:
    provider = LocalDockerProvider()
    base_url = provider.start_container(image_name)
    provider.wait_for_ready(base_url, timeout_s=startup_timeout_s)
    env = FlatmateRlEnv(base_url=base_url, provider=provider)
    await env.connect()
    return env


async def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--scenario-id",
        action="append",
        dest="scenario_ids",
        choices=sorted(SCENARIOS.keys()),
        help="Scenario id to run. May be provided multiple times.",
    )
    parser.add_argument("--max-steps", type=int, default=int(MAX_STEPS_ENV) if MAX_STEPS_ENV else None)
    parser.add_argument("--heuristic-only", action="store_true", help="Skip model calls and use only the heuristic policy.")
    parser.add_argument("--startup-timeout", type=float, default=90.0)
    parser.add_argument("--strict-eval", action="store_true", help="Hide scenario metadata and reward signals from the broker policy.")
    parser.add_argument(
        "--strict-parsing",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Reject malformed action JSON instead of coercing it. Use --no-strict-parsing for legacy coercion.",
    )
    parser.add_argument("--verbose", action="store_true", help="Print scenario checks, expected actions, and detailed state diagnostics.")
    args = parser.parse_args()

    scenario_ids = args.scenario_ids or [next(iter(SCENARIOS))]
    client = None if args.heuristic_only or not API_KEY else OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
    env = await connect_env(IMAGE_NAME, startup_timeout_s=args.startup_timeout)

    try:
        summaries = []
        for task_id in scenario_ids:
            summaries.append(
                await run_scenario(
                    env=env,
                    task_id=task_id,
                    client=client,
                    max_steps=args.max_steps,
                    strict_eval=args.strict_eval,
                    verbose=args.verbose,
                    strict_parsing=args.strict_parsing,
                )
            )
        print("[SUMMARY] " + json.dumps(summaries, ensure_ascii=False), flush=True)
    except ModelConfigurationError as exc:
        print(f"[ERROR] {exc}", flush=True)
        raise SystemExit(2) from exc
    finally:
        await env.close()


if __name__ == "__main__":
    asyncio.run(main())