from __future__ import annotations

import json
import sys
from pathlib import Path

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

from models import BDOAction
from server.bdo_environment import BDOEnvironment


PROMPT = """
Enter a JSON action batch, for example:
{"predicted_fraud_level": 0.4, "actions": [{"name": "allocate_funds", "params": {"village": "Loba", "amount": 5000}}, {"name": "approve_batch", "params": {"village": "Loba", "mode": "conservative"}}]}

You can paste either a compact one-line JSON object or multi-line JSON.
For multi-line input, paste the JSON and then submit an empty line.
Type `quit` to exit.
""".strip()


def print_summary(observation: dict) -> None:
    print(
        f"Month {observation['meta']['month']} | "
        f"budget={observation['treasury']['district_budget']} | "
        f"reserve={observation['treasury']['buffer_reserve']} | "
        f"hours={observation['meta']['hours_remaining']} | "
        f"event={observation['meta']['global_event']}"
    )
    for node in observation["nodes"]:
        print(
            f"- {node['village']}: reported={node['reported_demand']} "
            f"lag={node['report_lag_days']}d biometric={node['biometric_signal']} "
            f"stability={node['stability_index']}"
        )
    if observation["high_risk_queue"]:
        print("High-risk queue:")
        for item in observation["high_risk_queue"]:
            print(
                f"  {item['transfer_id']} {item['village']} amount={item['amount']} "
                f"flag={item['oversight_flag']}"
            )


def read_json_payload() -> str | None:
    first = input("> ").strip()
    if first.lower() in {"quit", "exit"}:
        return None
    if not first:
        return ""

    lines = [first]
    while True:
        candidate = "\n".join(lines)
        try:
            json.loads(candidate)
            return candidate
        except json.JSONDecodeError:
            try:
                next_line = input().rstrip()
            except EOFError:
                return candidate
            if not next_line.strip():
                return candidate
            lines.append(next_line)


def main() -> None:
    scenario = sys.argv[1] if len(sys.argv) > 1 else "calm_year"
    env = BDOEnvironment(scenario=scenario)
    observation = env.reset()

    print(f"Scenario: {scenario}")
    print("Episode info:", env.state.model_dump(mode="json", exclude_none=True))

    while not observation.done:
        print("\nCurrent district snapshot:")
        print_summary(observation.model_dump(mode="json", exclude_none=True))
        print()
        print(PROMPT)
        raw = read_json_payload()
        if raw is None:
            break

        try:
            json.loads(raw)
        except json.JSONDecodeError as exc:
            print(f"Invalid JSON: {exc}")
            continue

        observation = env.step(BDOAction.model_validate_json(raw))
        step_info = observation.info or observation.metadata
        print(f"Reward: {float(observation.reward or 0.0):.4f}")
        print("Executed actions:")
        for item in step_info["executed_actions"]:
            print(f"  - {item}")
        if step_info["truncated_actions"]:
            print("Truncated actions:", step_info["truncated_actions"])
        if step_info["errors"]:
            print("Errors:", step_info["errors"])
        print("Reward breakdown:", step_info["reward_breakdown"])


if __name__ == "__main__":
    main()