import argparse
import io
import json
import os
import sys
import traceback
from pathlib import Path
from typing import Iterable, List, Optional, Tuple

import cv2
import numpy as np

FILE = Path(__file__).resolve()
ROOT = FILE.parents[2]
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))

# Match the reference PDCL exporter behavior while still allowing callers to override
# credentials from their shell environment or `.env`.
os.environ.setdefault("STS_UID", "dis-uploader")
os.environ.setdefault("STS_SECRET_KEY", "277310cc09724d315514a79701fecb0f")

try:
    from dotenv import load_dotenv
except ImportError:
    def load_dotenv(*args, **kwargs):
        return False

from pdcl_pyclip.decoder_struct import StructDecoder
from pdcl_pyclip.msg_camera import VideoMessage
from pdcl_pyclip.reader import ClipReader

from tools.pdcl_inference.pdcl_clip_service import PDCLClipService
from tools.pdcl_inference.pipeline_types import ClipTask, TaskResult
from tools.pdcl_inference.status_store import StatusStore


CALIB_ATTACHMENT_NAMES = (
    "sigmastar.1/calibs/camera4.json",
    "test_data/calibs/camera4.json",
    "calibs/camera4.json",
)


def validate_pdcl_auth_env() -> None:
    required_vars = ["STS_UID", "STS_SECRET_KEY"]
    missing = [name for name in required_vars if not os.getenv(name)]
    if missing:
        raise ValueError(
            "Missing required PDCL auth env vars: "
            + ", ".join(missing)
            + ". Please export them in shell or set them in .env before running."
        )


def _plane_to_ndarray(plane) -> np.ndarray:
    stride = plane.line_size
    height = plane.height
    width = plane.width
    array = np.frombuffer(plane, dtype=np.uint8)
    if stride == width:
        return array.reshape(height, width)
    return array.reshape(height, stride)[:, :width]


def _yuvj420p_to_nv12(
    y_plane: np.ndarray,
    u_plane: np.ndarray,
    v_plane: np.ndarray,
) -> Tuple[np.ndarray, np.ndarray]:
    uv_height = u_plane.shape[0]
    uv_width = u_plane.shape[1]
    uv_plane = np.zeros((uv_height, uv_width * 2), dtype=np.uint8)
    uv_plane[:, 0::2] = u_plane
    uv_plane[:, 1::2] = v_plane
    return y_plane.copy(), uv_plane


def _h265_payload_to_bgr(payload: bytes) -> np.ndarray:
    try:
        import av
    except ImportError as exc:
        raise ImportError("PyAV is required for mcap decoding. Please install: pip install av") from exc

    container = av.open(io.BytesIO(payload))
    for frame in container.decode(video=0):
        y_plane = _plane_to_ndarray(frame.planes[0])
        u_plane = _plane_to_ndarray(frame.planes[1])
        v_plane = _plane_to_ndarray(frame.planes[2])
        y_nv12, uv_nv12 = _yuvj420p_to_nv12(y_plane, u_plane, v_plane)
        yuv_image = np.concatenate((y_nv12, uv_nv12), axis=0)
        return cv2.cvtColor(yuv_image, cv2.COLOR_YUV2BGR_NV12)
    raise ValueError("decode failed: no video frame in payload")


def resolve_calib_file_for_clip(clip_path: str) -> Optional[Path]:
    base = Path(clip_path).resolve()
    search_dir = base.parent
    candidates = []

    for _ in range(4):
        candidates.extend(
            [
                search_dir / "calibs" / "camera4.json",
                search_dir / "test_data" / "calibs" / "camera4.json",
                search_dir / "L2_calib" / "camera4.json",
                search_dir / "calib" / "L2_calib" / "camera4.json",
            ]
        )
        parent = search_dir.parent
        if parent == search_dir:
            break
        search_dir = parent

    for calib_path in candidates:
        if calib_path.exists():
            return calib_path
    return None


def load_embedded_calib_payload(clip_path: str) -> Optional[dict]:
    reader = ClipReader(clip_path)
    for attachment in reader.iter_attachments():
        if attachment.name in CALIB_ATTACHMENT_NAMES:
            return json.loads(attachment.data.decode("utf-8"))
    return None


