#!/usr/bin/env python3
"""Run selectable SeQUeNCe experiments against selectable topology JSON files."""

import argparse
import math
from pathlib import Path

from sequence.app.random_request import RandomRequestApp
from sequence.app.request_app import RequestApp
from sequence.app.teleport_app import TeleportApp
from sequence.constants import MILLISECOND
from sequence.topology.dqc_net_topo import DQCNetTopo
from sequence.topology.router_net_topo import RouterNetTopo
from sequence.constants import SINGLE_HERALDED, SECOND
from sequence.entanglement_management.generation import EntanglementGenerationA, EntanglementGenerationB


SCRIPT_DIR = Path(__file__).resolve().parent
TOPOLOGY_DIR = SCRIPT_DIR / "topologies"


def run_simple_request(args) -> None:
    EntanglementGenerationA.set_global_type(SINGLE_HERALDED)
    EntanglementGenerationB.set_global_type(SINGLE_HERALDED)

    topo = RouterNetTopo(str(TOPOLOGY_DIR / f"{args.topology}.json"))
    routers = topo.get_nodes_by_type(RouterNetTopo.QUANTUM_ROUTER)
    if len(routers) < 2:
        raise RuntimeError("simple_request requires at least two quantum routers")

    app_map = {router.name: RequestApp(router) for router in routers}
    initiator = routers[0].name
    responder = routers[1].name

    tl = topo.get_timeline()
    tl.stop_time = args.stop_time
    tl.init()
    app_map[initiator].start(
        responder=responder,
        start_t=args.request_start,
        end_t=args.request_end,
        memo_size=args.request_memory,
        fidelity=args.request_fidelity,
    )
    tl.run()

    print("RESULT: simple_request")
    print(f"  initiator={initiator} responder={responder}")
    print(f"  pair_count={app_map[initiator].memory_counter}")
    print(f"  throughput_pairs_per_sec={app_map[initiator].get_throughput()}")


def run_random_requests(args) -> None:
    topo = RouterNetTopo(str(TOPOLOGY_DIR / f"{args.topology}.json"))
    routers = topo.get_nodes_by_type(RouterNetTopo.QUANTUM_ROUTER)
    if len(routers) < 2:
        raise RuntimeError("random_requests requires at least two quantum routers")

    selected_count = min(max(args.random_request_count, 1), len(routers))
    selected = routers[:selected_count]
    names = [node.name for node in routers]

    apps = []
    for index, node in enumerate(selected):
        others = [name for name in names if name != node.name]
        app = RandomRequestApp(
            node,
            others,
            seed=index,
            min_dur=max(args.request_end - args.request_start, 10_000_000_000),
            max_dur=max(args.request_end - args.request_start, 10_000_000_000) * 2,
            min_size=max(args.request_memory, 1),
            max_size=max(args.request_memory + 1, 2),
            min_fidelity=min(max(args.request_fidelity, 0.01), 0.99),
            max_fidelity=1.0,
        )
        app.start()
        apps.append(app)

    tl = topo.get_timeline()
    tl.stop_time = args.stop_time
    tl.init()
    tl.run()

    print("RESULT: random_requests")
    print(f"  active_apps={len(apps)}")
    for app in apps:
        waits = app.get_wait_time()
        tputs = app.get_all_throughput()
        completed = min(len(waits), len(tputs))
        avg_wait = sum(waits[:completed]) / completed if completed else 0.0
        avg_tput = sum(tputs[:completed]) / completed if completed else 0.0
        print(
            f"  node={app.node.name} completed={completed} "
            f"avg_wait_ps={avg_wait:.2f} avg_throughput={avg_tput:.6f}"
        )


def run_teleport_demo(args) -> None:
    topo = DQCNetTopo(str(TOPOLOGY_DIR / f"{args.topology}.json"))
    nodes = topo.nodes[DQCNetTopo.DQC_NODE]
    alice = next((n for n in nodes if n.name == "alice"), None)
    bob = next((n for n in nodes if n.name == "bob"), None)
    if alice is None or bob is None:
        raise RuntimeError("teleport_demo requires alice and bob DQC nodes")

    data_memo_arr = alice.get_component_by_name(alice.data_memo_arr_name)
    amp = 1 / math.sqrt(2)
    data_memo_arr[0].update_state([amp, amp])

    alice_app = TeleportApp(alice)
    bob_app = TeleportApp(bob)

    alice_app.start(
        responder=bob.name,
        start_t=max(args.request_start // 1_000_000_000, 1) * MILLISECOND,
        end_t=max(args.request_end // 1_000_000_000, 10) * MILLISECOND,
        memory_size=max(args.request_memory, 1),
        fidelity=min(max(args.request_fidelity, 0.01), 0.99),
        data_memory_index=0,
    )

    tl = topo.tl
    tl.stop_time = args.stop_time
    tl.init()
    tl.run()

    if not bob_app.results:
        raise RuntimeError("Teleportation completed without a Bob-side result")

    result_state = bob_app.results[0][1]
    print("RESULT: teleport_demo")
    print(f"  bob_result_state={result_state}")
    print(f"  total_results={len(bob_app.results)}")


def parse_args():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--experiment", required=True, choices=["simple_request", "random_requests", "teleport_demo"])
    parser.add_argument("--topology", required=True, choices=["two_node", "star_network", "teleport_2node"])
    parser.add_argument("--stop-time", type=int, required=True)
    parser.add_argument("--request-start", type=int, required=True)
    parser.add_argument("--request-end", type=int, required=True)
    parser.add_argument("--request-memory", type=int, required=True)
    parser.add_argument("--request-fidelity", type=float, required=True)
    parser.add_argument("--random-request-count", type=int, required=True)
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    compatibility = {
        "simple_request": {"two_node", "star_network"},
        "random_requests": {"two_node", "star_network"},
        "teleport_demo": {"teleport_2node"},
    }
    if args.topology not in compatibility[args.experiment]:
        allowed = ", ".join(sorted(compatibility[args.experiment]))
        raise SystemExit(
            f"Topology '{args.topology}' is not allowed for experiment '{args.experiment}'. Allowed: {allowed}"
        )

    if args.experiment == "simple_request":
        run_simple_request(args)
    elif args.experiment == "random_requests":
        run_random_requests(args)
    else:
        run_teleport_demo(args)
