from __future__ import annotations

import argparse
import importlib
import json
import time
from pathlib import Path

from masque import LibraryError


READERS = {
    'gdsii': ('masque.file.gdsii', 'readfile'),
    'gdsii_arrow': ('masque.file.gdsii_arrow', 'readfile'),
    }


def _summarize(path: Path, elapsed_s: float, info: dict[str, object], lib: object) -> dict[str, object]:
    assert hasattr(lib, '__len__')
    assert hasattr(lib, 'tops')
    tops = lib.tops()      # type: ignore[no-any-return, attr-defined]
    try:
        unique_top = lib.top()      # type: ignore[no-any-return, attr-defined]
    except LibraryError:
        unique_top = None

    return {
        'path': str(path),
        'elapsed_s': elapsed_s,
        'library_name': info['name'],
        'cell_count': len(lib),     # type: ignore[arg-type]
        'topcells': tops,
        'topcell': unique_top,
        }


def build_arg_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description='Profile GDS readers with a stable end-to-end workload.')
    parser.add_argument('--reader', choices=sorted(READERS), required=True)
    parser.add_argument('--path', type=Path, required=True)
    parser.add_argument('--warmup', type=int, default=1)
    parser.add_argument('--repeat', type=int, default=1)
    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_arg_parser()
    args = parser.parse_args(argv)

    module_name, attr_name = READERS[args.reader]
    readfile = getattr(importlib.import_module(module_name), attr_name)
    path = args.path.expanduser().resolve()

    for _ in range(args.warmup):
        readfile(path)

    runs = []
    for _ in range(args.repeat):
        start = time.perf_counter()
        lib, info = readfile(path)
        elapsed_s = time.perf_counter() - start
        runs.append(_summarize(path, elapsed_s, info, lib))

    print(json.dumps({
        'reader': args.reader,
        'warmup': args.warmup,
        'repeat': args.repeat,
        'runs': runs,
        }, indent=2, sort_keys=True))
    return 0


if __name__ == '__main__':
    raise SystemExit(main())