def load_calib_payload_for_clip(
    clip_path: str,
    calib_file_override: str = "",
) -> Tuple[dict, str]:
    if calib_file_override:
        calib_src = Path(calib_file_override).resolve()
        if not calib_src.exists():
            raise FileNotFoundError(f"Calibration file {calib_src} does not exist.")
        with calib_src.open("r", encoding="utf-8") as file:
            return json.load(file), f"override:{calib_src}"

    calib_src = resolve_calib_file_for_clip(clip_path)
    if calib_src is not None:
        with calib_src.open("r", encoding="utf-8") as file:
            return json.load(file), f"file:{calib_src}"

    embedded_calib = load_embedded_calib_payload(clip_path)
    if embedded_calib is not None:
        return embedded_calib, "mcap_attachment"

    raise FileNotFoundError(
        "Calibration file camera4.json not found near clip and not found in mcap attachments. "
        "Please provide --calib-file explicitly."
    )


def decode_mcap_to_images(
    clip_uuid: str,
    clip_path: str,
    output_dir: Path,
    camera_topic: str,
    max_frames: int,
) -> Tuple[int, Optional[dict]]:
    output_dir.mkdir(parents=True, exist_ok=True)
    saved = 0
    camera4_json = load_embedded_calib_payload(clip_path)

    for frame_name, image in iter_decoded_clip_frames(
        clip_uuid=clip_uuid,
        clip_path=clip_path,
        camera_topic=camera_topic,
        max_frames=max_frames,
    ):
        if not cv2.imwrite(str(output_dir / frame_name), image):
            raise IOError(f"Failed to write image: {output_dir / frame_name}")
        saved += 1

    return saved, camera4_json


def iter_decoded_clip_frames(
    clip_uuid: str,
    clip_path: str,
    camera_topic: str,
    max_frames: int,
) -> Iterable[Tuple[str, np.ndarray]]:
    reader = ClipReader(clip_path)
    struct_decoder = StructDecoder()
    decoded = 0

    for schema, channel, msg in reader.iter_messages(topics=[camera_topic]):
        if schema.encoding != "struct":
            continue
        data = struct_decoder.decode(schema, channel, msg)
        if not isinstance(data, VideoMessage):
            continue

        frame_name = f"{clip_uuid}_{data.frame_id}.png"
        yield frame_name, _h265_payload_to_bgr(data.payload)

        decoded += 1
        if max_frames > 0 and decoded >= max_frames:
            break


def save_calib_for_clip(
    clip_path: str,
    calib_output_path: Path,
    calib_file_override: str,
    embedded_calib: Optional[dict],
) -> str:
    calib_output_path.parent.mkdir(parents=True, exist_ok=True)

    if calib_file_override:
        calib_src = Path(calib_file_override).resolve()
        if not calib_src.exists():
            raise FileNotFoundError(f"Calibration file {calib_src} does not exist.")
        calib_output_path.write_bytes(calib_src.read_bytes())
        return f"override:{calib_src}"

    calib_src = resolve_calib_file_for_clip(clip_path)
    if calib_src is not None:
        calib_output_path.write_bytes(calib_src.read_bytes())
        return f"file:{calib_src}"

    if embedded_calib is None:
        raise FileNotFoundError(
            "Calibration file camera4.json not found near clip and not found in mcap attachments. "
            "Please provide --calib-file explicitly."
        )

    with open(calib_output_path, "w") as file:
        json.dump(embedded_calib, file, indent=2, ensure_ascii=False)
    return "mcap_attachment"


def build_calib_summary(calib_payload: Optional[dict]) -> dict:
    """Return a compact summary that works for flat and combined camera4 calib JSON files."""
    if not isinstance(calib_payload, dict):
        return {
            "format": "unknown",
            "has_distortion": False,
        }

    intrinsics = calib_payload
    extrinsics = None
    calib_format = "flat_camera4"

    if "intrinsics" in calib_payload:
        intrinsics = calib_payload.get("intrinsics", {}).get("camera4.json", {}) or {}
        extrinsics = calib_payload.get("extrinsics", {}).get("camera4.json", {}) or {}
        calib_format = "combined_calibration"

    distort_coeffs = intrinsics.get("distort_coeffs", calib_payload.get("distort_coeffs", [])) or []
    return {
        "format": calib_format,
        "focal_u": intrinsics.get("focal_u"),
        "focal_v": intrinsics.get("focal_v"),
        "cu": intrinsics.get("cu"),
        "cv": intrinsics.get("cv"),
        "pitch": calib_payload.get("pitch", extrinsics.get("rpy", [None, None, None])[1] if extrinsics else None),
        "yaw": calib_payload.get("yaw", extrinsics.get("rpy", [None, None, None])[2] if extrinsics else None),
        "image_width": calib_payload.get("image_width", calib_payload.get("img_width")),
        "image_height": calib_payload.get("image_height", calib_payload.get("img_height")),
        "distort_coeff_count": len(distort_coeffs),
        "has_distortion": bool(distort_coeffs),
    }


