# retoor <retoor@molodetz.nl>
import argparse
import json

COMPLEXITY_SHARE = {"simple": 0.35, "moderate": 0.65, "complex": 1.0}
EXTENSION_SHARE = 0.25
MAX_ROUNDS = 8
MIN_QUERIES = 2
MAX_QUERIES = 5


def schedule(complexity, budget_seconds, max_rounds=MAX_ROUNDS):
    share = COMPLEXITY_SHARE[complexity]
    rounds = max(1, min(max_rounds, round(max_rounds * share)))
    allowance = budget_seconds * share
    per_round = allowance / rounds
    return {"complexity": complexity, "rounds": rounds, "allowance": allowance, "per_round": per_round}


def plan(queries, seen_urls, max_queries=MAX_QUERIES, min_queries=MIN_QUERIES):
    fresh = [query for query in queries if query not in seen_urls]
    if len(fresh) < min_queries:
        raise ValueError(f"need at least {min_queries} fresh follow-up queries, got {len(fresh)}")
    return fresh[:max_queries]


def should_extend(improved, extended_already, rounds_used):
    return bool(improved) and not extended_already and rounds_used < MAX_ROUNDS


def next_step(rounds_used, rounds_planned, improved, extended_already):
    if rounds_used >= MAX_ROUNDS:
        return "answer"
    if should_extend(improved, extended_already, rounds_used):
        return "extend"
    if rounds_used >= rounds_planned:
        return "answer"
    return "continue"


def state(complexity, budget_seconds, rounds_used=0, improved=False, extended_already=False):
    booked = schedule(complexity, budget_seconds)
    extend = should_extend(improved, extended_already, rounds_used)
    extra = budget_seconds * EXTENSION_SHARE if extend else 0.0
    return {
        "rounds_used": rounds_used,
        "rounds_planned": booked["rounds"],
        "allowance": booked["allowance"] + extra,
        "extended": extended_already or extend,
        "decision": next_step(rounds_used, booked["rounds"], improved, extended_already),
    }


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Plan follow-up iteration within round and wall-clock caps.")
    parser.add_argument("--complexity", choices=sorted(COMPLEXITY_SHARE), default="moderate")
    parser.add_argument("--budget", type=float, default=600.0, help="Total wall-clock seconds for research")
    parser.add_argument("--rounds-used", type=int, default=0, help="Rounds already completed")
    parser.add_argument("--improved", action="store_true", help="Last round improved the corpus")
    parser.add_argument("--extended", action="store_true", help="The 25 percent extension was already granted")
    parser.add_argument("--queries", nargs="*", default=[], help="Candidate follow-up queries to validate")
    parser.add_argument("--seen", nargs="*", default=[], help="Queries already asked")
    args = parser.parse_args()
    payload = state(args.complexity, args.budget, args.rounds_used, args.improved, args.extended)
    if args.queries:
        payload["queries"] = plan(args.queries, set(args.seen))
    print(json.dumps(payload))