"""Benchmark: lazy Sox loading vs seeding the cache from app.sab.

    python bench_sab_seed.py HOST [--port 1876] [--user admin] [--password PW]

Each method runs on a fresh connection, so it starts from an empty cache.

  lazy-seq    every component's tree, one request at a time (as used_kit_names does)
  lazy-tree   the tree folder by folder, each folder's children in one window
  lazy-full   lazy-tree plus config and links of every component: what app.sab holds
  lazy-full-seq  lazy-seq plus config and links, one request at a time
  sab         App.save, get app.sab, decode it
"""

import argparse
import time

import app_state  # noqa: F401  loads Sedona home
import schema
from device_session import DeviceSession
from sox import CONFIG
import app_backup
import sab_app


class Counter:
    """Counts Sox requests sent through the client."""

    def __init__(self, client):
        self.n = 0
        orig_requests = client.requests

        def requests(reqs, *a, **k):
            self.n += len(reqs)
            return orig_requests(reqs, *a, **k)

        client.requests = requests


def connect(args):
    ds = DeviceSession()
    # DeviceSession.connect also loads the root's children and their values; time
    # only the bare session here, so every method starts from the same place.
    from sox import SoxClient
    ds.sox_client = SoxClient(args.host, args.port, args.user, args.password)
    ds.sox_client.checksum_resolver = schema.resolve_truncated_checksum
    ds.sox_client.connect()
    parts = ds.sox_client.read_schema()
    ds.sox_client.read_version()
    ds._load_slot_registry(parts)
    return ds


def lazy_seq(ds):
    client = ds.sox_client
    pending, n = [0], 0
    while pending:
        comp = client.load_ids([pending.pop()], checked=False)[0]
        if comp is not None:
            n += 1
            pending.extend(comp.children)
    return n


def lazy_tree(ds):
    client = ds.sox_client
    client.load(0)
    pending, comps = [0], []
    while pending:
        cid = pending.pop()
        comps.append(client.cache_get(cid))
        for child in ds.load_children(cid):
            if child.children:
                pending.append(child.id)
            else:
                comps.append(child)
    return [c for c in comps if c is not None]


def lazy_full(ds):
    client = ds.sox_client
    comps = lazy_tree(ds)
    # In folder-sized windows, as open_wiresheet sends them.
    for i in range(0, len(comps), 16):
        client.update(comps[i:i + 16], CONFIG)
    for c in comps:
        c.links = client.read_links(c.id)
    return len(comps)


def lazy_full_seq(ds):
    """lazy-seq plus config and links, every request alone: a P-ION's pace."""
    client = ds.sox_client
    pending, n = [0], 0
    while pending:
        comp = client.load_ids([pending.pop()], checked=False)[0]
        if comp is None:
            continue
        n += 1
        pending.extend(comp.children)
        client.update([comp], CONFIG)
        comp.links = client.read_links(comp.id)
    return n


def sab(ds, timing):
    client = ds.sox_client
    t = time.perf_counter()
    app_backup.save_app(client)
    timing["save"] = time.perf_counter() - t
    t = time.perf_counter()
    data = client.get_file(app_backup.APP_SAB)
    timing["get"] = time.perf_counter() - t
    timing["bytes"] = len(data)
    t = time.perf_counter()
    app = sab_app.SabApp(data)
    timing["decode"] = time.perf_counter() - t
    timing["links"] = len(app.links)
    return len(app.comps)


def run(name, args, fn):
    try:
        ds = connect(args)
    except Exception as e:
        print(f"{name:10s}  connect failed: {e}")
        return
    counter = Counter(ds.sox_client)
    timing = {}
    t = time.perf_counter()
    try:
        result = fn(ds, timing) if name == "sab" else fn(ds)
        elapsed = time.perf_counter() - t
        n = result if isinstance(result, int) else len(result)
        extra = "  ".join(f"{k}={v:.2f}s" if isinstance(v, float) else f"{k}={v}"
                          for k, v in timing.items())
        print(f"{name:10s}  {elapsed:7.2f}s  comps={n:4d}  requests={counter.n:5d}  {extra}")
    except Exception as e:
        print(f"{name:10s}  FAILED after {time.perf_counter() - t:.2f}s "
              f"({counter.n} requests): {e}")
    finally:
        ds.disconnect()
        time.sleep(1.0)


def main():
    p = argparse.ArgumentParser()
    p.add_argument("host")
    p.add_argument("--port", type=int, default=1876)
    p.add_argument("--user", default="admin")
    p.add_argument("--password", default="")
    p.add_argument("--only", nargs="*",
                   default=["lazy-seq", "lazy-tree", "lazy-full", "sab"])
    args = p.parse_args()
    fns = {"lazy-seq": lazy_seq, "lazy-tree": lazy_tree,
           "lazy-full": lazy_full, "lazy-full-seq": lazy_full_seq, "sab": sab}
    print(f"== {args.host}:{args.port}")
    for name in args.only:
        run(name, args, fns[name])


if __name__ == "__main__":
    main()