#!/usr/bin/env python3
from __future__ import annotations

import argparse
import contextlib
from collections import defaultdict, deque
import glob
import json
import math
import os
import sys
import time
from pathlib import Path
from typing import Any

from camera_recorder_config import load_config, load_mapper_configs, resolve_default_model, resolve_path
from g1_camera_pose_mapper import PoseToG1Mapper, keypoints_to_serializable, load_home_pose


cv2 = None
YOLO = None


def clamp(value: float, low: float, high: float) -> float:
    return max(low, min(high, float(value)))


def ensure_runtime_dependencies(runtime_cfg: dict[str, Any]) -> None:
    global cv2, YOLO

    messages = runtime_cfg["dependency_messages"]

    if cv2 is None:
        try:
            import cv2 as _cv2
        except ModuleNotFoundError as exc:
            raise SystemExit(messages["opencv_missing"]) from exc
        cv2 = _cv2

    if YOLO is None:
        try:
            from ultralytics import YOLO as _YOLO
        except ModuleNotFoundError as exc:
            raise SystemExit(messages["ultralytics_missing"]) from exc
        YOLO = _YOLO


def parse_source(source: str) -> int | str:
    return int(source) if source.isdigit() else source


def wait_for_confirmation(prompt: str, enabled: bool = True) -> bool:
    if not enabled or not sys.stdin.isatty():
        return True
    try:
        input(prompt)
        return True
    except KeyboardInterrupt:
        print("\n[REC] Ctrl+C pressed.")
        return False


def ensure_unique_path(user_path: str, default_dir: Path, suffix: str) -> Path:
    path = Path(user_path)
    if not path.suffix:
        path = path.with_suffix(suffix)

    if len(path.parts) == 1:
        default_dir.mkdir(parents=True, exist_ok=True)
        candidate = default_dir / path.name
    else:
        candidate = path
        candidate.parent.mkdir(parents=True, exist_ok=True)

    if not candidate.exists():
        return candidate

    index = 1
    while True:
        alt = candidate.with_name(f"{candidate.stem}({index}){candidate.suffix}")
        if not alt.exists():
            return alt
        index += 1


def ensure_unique_output_pair(
    user_path: str,
    dataset_dir: Path,
    dataset_suffix: str,
    raw_pose_dir: Path,
    raw_pose_suffix: str,
) -> tuple[Path, Path]:
    dataset_path = Path(user_path)
    if not dataset_path.suffix:
        dataset_path = dataset_path.with_suffix(dataset_suffix)

    if len(dataset_path.parts) == 1:
        dataset_dir.mkdir(parents=True, exist_ok=True)
        dataset_candidate = dataset_dir / dataset_path.name
    else:
        dataset_candidate = dataset_path
        dataset_candidate.parent.mkdir(parents=True, exist_ok=True)

    raw_pose_dir.mkdir(parents=True, exist_ok=True)

    def raw_candidate_for(path: Path) -> Path:
        return raw_pose_dir / f"{path.stem}{raw_pose_suffix}"

    raw_candidate = raw_candidate_for(dataset_candidate)
    if not dataset_candidate.exists() and not raw_candidate.exists():
        return dataset_candidate, raw_candidate

    index = 1
    while True:
        dataset_alt = dataset_candidate.with_name(f"{dataset_candidate.stem}({index}){dataset_candidate.suffix}")
        raw_alt = raw_candidate_for(dataset_alt)
        if not dataset_alt.exists() and not raw_alt.exists():
            return dataset_alt, raw_alt
        index += 1


def _opencv_error_log_level() -> Any:
    if hasattr(cv2, "LOG_LEVEL_ERROR"):
        return cv2.LOG_LEVEL_ERROR
    if hasattr(cv2, "utils") and hasattr(cv2.utils, "logging"):
        return cv2.utils.logging.LOG_LEVEL_ERROR
    return None


