"""Simulator rollout helpers with scenario perturbations for evaluation."""

from __future__ import annotations

import os
from typing import Any

from app.agents.orchestrator import Orchestrator
from app.common.enums import DoseBucket
from app.env.env_core import PolyGuardEnv


def _apply_perturbation(env: PolyGuardEnv, perturbation: str | None) -> None:
    if not perturbation:
        return

    state = env.state
    meds = state.patient.medications

    if perturbation == "missing_labs":
        state.patient.labs.egfr = None
        state.patient.labs.ast = None
        state.patient.labs.alt = None
    elif perturbation == "noisy_dose_info":
        for idx, med in enumerate(meds):
            if idx % 2 == 0:
                med.dose_bucket = DoseBucket.HIGH if med.dose_bucket != DoseBucket.HIGH else DoseBucket.LOW
    elif perturbation == "conflicting_meds" and meds:
        meds.append(meds[0].model_copy())
    elif perturbation == "alias_noise" and meds:
        meds[0].drug = f"{meds[0].drug}_alias"
    elif perturbation == "hidden_duplicate" and meds:
        meds.append(meds[0].model_copy(update={"drug": meds[0].drug}))
    elif perturbation == "stale_evidence":
        state.unresolved_conflicts.append("evidence_stale")
    elif perturbation == "delayed_ade":
        state.patient.latent_confounders["delayed_ade"] = 0.8


def run_rollouts(
    episodes: int = 5,
    difficulty: str = "medium",
    sub_environment: str | None = None,
    perturbation: str | None = None,
    seed_offset: int = 900,
    policy_stack: str = "llm+bandit",
) -> list[dict[str, Any]]:
    previous_policy = os.getenv("POLYGUARD_POLICY_STACK")
    os.environ["POLYGUARD_POLICY_STACK"] = policy_stack

    env = PolyGuardEnv()
    orchestrator = Orchestrator(env)
    rows: list[dict[str, Any]] = []

    for i in range(episodes):
        env.reset(seed=seed_offset + i, difficulty=difficulty, sub_environment=sub_environment)
        _apply_perturbation(env, perturbation=perturbation)

        done = False
        while not done:
            out = orchestrator.run_step()
            done = bool(out.get("done"))
            info = out.get("info", {}) if isinstance(out.get("info", {}), dict) else {}
            critic = out.get("critic", {}) if isinstance(out.get("critic", {}), dict) else {}
            reward_breakdown = info.get("reward_breakdown", {}) if isinstance(info.get("reward_breakdown", {}), dict) else {}
            primary_channels = (
                info.get("primary_reward_channels", {})
                if isinstance(info.get("primary_reward_channels", {}), dict)
                else {}
            )
            final_action = out.get("final_action", {}) if isinstance(out.get("final_action", {}), dict) else {}

            rows.append(
                {
                    "episode": i,
                    "step": int(env.state.step_count),
                    "reward": float(out.get("reward", 0.0)),
                    "done": done,
                    "legal": bool(critic.get("legal", False)),
                    "severe_violation": len(critic.get("violations", [])) > 1,
                    "abstain": str(final_action.get("action_type", "")).startswith("REQUEST_"),
                    "termination_reason": info.get("termination_reason"),
                    "step_timeout": bool(info.get("step_timeout")),
                    "failure_reasons": info.get("failure_reasons", []),
                    "invalid_action_count": int(info.get("invalid_action_count", 0)),
                    "reward_breakdown": reward_breakdown,
                    "primary_reward_channels": primary_channels,
                    "policy_stack": policy_stack,
                    "difficulty": difficulty,
                    "sub_environment": sub_environment,
                    "perturbation": perturbation,
                }
            )

    if previous_policy is None:
        os.environ.pop("POLYGUARD_POLICY_STACK", None)
    else:
        os.environ["POLYGUARD_POLICY_STACK"] = previous_policy

    return rows