# Keyframe capture for server-side VSLAM
import asyncio
import base64
import subprocess
import time
import uuid

from .config import (
    log, _ros2_cache, _ros2_lock,
    _keyframe_buffer, _kf_wifi_cache, _kf_wifi_cache_ts,
    KF_MAX_BUFFER, KF_BATCH_SIZE, KF_MIN_DISPLACEMENT, KF_MIN_ROTATION, KF_MIN_INTERVAL,
    KF_THUMB_WIDTH, KF_THUMB_HEIGHT, KF_THUMB_QUALITY, KF_WIFI_SCAN_INTERVAL,
)
import ai_teammate_ros2_bridge.config as _cfg


def _should_capture_keyframe() -> bool:
    """Decide whether to capture a keyframe based on odometry change."""
    now = time.time()
    if now - _cfg._last_kf_ts < KF_MIN_INTERVAL:
        return False

    with _ros2_lock:
        odom = _ros2_cache.get("odometry")
        frame = _ros2_cache.get("camera_frame")
    if not odom or not frame:
        return False

    if _cfg._last_kf_pose is None:
        return True

    dx = odom.get("x", 0) - _cfg._last_kf_pose.get("x", 0)
    dy = odom.get("y", 0) - _cfg._last_kf_pose.get("y", 0)
    displacement = (dx * dx + dy * dy) ** 0.5
    dtheta = abs(odom.get("theta", 0) - _cfg._last_kf_pose.get("theta", 0))

    return displacement >= KF_MIN_DISPLACEMENT or dtheta >= KF_MIN_ROTATION


def _capture_keyframe():
    """Capture current frame + odometry into keyframe buffer."""
    import cv2
    import numpy as np

    with _ros2_lock:
        frame = _ros2_cache.get("camera_frame")
        odom = dict(_ros2_cache.get("odometry", {}))

    if not frame:
        return

    # Thumbnail (160x120, quality 40)
    img = cv2.imdecode(np.frombuffer(frame, np.uint8), cv2.IMREAD_COLOR)
    if img is not None:
        thumb = cv2.resize(img, (KF_THUMB_WIDTH, KF_THUMB_HEIGHT))
        _, jpeg = cv2.imencode('.jpg', thumb, [cv2.IMWRITE_JPEG_QUALITY, KF_THUMB_QUALITY])
        thumb_b64 = base64.b64encode(jpeg.tobytes()).decode("ascii")
    else:
        thumb_b64 = None

    # Wheel odometry scale
    scale_factor = None
    if _cfg._last_kf_pose:
        dx = odom.get("x", 0) - _cfg._last_kf_pose.get("x", 0)
        dy = odom.get("y", 0) - _cfg._last_kf_pose.get("y", 0)
        scale_factor = (dx * dx + dy * dy) ** 0.5

    kf = {
        "keyframe_id": str(uuid.uuid4())[:8],
        "timestamp": int(time.time() * 1000),
        "pose": odom,
        "features": [],
        "image_thumbnail": thumb_b64,
        "quality": {"feature_count": 0, "blur_score": 0.0},
        "scale_factor": scale_factor,
        "scale_confidence": 0.9 if scale_factor else None,
        "scale_source": "odometry" if scale_factor else None,
        "fingerprint": _cfg._kf_wifi_cache,
    }

    _keyframe_buffer.append(kf)
    if len(_keyframe_buffer) > KF_MAX_BUFFER:
        _keyframe_buffer.pop(0)

    _cfg._last_kf_pose = odom
    _cfg._last_kf_ts = time.time()
    _cfg._kf_total_captured += 1
    log.info(f"Keyframe captured #{_cfg._kf_total_captured} (buffer: {len(_keyframe_buffer)})")


async def _maybe_upload_keyframe_batch():
    """Upload batch when buffer reaches batch size."""
    if len(_keyframe_buffer) < KF_BATCH_SIZE:
        return
    await _flush_keyframe_buffer()


async def _flush_keyframe_buffer():
    """Upload keyframe buffer to Gateway and clear it."""
    from .ws_client import _send_service_request

    if not _keyframe_buffer:
        return

    batch = list(_keyframe_buffer)
    _keyframe_buffer.clear()

    for kf in batch:
        fp = kf.pop("fingerprint", None)
        if fp:
            kf.setdefault("metadata", {})["wifi_fingerprint"] = fp

    svc_data = {
        "service": "vslam_keyframe_batch",
        "radio_scale_enabled": True,
        "keyframes": batch,
        "merge_map": True,
        "device_type": "robot",
    }

    log.info(f"Uploading keyframe batch: {len(batch)} keyframes")
    result = await _send_service_request(svc_data, timeout=30.0)
    if result.get("status") == "ok":
        _cfg._kf_last_upload_ts = time.time()
        log.info("Keyframe batch uploaded successfully")
    else:
        log.warning(f"Keyframe batch upload failed: {result.get('message', '')}")
        _keyframe_buffer[:0] = batch


async def _kf_wifi_scan_loop():
    """Background WiFi scan — periodically refresh cache."""
    while _cfg._kf_capture_active:
        try:
            result = subprocess.run(
                ["nmcli", "-t", "-f", "SSID,SIGNAL,BSSID", "dev", "wifi", "list", "--rescan", "yes"],
                capture_output=True, text=True, timeout=10,
            )
            aps = []
            for line in result.stdout.strip().split("\n"):
                parts = line.split(":")
                if len(parts) >= 3:
                    aps.append({
                        "ssid": parts[0],
                        "signal": int(parts[1]) if parts[1].lstrip('-').isdigit() else 0,
                        "bssid": ":".join(parts[2:8]).replace("\\", ""),
                    })
            if aps:
                _cfg._kf_wifi_cache = {"wifi": aps}
                _cfg._kf_wifi_cache_ts = time.time()
        except Exception as e:
            log.debug(f"WiFi scan error: {e}")
        await asyncio.sleep(KF_WIFI_SCAN_INTERVAL)