@contextlib.contextmanager
def suppress_camera_probe_noise():
    old_level = None
    try:
        if hasattr(cv2, "getLogLevel"):
            old_level = cv2.getLogLevel()
        error_level = _opencv_error_log_level()
        if error_level is not None and hasattr(cv2, "setLogLevel"):
            cv2.setLogLevel(error_level)
    except Exception:
        old_level = None

    with open(os.devnull, "w", encoding="utf-8") as devnull:
        with contextlib.redirect_stderr(devnull):
            try:
                yield
            finally:
                try:
                    if old_level is not None and hasattr(cv2, "setLogLevel"):
                        cv2.setLogLevel(old_level)
                except Exception:
                    pass


def camera_indices_from_dev(max_index: int) -> list[int]:
    indices = []
    for path in glob.glob("/dev/video*"):
        suffix = path.replace("/dev/video", "")
        if suffix.isdigit():
            index = int(suffix)
            if index <= max_index:
                indices.append(index)
    return sorted(set(indices))


def resolve_camera_backend(runtime_cfg: dict[str, Any]) -> int:
    backend_name = runtime_cfg["camera_backend"]
    return int(getattr(cv2, backend_name, getattr(cv2, "CAP_ANY", 0)))


def _try_open_camera(index: int, runtime_cfg: dict[str, Any]) -> dict[str, float | int] | None:
    backend = resolve_camera_backend(runtime_cfg)
    probe_attempts = int(runtime_cfg["camera_probe_attempts"])
    fps_threshold = float(runtime_cfg["camera_fps_threshold"])

    cap = cv2.VideoCapture(index, backend)
    if not cap.isOpened():
        cap.release()
        return None

    ok = False
    frame = None
    for _ in range(probe_attempts):
        ok, frame = cap.read()
        if ok and frame is not None and frame.size > 0:
            break

    if not ok or frame is None:
        cap.release()
        return None

    fps = float(cap.get(cv2.CAP_PROP_FPS) or 0.0)
    if fps <= fps_threshold or math.isnan(fps):
        fps = 0.0

    info = {
        "index": index,
        "width": int(frame.shape[1]),
        "height": int(frame.shape[0]),
        "fps": fps,
    }
    cap.release()
    return info


def list_available_cameras(max_index: int, runtime_cfg: dict[str, Any]) -> list[dict[str, float | int]]:
    candidates = camera_indices_from_dev(max_index)
    if not candidates:
        candidates = list(range(max_index + 1))

    cameras = []
    with suppress_camera_probe_noise():
        for index in candidates:
            info = _try_open_camera(index, runtime_cfg)
            if info is not None:
                cameras.append(info)
    return cameras


def choose_camera(cameras: list[dict[str, float | int]], runtime_cfg: dict[str, Any]) -> int:
    if len(cameras) == 1:
        return int(cameras[0]["index"])

    fps_threshold = float(runtime_cfg["camera_fps_threshold"])

    print("\nAvailable cameras:")
    for option_num, cam in enumerate(cameras, start=1):
        fps = float(cam["fps"])
        fps_text = f"{fps:.1f}fps" if fps > fps_threshold and not math.isnan(fps) else "fps unknown"
        print(
            f"  {option_num}. camera {int(cam['index'])} "
            f"({int(cam['width'])}x{int(cam['height'])}, {fps_text})"
        )

    if not sys.stdin.isatty():
        return int(cameras[0]["index"])

    while True:
        try:
            choice = input("Choose camera option number: ").strip()
        except KeyboardInterrupt as exc:
            raise SystemExit("\n[REC] Camera selection cancelled by Ctrl+C.") from exc
        if choice.isdigit():
            selected = int(choice)
            if 1 <= selected <= len(cameras):
                return int(cameras[selected - 1]["index"])
        print("Invalid choice. Enter a valid option number.")


def open_capture(source: int | str, runtime_cfg: dict[str, Any]):
    backend = resolve_camera_backend(runtime_cfg)
    if isinstance(source, int):
        cap = cv2.VideoCapture(source, backend)
        if cap.isOpened():
            return cap
    return cv2.VideoCapture(source)


