from __future__ import annotations

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

from masque import LibraryError


READERS: dict[str, tuple[str, tuple[str, ...]]] = {
    'gdsii': ('masque.file.gdsii', ('readfile',)),
    'gdsii_arrow': ('masque.file.gdsii_arrow', ('readfile', 'arrow_import', 'arrow_convert')),
    }


def _summarize_library(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 _summarize_arrow_import(path: Path, elapsed_s: float, arrow_arr: Any) -> dict[str, object]:
    libarr = arrow_arr[0]
    return {
        'path': str(path),
        'elapsed_s': elapsed_s,
        'arrow_rows': len(arrow_arr),
        'library_name': libarr['lib_name'].as_py(),
        'cell_count': len(libarr['cells']),
        'layer_count': len(libarr['layers']),
        }


def _profile_stage(module: Any, stage: str, path: Path) -> dict[str, object]:
    start = time.perf_counter()

    if stage == 'readfile':
        lib, info = module.readfile(path)
        elapsed_s = time.perf_counter() - start
        return _summarize_library(path, elapsed_s, info, lib)

    if stage == 'arrow_import':
        arrow_arr = module._read_to_arrow(path)
        elapsed_s = time.perf_counter() - start
        return _summarize_arrow_import(path, elapsed_s, arrow_arr)

    if stage == 'arrow_convert':
        arrow_arr = module._read_to_arrow(path)
        libarr = arrow_arr[0]
        start = time.perf_counter()
        lib, info = module.read_arrow(libarr)
        elapsed_s = time.perf_counter() - start
        return _summarize_library(path, elapsed_s, info, lib)

    raise ValueError(f'Unsupported stage {stage!r}')


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('--stage', default='readfile')
    parser.add_argument('--path', type=Path, required=True)
    parser.add_argument('--warmup', type=int, default=1)
    parser.add_argument('--repeat', type=int, default=1)
    parser.add_argument('--output-json', type=Path)
    return parser


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

    module_name, stages = READERS[args.reader]
    if args.stage not in stages:
        parser.error(f'reader {args.reader!r} only supports stages: {", ".join(stages)}')

    module = importlib.import_module(module_name)
    path = args.path.expanduser().resolve()

    for _ in range(args.warmup):
        _profile_stage(module, args.stage, path)

    runs = []
    for _ in range(args.repeat):
        runs.append(_profile_stage(module, args.stage, path))

    payload = {
        'reader': args.reader,
        'stage': args.stage,
        'warmup': args.warmup,
        'repeat': args.repeat,
        'runs': runs,
        }
    rendered = json.dumps(payload, indent=2, sort_keys=True)
    if args.output_json is not None:
        args.output_json.parent.mkdir(parents=True, exist_ok=True)
        args.output_json.write_text(rendered + '\n')
    print(rendered)
    return 0


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