import socket
import struct
import csv
import os
import json
from datetime import datetime
from collections import deque

import numpy as np
import cv2

HOST = "0.0.0.0"
PORT = 9999
DATASET_ROOT = "data/iphone"
GRAVITY = 9.81  # m/s², used to convert accel from g-units to m/s²


def recv_exact(conn, n):
    data = b""
    while len(data) < n:
        chunk = conn.recv(n - len(data))
        if not chunk:
            raise ConnectionError("socket closed")
        data += chunk
    return data


def make_session_dirs():
    session_name = datetime.now().strftime("%Y%m%d_%H%M%S")
    session_dir = os.path.join(DATASET_ROOT, session_name)
    images_dir = os.path.join(session_dir, "images")
    os.makedirs(images_dir, exist_ok=True)
    return session_dir, images_dir


def interp_sample(buf, t_ns):
    """
    buf: deque of (timestamp_ns, x, y, z), sorted by time
    t_ns: target timestamp

    Return interpolated (x, y, z) if there exist samples bracketing t_ns.
    Otherwise return None.
    """
    if len(buf) < 2:
        return None

    for i in range(len(buf) - 1):
        t0, x0, y0, z0 = buf[i]
        t1, x1, y1, z1 = buf[i + 1]

        if t0 <= t_ns <= t1:
            if t1 == t0:
                return (x0, y0, z0)

            alpha = (t_ns - t0) / float(t1 - t0)
            x = (1.0 - alpha) * x0 + alpha * x1
            y = (1.0 - alpha) * y0 + alpha * y1
            z = (1.0 - alpha) * z0 + alpha * z1
            return (x, y, z)

    return None


def flush_synced_imu(accel_buf, gyro_buf, imu_writer, f_imu):
    """
    Emit all accel samples that can already be bracketed by gyro samples.
    Return number of emitted imu rows.
    """
    emitted = 0

    if len(accel_buf) == 0 or len(gyro_buf) < 2:
        return emitted

    # Only attempt interpolation if accel timestamp <= latest gyro timestamp
    latest_gyro_t = gyro_buf[-1][0]

    while len(accel_buf) > 0:
        t_a, ax, ay, az = accel_buf[0]

        # No future gyro sample yet — cannot interpolate
        if t_a > latest_gyro_t:
            break

        g = interp_sample(gyro_buf, t_a)
        if g is None:
            # accel may be earlier than the oldest gyro sample, or gap not yet covered.
            # If t_a < earliest gyro time, this accel sample can never be interpolated — discard it.
            if t_a < gyro_buf[0][0]:
                accel_buf.popleft()
                continue
            break
        # Write as [t, gx, gy, gz, ax_ms2, ay_ms2, az_ms2]:
        # swap accel/gyro columns and convert accel from g to m/s² (* -9.81)
        imu_writer.writerow([t_a, g[0], g[1], g[2], ax * (-GRAVITY), ay * (-GRAVITY), az * (-GRAVITY)])
        emitted += 1
        accel_buf.popleft()

    if emitted > 0:
        f_imu.flush()

    return emitted


def prune_gyro_buffer(gyro_buf, oldest_needed_t):
    """
    Keep at least one gyro sample before oldest_needed_t for interpolation continuity.
    """
    while len(gyro_buf) >= 2 and gyro_buf[1][0] < oldest_needed_t:
        gyro_buf.popleft()