def extract_people(result, kpt_conf: float) -> list[dict[str, Any]]:
    boxes = result.boxes
    kpts = result.keypoints
    if boxes is None or kpts is None or len(boxes) == 0:
        return []

    ids = boxes.id
    if ids is not None:
        track_ids = [int(value) for value in ids.int().cpu().tolist()]
    else:
        track_ids = [None] * len(boxes)

    xyxy = boxes.xyxy.cpu().tolist()
    box_conf = boxes.conf.cpu().tolist() if boxes.conf is not None else [0.0] * len(boxes)
    kpts_xy = kpts.xy.cpu().numpy()
    kpts_conf = kpts.conf.cpu().numpy() if kpts.conf is not None else None

    people = []
    for index in range(len(boxes)):
        x1, y1, x2, y2 = [float(v) for v in xyxy[index]]
        kp_map: dict[int, tuple[float, float, float]] = {}
        confs = []
        for kp_idx in range(len(kpts_xy[index])):
            conf = float(kpts_conf[index][kp_idx]) if kpts_conf is not None else 1.0
            x, y = float(kpts_xy[index][kp_idx][0]), float(kpts_xy[index][kp_idx][1])
            kp_map[kp_idx] = (x, y, conf)
            if conf >= kpt_conf:
                confs.append(conf)

        area = max(1.0, (x2 - x1) * (y2 - y1))
        people.append(
            {
                "track_id": track_ids[index],
                "bbox": [x1, y1, x2, y2],
                "det_conf": float(box_conf[index]),
                "mean_kpt_conf": float(sum(confs) / len(confs)) if confs else 0.0,
                "area": area,
                "keypoints": kp_map,
            }
        )
    return people


def select_person(
    people: list[dict[str, Any]],
    locked_track_id: int | None,
    selection_cfg: dict[str, Any],
) -> dict[str, Any] | None:
    if not people:
        return None

    if locked_track_id is not None:
        for person in people:
            if person["track_id"] == locked_track_id:
                return person

    bias = float(selection_cfg["area_confidence_bias"])
    return max(
        people,
        key=lambda person: float(person["area"]) * (bias + float(person["mean_kpt_conf"])),
    )


def draw_people(
    frame,
    people: list[dict[str, Any]],
    selected_person: dict[str, Any] | None,
    track_history: dict[str, deque[tuple[int, int]]],
    kpt_conf: float,
    pose_ui_cfg: dict[str, Any],
    runtime_cfg: dict[str, Any],
) -> Any:
    selected_key = str(runtime_cfg["selected_track_history_key"])
    selected_id = selected_person["track_id"] if selected_person else None
    joint_names = {int(idx): label for idx, label in pose_ui_cfg["joint_names"].items()}
    skeleton_edges = [tuple(int(v) for v in edge) for edge in pose_ui_cfg["skeleton_edges"]]

    for person in people:
        bbox = person["bbox"]
        keypoints = person["keypoints"]
        is_selected = person is selected_person or (
            selected_id is not None and person["track_id"] == selected_id
        )

        line_color = (0, 220, 255) if is_selected else (110, 110, 110)
        text_color = (255, 255, 255) if is_selected else (180, 180, 180)
        point_color = (0, 64, 255) if is_selected else (140, 140, 140)
        thickness = 2 if is_selected else 1

        x1, y1, x2, y2 = [int(v) for v in bbox]
        cv2.rectangle(frame, (x1, y1), (x2, y2), line_color, thickness)

        valid_points = []
        for a, b in skeleton_edges:
            pa = keypoints.get(a)
            pb = keypoints.get(b)
            if pa is None or pb is None:
                continue
            if pa[2] < kpt_conf or pb[2] < kpt_conf:
                continue
            p1 = (int(pa[0]), int(pa[1]))
            p2 = (int(pb[0]), int(pb[1]))
            valid_points.extend([p1, p2])
            cv2.line(frame, p1, p2, line_color, thickness)

        for idx, label in joint_names.items():
            point = keypoints.get(idx)
            if point is None or point[2] < kpt_conf:
                continue
            px, py = int(point[0]), int(point[1])
            valid_points.append((px, py))
            cv2.circle(frame, (px, py), 4 if is_selected else 3, point_color, -1)
            if is_selected:
                cv2.putText(
                    frame,
                    label,
                    (px + 5, py - 6),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.35,
                    text_color,
                    1,
                    cv2.LINE_AA,
                )

        if is_selected and valid_points:
            cx = int(sum(point[0] for point in valid_points) / len(valid_points))
            cy = int(sum(point[1] for point in valid_points) / len(valid_points))
            history = track_history[selected_key]
            history.append((cx, cy))
            while len(history) > history.maxlen:
                history.popleft()
            for index in range(1, len(history)):
                cv2.line(frame, history[index - 1], history[index], (0, 255, 0), 2)

        label = f"ID {person['track_id']}" if person["track_id"] is not None else "person"
        cv2.putText(
            frame,
            label,
            (x1 + 5, max(18, y1 - 8)),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.55,
            text_color,
            2 if is_selected else 1,
            cv2.LINE_AA,
        )

    return frame