def parse_clip_list(clip_list_file: str, vehicle_name: str) -> List[ClipTask]:
    tasks: List[ClipTask] = []
    with open(clip_list_file, "r") as file:
        for line in file:
            line = line.strip()
            if not line or line.startswith("#"):
                continue

            parts = line.split()
            clip_uuid = parts[0]
            date_name = parts[1] if len(parts) >= 2 else "unknown_date"
            item_vehicle_name = parts[2] if len(parts) >= 3 else vehicle_name

            clip_path = PDCLClipService.get_clip_path(clip_uuid)
            if clip_path is None:
                print(f"Skip clip_id={clip_uuid}: no files found")
                continue

            tasks.append(
                ClipTask(
                    clip_uuid=clip_uuid,
                    date_name=date_name,
                    vehicle_name=item_vehicle_name,
                    clip_path=clip_path,
                )
            )
    return tasks


def build_case_dir_name(prefix: str, clip_task: ClipTask) -> str:
    safe_clip_uuid = clip_task.clip_uuid.replace("/", "_")
    return f"{prefix}_{clip_task.vehicle_name}_{clip_task.date_name}_{safe_clip_uuid}"


def save_manifest(
    output_dir: Path,
    clip_task: ClipTask,
    frame_count: int,
    calib_source: str,
    camera_topic: str,
    calib_summary: Optional[dict] = None,
) -> None:
    manifest = {
        "clip_uuid": clip_task.clip_uuid,
        "date_name": clip_task.date_name,
        "vehicle_name": clip_task.vehicle_name,
        "clip_path": clip_task.clip_path,
        "camera_topic": camera_topic,
        "frame_count": frame_count,
        "calib_source": calib_source,
        "calib_summary": calib_summary,
    }
    with open(output_dir / "manifest.json", "w") as file:
        json.dump(manifest, file, indent=2, ensure_ascii=False)


def save_calib_summary(output_dir: Path, calib_source: str, calib_summary: dict) -> None:
    payload = {
        "calib_source": calib_source,
        "calib_summary": calib_summary,
    }
    with open(output_dir / "calib_summary.json", "w") as file:
        json.dump(payload, file, indent=2, ensure_ascii=False)


def export_one_clip(
    clip_task: ClipTask,
    args: argparse.Namespace,
    status: StatusStore,
) -> TaskResult:
    task_id = clip_task.task_id
    if args.skip_done and status.is_done(task_id):
        info = status.get(task_id) or {}
        return TaskResult(task_id=task_id, success=True, message=f"skip done: {info.get('detail', '')}")

    case_dir = Path(args.output_root) / build_case_dir_name(args.output_prefix, clip_task)
    images_dir = case_dir / "images"
    calib_path = case_dir / "calib" / "L2_calib" / "camera4.json"
    status.mark_running(task_id, f"running source={clip_task.clip_path}")

    try:
        frame_count, embedded_calib = decode_mcap_to_images(
            clip_uuid=clip_task.clip_uuid,
            clip_path=clip_task.clip_path,
            output_dir=images_dir,
            camera_topic=args.camera_topic,
            max_frames=args.max_frames_per_clip,
        )
        calib_source = save_calib_for_clip(
            clip_path=clip_task.clip_path,
            calib_output_path=calib_path,
            calib_file_override=args.calib_file,
            embedded_calib=embedded_calib,
        )
        with open(calib_path, "r", encoding="utf-8") as file:
            calib_payload = json.load(file)
        calib_summary = build_calib_summary(calib_payload)
        save_manifest(
            output_dir=case_dir,
            clip_task=clip_task,
            frame_count=frame_count,
            calib_source=calib_source,
            camera_topic=args.camera_topic,
            calib_summary=calib_summary,
        )
        save_calib_summary(case_dir, calib_source=calib_source, calib_summary=calib_summary)
        status.mark_done(task_id, str(case_dir))
        return TaskResult(task_id=task_id, success=True, message=f"ok frames={frame_count}", output_dir=str(case_dir))
    except Exception as exc:
        message = f"{type(exc).__name__}: {exc}"
        status.mark_failed(task_id, message)
        err_log = Path(args.output_root) / "_status" / "errors.log"
        err_log.parent.mkdir(parents=True, exist_ok=True)
        with open(err_log, "a") as file:
            file.write(f"[{task_id}] {message}\n{traceback.format_exc()}\n")
        return TaskResult(task_id=task_id, success=False, message=message)


