from __future__ import annotations

import argparse
import json
import sys
import time
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, TextIO

from .bench_connect_lcd import (
    BenchLogger,
    FrameDetector,
    add_serial_format_args,
    _import_serial,
    open_device_serial,
    _read_for,
    _relay_command,
    _relay_settle,
    _send_frame,
    _wait_for_ready,
    format_frame,
    frame_checksum_ok,
    parse_frame,
    serial_format_label,
)
from .serial_table_dump import build_read_frame, decode_table_read_response


DEFAULT_ACK_TARGET = bytes.fromhex("07804020902D")
DEFAULT_ACK_FRAME = bytes.fromhex("05004000001F")


@dataclass
class ScenarioContext:
    args: argparse.Namespace
    logger: BenchLogger
    detector: FrameDetector
    device: Any
    relay: Any | None = None
    table_rows: list[dict[str, Any]] = field(default_factory=list)
    target_counts: dict[str, int] = field(default_factory=dict)
    tx_records: list[dict[str, Any]] = field(default_factory=list)
    ack_sent: int = 0
    abort_requested: bool = False


def default_log_path(scenario: dict[str, Any]) -> Path:
    name = str(scenario.get("name") or "serial-scenario").strip() or "serial-scenario"
    safe_name = "".join(char if char.isalnum() or char in "-_" else "-" for char in name)
    return Path("captures") / f"{safe_name}-{datetime.now().strftime('%Y%m%d-%H%M%S')}.txt"


def build_arg_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Run JSON-described serial bench scenarios against the real RCP."
    )
    parser.add_argument("scenario", type=Path, help="JSON scenario file")
    parser.add_argument("--port", default="COM5", help="RS232 serial port connected to the RCP")
    parser.add_argument("--baud", type=int, default=38400, help="RCP serial baud rate")
    add_serial_format_args(parser)
    parser.add_argument("--relay-port", default="COM6", help="Pico relay serial port")
    parser.add_argument("--relay-baud", type=int, default=115200, help="Pico relay serial baud rate")
    parser.add_argument("--no-power-cycle", action="store_true", help="skip power_cycle actions")
    parser.add_argument("--power-off-command", default="off", help="relay command used to remove DUT power")
    parser.add_argument("--power-on-command", default="on", help="relay command used to apply DUT power")
    parser.add_argument("--relay-settle", type=float, default=2.0, help="seconds to wait after opening the relay port")
    parser.add_argument("--sync", choices=("checksum", "fixed"), default="checksum", help="RX frame sync strategy")
    parser.add_argument("--log", type=Path, help="capture log path")
    parser.add_argument("--result-json", type=Path, help="write machine-readable scenario summary")
    parser.add_argument("--dry-run", action="store_true", help="print the plan without opening serial ports")
    return parser


def main(argv: list[str] | None = None, *, stdout: TextIO = sys.stdout) -> int:
    args = build_arg_parser().parse_args(argv)
    scenario = load_scenario(args.scenario)
    log_path = args.log or default_log_path(scenario)

    if args.dry_run:
        _print_dry_run(args, scenario, log_path, stdout)
        return 0

    serial = _import_serial()
    logger = BenchLogger(log_path, stdout=stdout)
    detector = FrameDetector(sync_mode=args.sync)
    try:
        logger.emit("Serial bench scenario")
        logger.emit(f"name={scenario.get('name', args.scenario.stem)}")
        logger.emit(
            f"device={args.port} {args.baud} {serial_format_label(args)} "
            f"relay={args.relay_port} {args.relay_baud} sync={args.sync}"
        )
        logger.emit(f"log={log_path}")
        with open_device_serial(serial, args) as device:
            ctx = ScenarioContext(args=args, logger=logger, detector=detector, device=device)
            try:
                for index, step in enumerate(_scenario_steps(scenario), start=1):
                    if ctx.abort_requested:
                        logger.event("SCENARIO_ABORT requested by prior step")
                        break
                    action, spec = _normalize_step(step)
                    logger.event(f"STEP {index} {action}")
                    _run_step(ctx, action, spec)
            finally:
                if ctx.relay is not None:
                    ctx.relay.close()
        _emit_summary(ctx, logger)
        if args.result_json:
            _write_result_json(args.result_json, scenario, log_path, ctx)
        return 0
    finally:
        logger.close()


