from __future__ import annotations

import json
import sys
from pathlib import Path
from statistics import mean

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from bdo_ai_env.baseline_agent import greedy_policy, stingy_policy
from models import BDOAction
from server.bdo_environment import BDOEnvironment


def run_policy(policy_name: str, scenario: str):
    policy_fn = {"greedy": greedy_policy, "stingy": stingy_policy}[policy_name]
    env = BDOEnvironment(scenario=scenario)
    observation = env.reset()
    monthly = []
    while not observation.done:
        action_batch = policy_fn(observation.model_dump(mode="json", exclude_none=True))
        observation = env.step(BDOAction.model_validate(action_batch))
        step_info = observation.info or observation.metadata
        monthly.append(
            {
                "month": observation.meta.month - 1,
                "reward": observation.reward,
                **step_info["reward_breakdown"],
            }
        )

    final = monthly[-1]
    return {
        "policy": policy_name,
        "scenario": scenario,
        "shock_month": env.state.shock_month,
        "final_reward": round(final["total"], 4),
        "avg_reward": round(mean(float(item["reward"] or 0.0) for item in monthly), 4),
        "min_reward": round(min(float(item["reward"] or 0.0) for item in monthly), 4),
        "final_coverage": final["coverage"],
        "final_security": final["fraud_prevention"],
        "final_solvency": final["solvency"],
        "final_stability": final["stability"],
        "final_belief": final["belief_accuracy"],
        "monthly": monthly,
    }


def main() -> None:
    scenario = sys.argv[1] if len(sys.argv) > 1 else "black_swan"
    output_path = sys.argv[2] if len(sys.argv) > 2 else None
    report = {
        "scenario": scenario,
        "results": [
            run_policy("greedy", scenario),
            run_policy("stingy", scenario),
        ],
    }
    rendered = json.dumps(report, indent=2)
    if output_path:
        output = Path(output_path)
        output.parent.mkdir(parents=True, exist_ok=True)
        output.write_text(rendered, encoding="utf-8")
    print(rendered)


if __name__ == "__main__":
    main()