from __future__ import annotations

import argparse
import json
import sys
import time
from dataclasses import dataclass
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,
    label_frame,
    parse_frame,
    serial_format_label,
)
from .connect_ok_matrix import FRAME_80_OK


ACK_0040 = bytes.fromhex("05004000001F")
REFRESH_OK = bytes.fromhex("0400008000DE")
ACK_006B = bytes.fromhex("05006B000034")
ACK_006C = bytes.fromhex("05006C000033")
ACK_006D = bytes.fromhex("05006D000032")
ACK_006E = bytes.fromhex("05006E000031")
ACK_0096 = bytes.fromhex("050116000048")
ACK_0097 = bytes.fromhex("050117000049")
ACK_00C6 = bytes.fromhex("050146000018")
ACK_00F8 = bytes.fromhex("050178000026")

DEFAULT_BASELINE = (FRAME_80_OK, FRAME_80_OK)
HEARTBEAT_FRAME = bytes.fromhex("0000000080DA")


@dataclass(frozen=True)
class AdvanceCase:
    name: str
    frame: bytes
    note: str = ""


def default_log_path(suite: str) -> Path:
    safe_suite = "".join(char if char.isalnum() or char in "-_" else "-" for char in suite)
    return Path("captures") / f"connect-ok-advance-{safe_suite}-{datetime.now().strftime('%Y%m%d-%H%M%S')}.txt"