def load_scenario(path: Path) -> dict[str, Any]:
    with path.open("r", encoding="utf-8") as handle:
        scenario = json.load(handle)
    if not isinstance(scenario, dict):
        raise SystemExit("scenario root must be a JSON object")
    if not isinstance(scenario.get("steps"), list):
        raise SystemExit("scenario must contain a steps array")
    return scenario


def _scenario_steps(scenario: dict[str, Any]) -> list[Any]:
    steps = scenario.get("steps", [])
    if not isinstance(steps, list):
        raise SystemExit("scenario steps must be an array")
    return steps


def _normalize_step(step: Any) -> tuple[str, dict[str, Any]]:
    if isinstance(step, str):
        return step, {}
    if not isinstance(step, dict):
        raise SystemExit(f"invalid scenario step {step!r}")
    if "action" in step:
        spec = dict(step)
        action = str(spec.pop("action"))
        return action, spec
    if len(step) == 1:
        action, value = next(iter(step.items()))
        if value is None:
            return str(action), {}
        if isinstance(value, dict):
            return str(action), dict(value)
        return str(action), {"value": value}
    raise SystemExit(f"scenario step needs an action: {step!r}")


def _run_step(ctx: ScenarioContext, action: str, spec: dict[str, Any]) -> None:
    if action == "power_cycle":
        _step_power_cycle(ctx, spec)
    elif action == "wait_ready":
        ready = _wait_for_ready(
            ctx.device,
            ctx.detector,
            ctx.logger,
            float(spec.get("timeout", 10.0)),
            int(spec.get("heartbeats", 2)),
        )
        if spec.get("require", False) and not ready:
            raise SystemExit(2)
    elif action in {"drain", "listen", "wait"}:
        _listen(ctx, float(spec.get("seconds", spec.get("value", 0.0))))
    elif action == "send":
        frame = _parse_required_frame(spec.get("frame"))
        label = str(spec.get("label", "send"))
        _send_and_record(ctx, frame, label)
        if float(spec.get("listen", 0.0)) > 0:
            _listen(ctx, float(spec.get("listen", 0.0)))
    elif action == "wait_for":
        _step_wait_for(ctx, spec)
    elif action == "table_sweep":
        _step_table_sweep(ctx, spec)
    elif action == "repeat":
        _step_repeat(ctx, spec)
    else:
        raise SystemExit(f"unknown scenario action {action!r}")


def _step_power_cycle(ctx: ScenarioContext, spec: dict[str, Any]) -> None:
    if ctx.args.no_power_cycle:
        ctx.logger.event("POWER_CYCLE skipped by --no-power-cycle")
        ctx.device.reset_input_buffer()
        ctx.detector = FrameDetector(sync_mode=ctx.args.sync)
        return
    serial = _import_serial()
    if ctx.relay is None:
        ctx.relay = serial.Serial(ctx.args.relay_port, ctx.args.relay_baud, timeout=0.25)
        _relay_settle(ctx.relay, float(spec.get("relay_settle", ctx.args.relay_settle)), ctx.logger)
    off_command = str(spec.get("off_command", ctx.args.power_off_command))
    on_command = str(spec.get("on_command", ctx.args.power_on_command))
    _relay_command(ctx.relay, off_command, ctx.logger)
    time.sleep(float(spec.get("off_seconds", 1.5)))
    ctx.device.reset_input_buffer()
    ctx.detector = FrameDetector(sync_mode=ctx.args.sync)
    _relay_command(ctx.relay, on_command, ctx.logger)


def _step_wait_for(ctx: ScenarioContext, spec: dict[str, Any]) -> None:
    targets = _parse_frame_list(spec.get("frames", spec.get("frame")))
    timeout = float(spec.get("timeout", 1.0))
    require = bool(spec.get("require", False))
    ctx.logger.event(
        "WAIT_FOR "
        + ",".join(format_frame(frame) for frame in targets)
        + f" timeout={timeout:.3f}s"
    )
    found = _listen_until(ctx, timeout, targets)
    if require and not found:
        raise SystemExit(3)


def _step_repeat(ctx: ScenarioContext, spec: dict[str, Any]) -> None:
    count = max(0, int(spec.get("count", 1)))
    steps = spec.get("steps", [])
    if not isinstance(steps, list):
        raise SystemExit("repeat step requires a steps array")
    for repeat_index in range(count):
        ctx.logger.event(f"REPEAT {repeat_index + 1}/{count}")
        for step in steps:
            action, child_spec = _normalize_step(step)
            _run_step(ctx, action, child_spec)