def add_overlay(
    frame,
    *,
    elapsed: float,
    samples: int,
    selected_person: dict[str, Any] | None,
    diagnostics: dict[str, Any],
    output_name: str,
) -> None:
    lines = [
        f"REC {elapsed:6.2f}s  samples={samples}",
        f"target={output_name}",
        (
            f"track={selected_person['track_id']} "
            f"left={int(bool(diagnostics.get('left_ready')))} "
            f"right={int(bool(diagnostics.get('right_ready')))}"
        )
        if selected_person is not None
        else "track=none  waiting for pose",
        "Press q or ESC to stop",
    ]

    for idx, line in enumerate(lines):
        y = 28 + idx * 24
        cv2.putText(frame, line, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.65, (16, 16, 16), 4, cv2.LINE_AA)
        cv2.putText(frame, line, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.65, (255, 255, 255), 1, cv2.LINE_AA)


def build_parser(config: dict[str, Any], config_path: str) -> argparse.ArgumentParser:
    defaults = config["defaults"]
    parser = argparse.ArgumentParser(
        description=(
            "Capture human pose from a camera with YOLO and save a replay-compatible G1 dataset."
        )
    )
    parser.add_argument("--config", default=config_path, help="Path to camera_recorder_config.json")
    parser.add_argument(
        "--source",
        default=defaults["source"],
        help="Video source: webcam index, path, or 'ask' to scan available cameras.",
    )
    parser.add_argument(
        "--camera-scan-max",
        type=int,
        default=int(defaults["camera_scan_max"]),
        help="Max camera index to scan when --source ask is used.",
    )
    parser.add_argument(
        "--model",
        default=None,
        help="YOLO pose model path/name. Defaults to the first model candidate found in camera_recorder_config.json.",
    )
    parser.add_argument(
        "--output",
        "--output-file",
        "--output-name",
        dest="output",
        default=defaults["output"],
        help="Output dataset name/path.",
    )
    parser.add_argument(
        "--home-pose",
        "--home",
        "--home-file",
        default=config["paths"]["default_home_pose"],
        dest="home_pose",
        help="Home pose jsonl used for lower body and fallback upper body values.",
    )
    parser.add_argument("--seconds", type=float, default=float(defaults["seconds"]), help="0 = record until q / ESC.")
    parser.add_argument("--record-hz", type=float, default=float(defaults["record_hz"]), help="Max dataset sample rate.")
    parser.add_argument("--conf", type=float, default=float(defaults["conf"]), help="Detection confidence threshold.")
    parser.add_argument("--kpt-conf", type=float, default=float(defaults["kpt_conf"]), help="Keypoint confidence threshold.")
    parser.add_argument("--imgsz", type=int, default=int(defaults["imgsz"]), help="Inference image size.")
    parser.add_argument("--smoothing", type=float, default=float(defaults["smoothing"]), help="Low-pass blend for mapped joints.")
    parser.add_argument("--window-width", type=int, default=int(defaults["window_width"]), help="GUI window width.")
    parser.add_argument("--window-height", type=int, default=int(defaults["window_height"]), help="GUI window height.")
    parser.add_argument("--show", dest="show", action="store_true", help="Show the live camera window.")
    parser.add_argument("--no-show", dest="show", action="store_false", help="Disable the live GUI.")
    parser.add_argument(
        "--no-prompt",
        action="store_true",
        default=bool(defaults.get("no_prompt", False)),
        help="Skip the Enter-to-begin confirmation prompt.",
    )
    parser.set_defaults(show=bool(defaults["show"]))
    return parser