def save_run_manifest(args: argparse.Namespace, clip_tasks: List[ClipTask]) -> None:
    manifest_path = Path(args.output_root) / "_status" / "run_manifest.json"
    manifest_path.parent.mkdir(parents=True, exist_ok=True)
    payload = {
        "clip_list_file": args.clip_list_file,
        "output_root": args.output_root,
        "camera_topic": args.camera_topic,
        "max_frames_per_clip": args.max_frames_per_clip,
        "num_tasks": len(clip_tasks),
        "tasks": [
            {
                "task_id": task.task_id,
                "clip_uuid": task.clip_uuid,
                "date_name": task.date_name,
                "vehicle_name": task.vehicle_name,
                "clip_path": task.clip_path,
            }
            for task in clip_tasks
        ],
    }
    with open(manifest_path, "w") as file:
        json.dump(payload, file, indent=2, ensure_ascii=False)


def parse_opt() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Batch decode PDCL MCAP clips to images and save matching camera calibration"
    )
    parser.add_argument(
        "--clip-list-file",
        type=str,
        default=str(FILE.parent / "clips_6286.txt"),
        help="Path to text file: '<clip_id> [date] [vehicle_name]' per line",
    )
    parser.add_argument(
        "--vehicle-name",
        type=str,
        default="G1M3",
        help="Default vehicle name when clip-list does not include a third column",
    )
    parser.add_argument(
        "--output-root",
        type=str,
        default=str(ROOT / "runs" / "pdcl_clip_exports"),
        help="Root directory used to save decoded frames and calibration",
    )
    parser.add_argument(
        "--output-prefix",
        type=str,
        default="clip_export",
        help="Prefix used in each clip output directory name",
    )
    parser.add_argument("--camera-topic", type=str, default="camera4", help="Camera topic name in mcap")
    parser.add_argument(
        "--max-frames-per-clip",
        type=int,
        default=0,
        help="Max frames decoded for each clip; 0 means all",
    )
    parser.add_argument(
        "--calib-file",
        type=str,
        default="",
        help="Optional override path to camera4.json for all exported clips",
    )
    parser.add_argument("--skip-done", action="store_true", help="Skip clips already marked done in status file")
    parser.add_argument("--limit-clips", type=int, default=0, help="Limit number of clips; 0 means all")
    return parser.parse_args()


def main() -> None:
    load_dotenv()
    validate_pdcl_auth_env()

    args = parse_opt()
    os.makedirs(args.output_root, exist_ok=True)

    clip_tasks = parse_clip_list(args.clip_list_file, args.vehicle_name)
    if args.limit_clips > 0:
        clip_tasks = clip_tasks[: args.limit_clips]
    if not clip_tasks:
        print("No clip tasks discovered. Exit.")
        return

    status_path = Path(args.output_root) / "_status" / "task_status.json"
    status = StatusStore(status_path)
    save_run_manifest(args, clip_tasks)

    print(f"Discovered {len(clip_tasks)} clips.")
    print(f"Output root: {args.output_root}")
    print(f"Status file: {status_path}")

    success = 0
    fail = 0
    for idx, clip_task in enumerate(clip_tasks, start=1):
        print(
            f"[{idx}/{len(clip_tasks)}] clip_id={clip_task.clip_uuid} "
            f"source={clip_task.clip_path}"
        )
        result = export_one_clip(clip_task, args, status)
        if result.success:
            success += 1
        else:
            fail += 1
        print(f"  -> {result.message}")

    print("\nBatch export done.")
    print(f"success={success}, fail={fail}")
    print(f"status_summary={status.summary()}")


if __name__ == "__main__":
    main()