def _step_table_sweep(ctx: ScenarioContext, spec: dict[str, Any]) -> None:
    selectors = _selector_list(spec)
    gap = float(spec.get("gap", 0.080))
    ack = _ack_config(spec.get("ack_on", {}))
    ack_note = (
        "ack=disabled"
        if not ack["enabled"]
        else f"ack_targets={len(ack['targets'])} ack_frame={format_frame(ack['frame'])}"
    )
    ctx.logger.event(
        f"TABLE_SWEEP selectors={len(selectors)} gap={gap:.3f}s {ack_note}"
    )
    for selector in selectors:
        if ctx.abort_requested:
            ctx.logger.event("TABLE_SWEEP_ABORT stopping before next selector")
            break
        frame = build_read_frame(selector)
        ctx.logger.event(f"READ selector=0x{selector:03X} frame={format_frame(frame)}")
        _send_and_record(ctx, frame, f"read_0x{selector:03X}")
        _listen_with_ack(ctx, gap, selector, ack)


def _ack_config(raw: Any) -> dict[str, Any]:
    spec = raw if isinstance(raw, dict) else {}
    enabled = bool(spec.get("enabled", True))
    targets = _parse_frame_list(spec.get("frames", spec.get("frame", DEFAULT_ACK_TARGET)))
    if not enabled:
        targets = set()
    return {
        "targets": set(targets),
        "frame": _parse_optional_frame(spec.get("ack_frame"), DEFAULT_ACK_FRAME),
        "guard": float(spec.get("ack_guard", 0.020)),
        "poll_interval": float(spec.get("poll_interval", 0.005)),
        "post_read": float(spec.get("post_ack_read", 0.250)),
        "once_per_selector": bool(spec.get("once_per_selector", True)),
        "enabled": enabled,
        "max_acks": _optional_int(spec.get("max_acks")),
        "max_target_hits": _optional_int(spec.get("max_target_hits")),
        "abort_on_limit": bool(spec.get("abort_on_limit", True)),
    }


def _listen_with_ack(
    ctx: ScenarioContext,
    seconds: float,
    selector: int,
    ack: dict[str, Any],
) -> list[bytes]:
    deadline = time.monotonic() + max(0.0, seconds)
    observed: list[bytes] = []
    acked_targets: set[bytes] = set()
    while time.monotonic() < deadline:
        frames = _read_available(ctx, selector=selector)
        observed.extend(frames)
        if not frames:
            sleep_for = min(max(0.001, ack["poll_interval"]), max(0.0, deadline - time.monotonic()))
            if sleep_for > 0:
                time.sleep(sleep_for)
            continue
        if not ack["enabled"]:
            continue
        for frame in frames:
            if frame not in ack["targets"]:
                continue
            _count_target(ctx, frame)
            if _ack_limit_reached(ctx, ack):
                ctx.logger.event("ACK_LIMIT reached before ACK send")
                if ack["abort_on_limit"]:
                    ctx.abort_requested = True
                    return observed
                continue
            if ack["once_per_selector"] and frame in acked_targets:
                continue
            acked_targets.add(frame)
            if ack["guard"] > 0:
                observed.extend(_listen(ctx, ack["guard"], selector=selector))
            _send_and_record(ctx, ack["frame"], "ack")
            ctx.ack_sent += 1
            if _ack_limit_reached(ctx, ack):
                ctx.logger.event("ACK_LIMIT reached after ACK send")
                if ack["abort_on_limit"]:
                    ctx.abort_requested = True
            if ack["post_read"] > 0:
                observed.extend(_listen(ctx, ack["post_read"], selector=selector))
            if ctx.abort_requested:
                return observed
    return observed


def _listen_until(ctx: ScenarioContext, seconds: float, targets: set[bytes]) -> bool:
    deadline = time.monotonic() + max(0.0, seconds)
    while time.monotonic() < deadline:
        interval = min(0.050, max(0.0, deadline - time.monotonic()))
        if interval <= 0:
            break
        for frame in _listen(ctx, interval):
            if frame in targets:
                ctx.logger.event(f"WAIT_FOR_MATCH {format_frame(frame)}")
                return True
    return False


def _listen(ctx: ScenarioContext, seconds: float, *, selector: int | None = None) -> list[bytes]:
    before = len(ctx.detector.frames)
    _read_for(ctx.device, ctx.detector, ctx.logger, seconds)
    frames = ctx.detector.frames[before:]
    _record_table_rows(ctx, frames, selector)
    return frames