def build_arg_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=(
            "Recover the RCP to CONNECT: OK, then sweep one candidate continuation/ACK "
            "frame from the active report window."
        )
    )
    parser.add_argument("--suite", choices=("core", "special", "latch", "all"), default="core")
    parser.add_argument("--case", action="append", help="run only matching case names; repeatable")
    parser.add_argument("--limit", type=int, help="run only the first N selected cases")
    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="do not power-cycle before the sweep starts")
    parser.add_argument(
        "--power-cycle-between-cases",
        action="store_true",
        help="power-cycle before each candidate instead of recovering from the current state",
    )
    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("--off-seconds", type=float, default=1.5, help="seconds to hold DUT power off")
    parser.add_argument("--relay-settle", type=float, default=2.0, help="seconds to wait after opening the relay port")
    parser.add_argument("--ready-timeout", type=float, default=10.0, help="seconds to wait for heartbeat")
    parser.add_argument("--ready-heartbeats", type=int, default=2, help="heartbeat frames to observe before a case")
    parser.add_argument("--require-ready", action="store_true", help="abort a case if readiness heartbeats are not observed")
    parser.add_argument("--pre-case-drain", type=float, default=0.250, help="seconds to drain/log RX before baseline")
    parser.add_argument(
        "--baseline-frame",
        action="append",
        type=parse_frame,
        help="override baseline with custom frame; repeatable; default is two selector-zero OK frames",
    )
    parser.add_argument("--baseline-gap", type=float, default=0.700, help="seconds to listen between baseline frames")
    parser.add_argument(
        "--target-mode",
        choices=("active", "connect-ok", "non-heartbeat", "none"),
        default="active",
        help="which device frame opens the candidate-send window",
    )
    parser.add_argument("--target-timeout", type=float, default=2.5, help="seconds to wait for the target window")
    parser.add_argument("--candidate-guard", type=float, default=0.020, help="seconds to wait after target before candidate")
    parser.add_argument(
        "--send-on-target-timeout",
        action="store_true",
        help="send the candidate even if no target-mode frame was observed",
    )
    parser.add_argument("--post-candidate-read", type=float, default=5.0, help="seconds to listen after candidate")
    parser.add_argument("--candidate", action="append", type=_parse_candidate, help="custom candidate as name=frame or frame")
    parser.add_argument("--sync", choices=("checksum", "fixed"), default="checksum", help="RX frame sync strategy")
    parser.add_argument("--prompt-observation", action="store_true", help="prompt for observed LCD/lamp state after each case")
    parser.add_argument("--pause-between-cases", action="store_true", help="wait for Enter before the next case")
    parser.add_argument("--log", type=Path, help="capture log path")
    parser.add_argument("--result-json", type=Path, help="write machine-readable case summary")
    parser.add_argument("--dry-run", action="store_true", help="print selected cases 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)
    cases = select_cases(args)
    baseline = tuple(args.baseline_frame or DEFAULT_BASELINE)
    log_path = args.log or default_log_path(args.suite)

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

    serial = _import_serial()
    logger = BenchLogger(log_path, stdout=stdout)
    results: list[dict[str, Any]] = []
    try:
        logger.emit("CONNECT: OK advance sweep")
        logger.emit(
            f"suite={args.suite} cases={len(cases)} 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}")
        _emit_baseline_plan(logger, baseline, args.baseline_gap)
        with open_device_serial(serial, args) as device:
            relay = None
            try:
                if not args.no_power_cycle or args.power_cycle_between_cases:
                    relay = serial.Serial(args.relay_port, args.relay_baud, timeout=0.25)
                    _relay_settle(relay, args.relay_settle, logger)
                if not args.no_power_cycle and not args.power_cycle_between_cases:
                    _power_cycle(args, device, relay, logger)

                for index, case in enumerate(cases, start=1):
                    if args.pause_between_cases and index > 1:
                        input(f"Press Enter to start case {index}/{len(cases)}: {case.name}")
                    result = _run_case(args, device, relay, logger, case, baseline, index, len(cases))
                    results.append(result)
                    if args.require_ready and not result["ready"]:
                        logger.event("ABORT readiness was required")
                        break
            finally:
                if relay is not None:
                    relay.close()
        _emit_summary(logger, results)
        if args.result_json:
            _write_result_json(args.result_json, log_path, args, baseline, results)
        return 0
    finally:
        logger.close()


def select_cases(args: argparse.Namespace) -> list[AdvanceCase]:
    cases = list(args.candidate or build_cases(args.suite))
    if args.case:
        filters = [item.lower() for item in args.case]
        cases = [case for case in cases if any(fragment in case.name.lower() for fragment in filters)]
    if args.limit is not None:
        cases = cases[: max(0, args.limit)]
    if not cases:
        raise SystemExit("no advance cases selected")
    return cases


def build_cases(suite: str) -> list[AdvanceCase]:
    core = [
        AdvanceCase("ack-0040", ACK_0040, "pure command-5 continuation ACK candidate"),
        AdvanceCase("refresh-ok", REFRESH_OK, "command-4 continuation ACK plus selector-zero 0x8080 refresh"),
    ]
    special = [
        AdvanceCase("ack-006c", ACK_006C, "command-5 special BE70 selector candidate"),
        AdvanceCase("ack-006d", ACK_006D, "command-5 special BE70 selector candidate"),
        AdvanceCase("ack-006e", ACK_006E, "command-5 special BE70 selector candidate"),
    ]
    latch = [
        AdvanceCase("ack-006b", ACK_006B, "command-5 latch-clear special selector candidate"),
        AdvanceCase("ack-0096", ACK_0096, "command-5 F731/F790 latch-clear candidate"),
        AdvanceCase("ack-0097", ACK_0097, "command-5 F731/F790 latch-clear candidate"),
        AdvanceCase("ack-00c6", ACK_00C6, "command-5 F731/F790 latch-clear candidate"),
        AdvanceCase("ack-00f8", ACK_00F8, "command-5 F731/F790 latch-clear candidate"),
    ]
    suites = {
        "core": core,
        "special": special,
        "latch": latch,
        "all": core + special + latch,
    }
    return suites[suite]


def _run_case(
    args: argparse.Namespace,
    device: Any,
    relay: Any | None,
    logger: BenchLogger,
    case: AdvanceCase,
    baseline: tuple[bytes, ...],
    index: int,
    total: int,
) -> dict[str, Any]:
    detector = FrameDetector(sync_mode=args.sync)
    logger.emit()
    logger.emit(f"CASE {index}/{total} {case.name}")
    logger.emit(f"note={case.note or '(none)'}")
    logger.emit(f"candidate={format_frame(case.frame)} checksum_ok={int(frame_checksum_ok(case.frame))}")
    logger.emit(
        f"target_mode={args.target_mode} target_timeout={args.target_timeout:.3f}s "
        f"candidate_guard={args.candidate_guard:.3f}s"
    )

    if args.power_cycle_between_cases:
        _power_cycle(args, device, relay, logger)
    else:
        device.reset_input_buffer()
        logger.event("POWER_CYCLE skipped for case; recovering with baseline")

    ready = _wait_for_ready(device, detector, logger, args.ready_timeout, args.ready_heartbeats)
    if args.require_ready and not ready:
        observation = _prompt_observation(args, logger, case)
        return _case_result(case, detector, ready=ready, target=None, candidate_sent=False, observation=observation)

    if args.pre_case_drain > 0:
        logger.event(f"DRAIN before baseline {args.pre_case_drain:.3f}s")
        _read_for(device, detector, logger, args.pre_case_drain)

    for frame_index, frame in enumerate(baseline, start=1):
        _send_frame(device, frame, logger, f"{case.name}.baseline{frame_index}")
        _read_for(device, detector, logger, args.baseline_gap)

    target = _wait_for_target_mode(device, detector, logger, args.target_mode, args.target_timeout)
    candidate_sent = target is not None or args.target_mode == "none" or args.send_on_target_timeout
    if candidate_sent:
        if args.candidate_guard > 0:
            logger.event(f"CANDIDATE guard {args.candidate_guard:.3f}s")
            _read_for(device, detector, logger, args.candidate_guard)
        _send_frame(device, case.frame, logger, f"{case.name}.candidate")
    else:
        logger.event("CANDIDATE skipped because target window was not observed")

    if args.post_candidate_read > 0:
        logger.event(f"POST_CANDIDATE_READ {args.post_candidate_read:.3f}s")
        _read_for(device, detector, logger, args.post_candidate_read)

    observation = _prompt_observation(args, logger, case)
    result = _case_result(case, detector, ready=ready, target=target, candidate_sent=candidate_sent, observation=observation)
    logger.event(
        f"CASE_RESULT {case.name} ready={int(ready)} target_seen={int(target is not None)} "
        f"candidate_sent={int(candidate_sent)} rx_frames={result['rx_frames']} "
        f"labels={json.dumps(result['labels'], sort_keys=True)}"
    )
    return result


def _wait_for_target_mode(
    device: Any,
    detector: FrameDetector,
    logger: BenchLogger,
    mode: str,
    timeout_seconds: float,
) -> bytes | None:
    if mode == "none":
        logger.event("TARGET disabled")
        return None
    logger.event(f"WAIT target_mode={mode} timeout={timeout_seconds:.3f}s")
    start_index = len(detector.frames)
    deadline = time.monotonic() + max(0.0, timeout_seconds)
    while time.monotonic() < deadline:
        before = len(detector.frames)
        _read_for(device, detector, logger, 0.050)
        for frame in detector.frames[max(start_index, before) :]:
            if _matches_target(frame, mode):
                logger.event(f"TARGET seen {format_frame(frame)} label={label_frame(frame)}")
                return frame
    logger.event("TARGET_TIMEOUT")
    return None


def _matches_target(frame: bytes, mode: str) -> bool:
    label = label_frame(frame)
    if mode == "connect-ok":
        return label in {"connect_ok_path_response_candidate", "connect_c0_path_response_candidate"}
    if mode == "non-heartbeat":
        return frame_checksum_ok(frame) and frame != HEARTBEAT_FRAME
    if mode == "active":
        return frame_checksum_ok(frame) and frame != HEARTBEAT_FRAME and frame[0] in {0x01, 0x02, 0x07}
    raise ValueError(f"unknown target mode {mode!r}")


def _power_cycle(args: argparse.Namespace, device: Any, relay: Any | None, logger: BenchLogger) -> None:
    if relay is None:
        raise SystemExit("relay was not opened")
    _relay_command(relay, args.power_off_command, logger)
    time.sleep(max(0.0, args.off_seconds))
    device.reset_input_buffer()
    _relay_command(relay, args.power_on_command, logger)


def _case_result(
    case: AdvanceCase,
    detector: FrameDetector,
    *,
    ready: bool,
    target: bytes | None,
    candidate_sent: bool,
    observation: str,
) -> dict[str, Any]:
    return {
        "name": case.name,
        "note": case.note,
        "frame": format_frame(case.frame),
        "ready": ready,
        "target_seen": target is not None,
        "target": format_frame(target) if target is not None else "",
        "candidate_sent": candidate_sent,
        "rx_frames": len(detector.frames),
        "labels": dict(detector.labels),
        "resync_events": detector.resync_events,
        "dropped_bytes": detector.dropped_bytes,
        "trailing_unframed_bytes": len(detector.buffer),
        "observation": observation,
    }


def _prompt_observation(args: argparse.Namespace, logger: BenchLogger, case: AdvanceCase) -> str:
    if not args.prompt_observation:
        return ""
    prompt = f"{case.name}: LCD/lamps/readouts observation, or Enter to skip: "
    observation = input(prompt).strip()
    logger.event(f"OBSERVATION {case.name}: {observation or '(no note)'}")
    return observation


def _emit_baseline_plan(logger: BenchLogger, baseline: tuple[bytes, ...], gap: float) -> None:
    logger.emit(f"baseline_gap={gap:.3f}s baseline_frames={len(baseline)}")
    for index, frame in enumerate(baseline, start=1):
        logger.emit(f"baseline[{index}]={format_frame(frame)} checksum_ok={int(frame_checksum_ok(frame))}")


def _emit_summary(logger: BenchLogger, results: list[dict[str, Any]]) -> None:
    logger.emit()
    logger.emit("Advance Sweep Summary")
    logger.emit(f"cases={len(results)}")
    for result in results:
        labels = ", ".join(f"{key}={value}" for key, value in sorted(result["labels"].items())) or "no_rx_frames"
        note = result["observation"] or "(no observation)"
        logger.emit(
            f"{result['name']}: ready={int(result['ready'])} target_seen={int(result['target_seen'])} "
            f"candidate_sent={int(result['candidate_sent'])} rx_frames={result['rx_frames']} "
            f"{labels} observation={note}"
        )


def _write_result_json(
    path: Path,
    log_path: Path,
    args: argparse.Namespace,
    baseline: tuple[bytes, ...],
    results: list[dict[str, Any]],
) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    payload = {
        "suite": args.suite,
        "log": str(log_path),
        "serial_format": f"{args.baud} {serial_format_label(args)}",
        "baseline": [format_frame(frame) for frame in baseline],
        "baseline_gap": args.baseline_gap,
        "target_mode": args.target_mode,
        "candidate_guard": args.candidate_guard,
        "cases": results,
    }
    path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")


def _print_dry_run(
    args: argparse.Namespace,
    cases: list[AdvanceCase],
    baseline: tuple[bytes, ...],
    log_path: Path,
    stdout: TextIO,
) -> None:
    print(f"suite={args.suite}", file=stdout)
    print(f"cases={len(cases)}", 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"power_cycle_start={int(not args.no_power_cycle)}", file=stdout)
    print(f"power_cycle_between_cases={int(args.power_cycle_between_cases)}", file=stdout)
    print(f"target_mode={args.target_mode} target_timeout={args.target_timeout:.3f}s", file=stdout)
    print(f"candidate_guard={args.candidate_guard:.3f}s post_candidate_read={args.post_candidate_read:.3f}s", file=stdout)
    print(f"baseline_gap={args.baseline_gap:.3f}s", file=stdout)
    for index, frame in enumerate(baseline, start=1):
        print(f"baseline[{index}]={format_frame(frame)} checksum_ok={int(frame_checksum_ok(frame))}", file=stdout)
    print(f"log={log_path}", file=stdout)
    for index, case in enumerate(cases, start=1):
        print(f"case[{index}]={case.name} frame={format_frame(case.frame)} checksum_ok={int(frame_checksum_ok(case.frame))}", file=stdout)
        if case.note:
            print(f"  note={case.note}", file=stdout)


def _parse_candidate(text: str) -> AdvanceCase:
    if "=" in text:
        name, frame_text = text.split("=", 1)
        name = name.strip()
        if not name:
            raise argparse.ArgumentTypeError("custom candidate name cannot be empty")
    else:
        name = "custom"
        frame_text = text
    return AdvanceCase(name=name, frame=parse_frame(frame_text), note="custom candidate")


__all__ = [
    "ACK_0040",
    "ACK_006B",
    "ACK_006C",
    "ACK_006D",
    "ACK_006E",
    "ACK_0096",
    "ACK_0097",
    "ACK_00C6",
    "ACK_00F8",
    "AdvanceCase",
    "REFRESH_OK",
    "build_arg_parser",
    "build_cases",
    "main",
    "select_cases",
]