def main():
    session_dir, images_dir = make_session_dirs()
    imu_csv_path = os.path.join(session_dir, "imu.csv")
    accel_raw_csv_path = os.path.join(session_dir, "accel_raw.csv")
    gyro_raw_csv_path = os.path.join(session_dir, "gyro_raw.csv")
    camera_json_path = os.path.join(session_dir, "camera.json")
    calib_txt_path   = os.path.join(session_dir, "calib.txt")

    print(f"[INFO] Dataset root:   {DATASET_ROOT}")
    print(f"[INFO] Session dir:     {session_dir}")
    print(f"[INFO] Images dir:      {images_dir}")
    print(f"[INFO] IMU csv path:    {imu_csv_path}")
    print(f"[INFO] Accel raw path:  {accel_raw_csv_path}")
    print(f"[INFO] Gyro raw path:   {gyro_raw_csv_path}")
    print(f"[INFO] Camera path:     {camera_json_path}")
    print(f"[INFO] Calib txt path:  {calib_txt_path}")
    print("[INFO] IMU mode: gyro interpolated to accel timestamps")

    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server.bind((HOST, PORT))
    server.listen(1)

    print(f"[INFO] Listening on {HOST}:{PORT}")
    print("[INFO] Waiting for iPhone connection...")

    conn, addr = server.accept()
    print(f"[INFO] Connected from {addr}")

    camera_saved = False
    num_images = 0
    num_imu = 0
    num_accel_raw = 0
    num_gyro_raw = 0

    accel_buf = deque()   # (t_ns, ax, ay, az)
    gyro_buf = deque()    # (t_ns, gx, gy, gz)

    with open(imu_csv_path, "w", newline="") as f_imu, \
         open(accel_raw_csv_path, "w", newline="") as f_accel, \
         open(gyro_raw_csv_path, "w", newline="") as f_gyro:

        imu_writer = csv.writer(f_imu)
        accel_writer = csv.writer(f_accel)
        gyro_writer = csv.writer(f_gyro)

        imu_writer.writerow(["# t_ns", "gx", "gy", "gz", "ax", "ay", "az"])
        accel_writer.writerow(["# t_ns", "ax", "ay", "az"])
        gyro_writer.writerow(["# t_ns", "gx", "gy", "gz"])

        while True:
            try:
                header = recv_exact(conn, 4)
                msg_len = struct.unpack(">I", header)[0]
                msg = recv_exact(conn, msg_len)

                msg_type = msg[0]
                payload = msg[1:]

                if msg_type == 1:
                    if len(payload) < 20:
                        continue

                    t_ns, w, h, jpeg_size = struct.unpack(">QIII", payload[:20])
                    jpeg_bytes = payload[20:20 + jpeg_size]

                    if len(jpeg_bytes) != jpeg_size:
                        continue

                    image_path = os.path.join(images_dir, f"{t_ns}.jpg")
                    with open(image_path, "wb") as img_f:
                        img_f.write(jpeg_bytes)

                    num_images += 1

                    arr = np.frombuffer(jpeg_bytes, dtype=np.uint8)
                    img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
                    if img is not None:
                        cv2.putText(
                            img,
                            f"t={t_ns}",
                            (20, 40),
                            cv2.FONT_HERSHEY_SIMPLEX,
                            0.8,
                            (0, 255, 0),
                            2,
                            cv2.LINE_AA,
                        )
                        try:
                            cv2.imshow("iPhone RGB", img)
                            if cv2.waitKey(1) == 27:
                                print("[INFO] ESC pressed, stopping.")
                                break
                        except cv2.error:
                            pass  # headless — skip display

                elif msg_type == 2:
                    if len(payload) != 8 + 1 + 8 * 3:
                        continue

                    t_ns = struct.unpack(">Q", payload[:8])[0]
                    sensor_type = payload[8]
                    x, y, z = struct.unpack(">ddd", payload[9:])

                    if sensor_type == 0:
                        # accel
                        accel_buf.append((t_ns, x, y, z))
                        accel_writer.writerow([t_ns, x, y, z])
                        f_accel.flush()
                        num_accel_raw += 1

                    elif sensor_type == 1:
                        # gyro
                        gyro_buf.append((t_ns, x, y, z))
                        gyro_writer.writerow([t_ns, x, y, z])
                        f_gyro.flush()
                        num_gyro_raw += 1
                    else:
                        continue

                    # On every new sample, try to flush any accel samples that are now syncable
                    emitted_now = flush_synced_imu(accel_buf, gyro_buf, imu_writer, f_imu)
                    num_imu += emitted_now

                    # Prune gyro buffer, keeping the minimum history needed for interpolation
                    if len(accel_buf) > 0:
                        oldest_needed_t = accel_buf[0][0]
                        prune_gyro_buffer(gyro_buf, oldest_needed_t)
                    elif len(gyro_buf) > 2:
                        # No pending accel samples — conservatively keep the last two gyro points
                        while len(gyro_buf) > 2:
                            gyro_buf.popleft()

                elif msg_type == 3:
                    if len(payload) != 4 + 4 + 4 * 5:
                        continue

                    width, height = struct.unpack(">II", payload[:8])
                    fx, fy, cx, cy, lens_position = struct.unpack(">fffff", payload[8:])

                    camera_dict = {
                        "device_model": "iPhone17Pro",
                        "camera_name": "back_wide",
                        "image_width": width,
                        "image_height": height,
                        "fx": fx,
                        "fy": fy,
                        "cx": cx,
                        "cy": cy,
                        "lens_position_locked": lens_position,
                        "orientation": "landscape",
                        "video_rotation_angle": 0,
                        "imu_sync_mode": "gyro_interpolated_to_accel_timestamps"
                    }

                    if not camera_saved:
                        with open(camera_json_path, "w") as jf:
                            json.dump(camera_dict, jf, indent=2)
                        with open(calib_txt_path, "w") as cf:
                            cf.write(f"{fx} {fy} {cx} {cy}\n")
                        camera_saved = True
                        print(f"[INFO] Saved camera.json: {camera_dict}")
                        print(f"[INFO] Saved calib.txt: fx={fx} fy={fy} cx={cx} cy={cy}")

                else:
                    print(f"[WARN] Unknown message type: {msg_type}")

            except KeyboardInterrupt:
                print("[INFO] Interrupted by user.")
                break
            except Exception as e:
                print(f"[ERROR] {e}")
                break

    conn.close()
    server.close()
    cv2.destroyAllWindows()

    print(f"[INFO] Saved {num_images} images.")
    print(f"[INFO] Saved {num_accel_raw} raw accel rows.")
    print(f"[INFO] Saved {num_gyro_raw} raw gyro rows.")
    print(f"[INFO] Saved {num_imu} synced imu rows.")
    print(f"[INFO] Dataset saved to: {session_dir}")


if __name__ == "__main__":
    main()