def _read_available(ctx: ScenarioContext, *, selector: int | None = None) -> list[bytes]:
    waiting = getattr(ctx.device, "in_waiting", 0)
    if not waiting:
        return []
    dropped_before = ctx.detector.dropped_bytes
    data = ctx.device.read(waiting)
    if not data:
        return []
    ctx.logger.chunk("RX", data)
    detected = ctx.detector.feed(data)
    for frame, label in detected:
        ctx.logger.event(f"DETECT {label} {format_frame(frame)}")
    dropped_now = ctx.detector.dropped_bytes - dropped_before
    if dropped_now:
        ctx.logger.event(
            f"RESYNC dropped_bytes={dropped_now} total_dropped={ctx.detector.dropped_bytes} "
            f"buffered={len(ctx.detector.buffer)}"
        )
    frames = [frame for frame, _label in detected]
    _record_table_rows(ctx, frames, selector)
    return frames


def _record_table_rows(ctx: ScenarioContext, frames: list[bytes], selector: int | None) -> None:
    for frame in frames:
        decoded = decode_table_read_response(frame)
        if decoded is None or selector is None:
            continue
        echo, value = decoded
        row = {
            "selector": selector,
            "echo": echo,
            "value": value,
            "frame": format_frame(frame),
        }
        ctx.table_rows.append(row)
        ctx.logger.event(f"TABLE selector=0x{selector:03X} echo={echo:02X} value={value:04X}")


def _send_and_record(ctx: ScenarioContext, frame: bytes, label: str) -> None:
    _send_frame(ctx.device, frame, ctx.logger, label)
    ctx.tx_records.append(
        {
            "label": label,
            "frame": format_frame(frame),
            "checksum_ok": frame_checksum_ok(frame),
        }
    )


def _count_target(ctx: ScenarioContext, frame: bytes) -> None:
    text = format_frame(frame)
    ctx.target_counts[text] = ctx.target_counts.get(text, 0) + 1
    ctx.logger.event(f"ACK_TARGET {text} count={ctx.target_counts[text]}")


def _selector_list(spec: dict[str, Any]) -> list[int]:
    if "selectors" in spec:
        raw_selectors = spec["selectors"]
        if not isinstance(raw_selectors, list):
            raise SystemExit("table_sweep selectors must be an array")
        return [_int_value(selector) & 0x01FF for selector in raw_selectors]
    start = _int_value(spec.get("start", 0))
    count = max(0, _int_value(spec.get("count", 0x80)))
    return [((start + offset) & 0x01FF) for offset in range(count)]


def _parse_frame_list(raw: Any) -> set[bytes]:
    if raw is None:
        return set()
    values = raw if isinstance(raw, list) else [raw]
    return {_parse_required_frame(value) for value in values}


def _parse_required_frame(raw: Any) -> bytes:
    if raw is None:
        raise SystemExit("frame is required")
    if isinstance(raw, bytes):
        return raw
    if not isinstance(raw, str):
        raise SystemExit(f"frame must be a hex string, got {raw!r}")
    return parse_frame(raw)


def _parse_optional_frame(raw: Any, default: bytes) -> bytes:
    if raw is None:
        return default
    return _parse_required_frame(raw)


def _int_value(raw: Any) -> int:
    if isinstance(raw, int):
        return raw
    if isinstance(raw, str):
        return int(raw, 0)
    raise SystemExit(f"expected integer value, got {raw!r}")


def _optional_int(raw: Any) -> int | None:
    if raw is None:
        return None
    return _int_value(raw)


def _ack_limit_reached(ctx: ScenarioContext, ack: dict[str, Any]) -> bool:
    max_acks = ack.get("max_acks")
    if max_acks is not None and ctx.ack_sent >= max_acks:
        return True
    max_target_hits = ack.get("max_target_hits")
    if max_target_hits is not None and sum(ctx.target_counts.values()) >= max_target_hits:
        return True
    return False


def _print_dry_run(args: argparse.Namespace, scenario: dict[str, Any], log_path: Path, stdout: TextIO) -> None:
    print(f"scenario={scenario.get('name', args.scenario.stem)}", file=stdout)
    print(f"device={args.port} {args.baud} {serial_format_label(args)}", file=stdout)
    print(f"relay={args.relay_port} {args.relay_baud}", file=stdout)
    print(f"sync={args.sync}", file=stdout)
    print(f"log={log_path}", file=stdout)
    for index, step in enumerate(_scenario_steps(scenario), start=1):
        action, spec = _normalize_step(step)
        print(f"step[{index}]={action}", file=stdout)
        _print_step_dry_run(action, spec, stdout)