def parse_args_and_config() -> tuple[argparse.Namespace, dict[str, Any], Path]:
    pre_parser = argparse.ArgumentParser(add_help=False)
    pre_parser.add_argument("--config", default="camera_recorder_config.json")
    pre_args, _ = pre_parser.parse_known_args()

    config, config_dir = load_config(pre_args.config)
    parser = build_parser(config, pre_args.config)
    args = parser.parse_args()
    return args, config, config_dir


def main() -> None:
    args, config, config_dir = parse_args_and_config()
    runtime_cfg = config["runtime"]
    output_cfg = config["output"]
    selection_cfg = config["selection"]
    pose_ui_cfg = config["pose_ui"]
    joint_cfg, pose_mapping_cfg, joint_cfg_path, pose_mapping_cfg_path = load_mapper_configs(config, config_dir)

    ensure_runtime_dependencies(runtime_cfg)

    data_dir = resolve_path(config_dir, config["paths"]["data_dir"])
    raw_pose_dir = resolve_path(config_dir, config["paths"]["raw_pose_dir"])
    model_path = args.model if args.model else resolve_default_model(config, config_dir)
    home_pose_path = resolve_path(config_dir, args.home_pose)

    output_path, raw_output_path = ensure_unique_output_pair(
        args.output,
        data_dir,
        output_cfg["dataset_suffix"],
        raw_pose_dir,
        output_cfg["raw_pose_suffix"],
    )

    display_env_var = str(runtime_cfg["display_env_var"])
    if args.show and os.name != "nt" and not os.environ.get(display_env_var):
        print(f"{display_env_var} not found. Disabling GUI window.")
        args.show = False

    home_q = load_home_pose(home_pose_path, int(joint_cfg["g1_num_motor"]))
    mapper = PoseToG1Mapper(
        home_q,
        joint_cfg,
        pose_mapping_cfg,
        min_conf=args.kpt_conf,
        smoothing=args.smoothing,
    )
    model = YOLO(model_path)

    if str(args.source).lower() == "ask":
        cameras = list_available_cameras(args.camera_scan_max, runtime_cfg)
        if not cameras:
            raise SystemExit("No camera found. Connect a camera or pass --source <path/index>.")
        source = choose_camera(cameras, runtime_cfg)
        print(f"Using camera index: {source}")
    else:
        source = parse_source(str(args.source))

    cap = open_capture(source, runtime_cfg)
    if not cap.isOpened():
        raise SystemExit(f"Could not open source: {source}")

    if args.show:
        cv2.namedWindow(str(runtime_cfg["window_name"]), cv2.WINDOW_NORMAL)
        cv2.resizeWindow(
            str(runtime_cfg["window_name"]),
            max(320, args.window_width),
            max(240, args.window_height),
        )

    meta = {
        "meta": {
            "format": output_cfg["dataset_meta_format"],
            "created_unix": time.time(),
            "motors": mapper.num_motors,
            "source": str(source),
            "model": str(model_path),
            "record_hz": float(args.record_hz),
            "home_pose": str(home_pose_path),
            "raw_pose_file": str(raw_output_path.resolve()),
            "config_file": str(resolve_path(config_dir, args.config)),
            "joint_config_file": str(joint_cfg_path),
            "pose_mapping_file": str(pose_mapping_cfg_path),
            "notes": output_cfg["dataset_notes"],
        }
    }
    raw_meta = {
        "meta": {
            "format": output_cfg["raw_meta_format"],
            "created_unix": time.time(),
            "source": str(source),
            "model": str(model_path),
            "kpt_conf": float(args.kpt_conf),
            "dataset_file": str(output_path.resolve()),
            "config_file": str(resolve_path(config_dir, args.config)),
            "joint_config_file": str(joint_cfg_path),
            "pose_mapping_file": str(pose_mapping_cfg_path),
        }
    }

    selected_track_id = None
    history_key = str(runtime_cfg["selected_track_history_key"])
    history_len = int(runtime_cfg["track_history_maxlen"])
    track_history: dict[str, deque[tuple[int, int]]] = defaultdict(lambda: deque(maxlen=history_len))
    last_sample_wall = 0.0
    last_status_wall = 0.0
    samples = 0
    started_wall = time.time()

    print(f"[REC] Dataset  -> {output_path}")
    print(f"[REC] Raw pose -> {raw_output_path}")
    print("[REC] Press q or ESC to stop.")

    if not wait_for_confirmation("👉 Press Enter to Begin recording... ", enabled=not args.no_prompt):
        cap.release()
        if args.show:
            cv2.destroyAllWindows()
        print("[REC] Recording cancelled before start.")
        return

    with output_path.open("w", encoding="utf-8") as dataset_file, raw_output_path.open("w", encoding="utf-8") as raw_file:
        dataset_file.write(json.dumps(meta) + "\n")
        raw_file.write(json.dumps(raw_meta) + "\n")

        try:
            while True:
                ok, frame = cap.read()
                if not ok:
                    break

                result = model.track(
                    source=frame,
                    conf=args.conf,
                    imgsz=args.imgsz,
                    persist=True,
                    verbose=False,
                )[0]
                people = extract_people(result, args.kpt_conf)
                selected_person = select_person(people, selected_track_id, selection_cfg)
                if selected_person is not None and selected_person["track_id"] is not None:
                    selected_track_id = int(selected_person["track_id"])
                elif not people:
                    selected_track_id = None
                    track_history[history_key].clear()

                keypoints = selected_person["keypoints"] if selected_person is not None else None
                q_frame, diagnostics = mapper.convert(keypoints)

                now = time.time()
                elapsed = now - started_wall

                sample_dt = 0.0 if args.record_hz <= 0 else (1.0 / args.record_hz)
                should_store = sample_dt == 0.0 or (now - last_sample_wall) >= sample_dt
                if should_store:
                    dataset_row = {
                        "t": round(elapsed, 6),
                        "q": [round(float(v), 6) for v in q_frame],
                        "tracked": bool(diagnostics["tracked"]),
                        "track_id": selected_person["track_id"] if selected_person is not None else None,
                    }
                    raw_row = {
                        "t": round(elapsed, 6),
                        "tracked": bool(diagnostics["tracked"]),
                        "track_id": selected_person["track_id"] if selected_person is not None else None,
                        "bbox": selected_person["bbox"] if selected_person is not None else None,
                        "det_conf": selected_person["det_conf"] if selected_person is not None else 0.0,
                        "mean_kpt_conf": selected_person["mean_kpt_conf"] if selected_person is not None else 0.0,
                        "diagnostics": diagnostics,
                        "keypoints": keypoints_to_serializable(keypoints, mapper.serializable_indices),
                    }
                    dataset_file.write(json.dumps(dataset_row) + "\n")
                    raw_file.write(json.dumps(raw_row) + "\n")
                    samples += 1
                    last_sample_wall = now

                if args.show:
                    display = draw_people(
                        frame,
                        people,
                        selected_person,
                        track_history,
                        args.kpt_conf,
                        pose_ui_cfg,
                        runtime_cfg,
                    )
                    add_overlay(
                        display,
                        elapsed=elapsed,
                        samples=samples,
                        selected_person=selected_person,
                        diagnostics=diagnostics,
                        output_name=output_path.name,
                    )
                    cv2.imshow(str(runtime_cfg["window_name"]), display)
                    key = cv2.waitKey(1) & 0xFF
                    if key in (27, ord("q")):
                        break
                elif now - last_status_wall >= 1.0:
                    status = "tracked" if diagnostics["tracked"] else "waiting"
                    print(f"[REC] {elapsed:6.2f}s samples={samples} state={status}")
                    last_status_wall = now

                if args.seconds > 0 and elapsed >= args.seconds:
                    break

        except KeyboardInterrupt:
            print("\n[REC] Stopped by user.")
        finally:
            cap.release()
            if args.show:
                cv2.destroyAllWindows()

    print(f"[REC] Saved {samples} samples to {output_path}")
    print(f"[REC] Saved raw pose log to {raw_output_path}")


if __name__ == "__main__":
    main()