def _print_step_dry_run(action: str, spec: dict[str, Any], stdout: TextIO, *, indent: str = "  ") -> None:
    if action == "send":
        frame = _parse_required_frame(spec.get("frame"))
        print(f"{indent}frame={format_frame(frame)} checksum_ok={int(frame_checksum_ok(frame))}", file=stdout)
        if float(spec.get("listen", 0.0)) > 0:
            print(f"{indent}listen={float(spec.get('listen', 0.0)):.3f}s", file=stdout)
    elif action in {"drain", "listen", "wait"}:
        print(f"{indent}seconds={float(spec.get('seconds', spec.get('value', 0.0))):.3f}", file=stdout)
    elif action == "wait_ready":
        print(
            f"{indent}heartbeats={int(spec.get('heartbeats', 2))} "
            f"timeout={float(spec.get('timeout', 10.0)):.3f}s "
            f"require={int(bool(spec.get('require', False)))}",
            file=stdout,
        )
    elif action == "table_sweep":
        selectors = _selector_list(spec)
        ack = _ack_config(spec.get("ack_on", {}))
        if selectors:
            first = selectors[0]
            last = selectors[-1]
            print(f"{indent}selectors={len(selectors)} first=0x{first:03X} last=0x{last:03X}", file=stdout)
        else:
            print(f"{indent}selectors=0", file=stdout)
        print(f"{indent}gap={float(spec.get('gap', 0.080)):.3f}", file=stdout)
        if not ack["enabled"]:
            print(f"{indent}ack=disabled", file=stdout)
        else:
            for target in sorted(ack["targets"]):
                print(f"{indent}ack_target={format_frame(target)}", file=stdout)
            print(f"{indent}ack_frame={format_frame(ack['frame'])}", file=stdout)
            print(f"{indent}max_acks={ack['max_acks']} max_target_hits={ack['max_target_hits']}", file=stdout)
    elif action == "repeat":
        count = max(0, int(spec.get("count", 1)))
        steps = spec.get("steps", [])
        step_count = len(steps) if isinstance(steps, list) else 0
        print(f"{indent}count={count} steps={step_count}", file=stdout)
        if not isinstance(steps, list):
            return
        for child_index, child in enumerate(steps, start=1):
            child_action, child_spec = _normalize_step(child)
            print(f"{indent}child[{child_index}]={child_action}", file=stdout)
            _print_step_dry_run(child_action, child_spec, stdout, indent=indent + "  ")


def _emit_summary(ctx: ScenarioContext, logger: BenchLogger) -> None:
    logger.emit()
    logger.emit("Summary")
    logger.emit(f"rx_frames={len(ctx.detector.frames)} trailing_unframed_bytes={len(ctx.detector.buffer)}")
    logger.emit(f"resync_events={ctx.detector.resync_events} dropped_bytes={ctx.detector.dropped_bytes}")
    logger.emit(f"tx_frames={len(ctx.tx_records)} ack_sent={ctx.ack_sent} table_response_rows={len(ctx.table_rows)}")
    logger.emit(f"abort_requested={int(ctx.abort_requested)}")
    for target, count in sorted(ctx.target_counts.items()):
        logger.emit(f"ack_target {target}={count}")
    for label, count in sorted(ctx.detector.labels.items()):
        logger.emit(f"{label}={count}")
    for row in ctx.table_rows:
        logger.emit(
            f"table selector=0x{row['selector']:03X} echo=0x{row['echo']:02X} value=0x{row['value']:04X}"
        )


def _write_result_json(path: Path, scenario: dict[str, Any], log_path: Path, ctx: ScenarioContext) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    result = {
        "scenario": scenario.get("name", ""),
        "log": str(log_path),
        "rx_frames": len(ctx.detector.frames),
        "trailing_unframed_bytes": len(ctx.detector.buffer),
        "resync_events": ctx.detector.resync_events,
        "dropped_bytes": ctx.detector.dropped_bytes,
        "labels": dict(ctx.detector.labels),
        "tx_frames": ctx.tx_records,
        "ack_sent": ctx.ack_sent,
        "abort_requested": ctx.abort_requested,
        "ack_targets": ctx.target_counts,
        "table_rows": ctx.table_rows,
    }
    path.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8")


__all__ = [
    "DEFAULT_ACK_FRAME",
    "DEFAULT_ACK_TARGET",
    "build_arg_parser",
    "load_scenario",
    "main",
]
