from __future__ import annotations

import asyncio
import os

import cv2
from aiortc.mediastreams import MediaStreamError
from loguru import logger

from app.audio.rep_announcer import RepAnnouncer
from app.exercises.dead_bug.detector import DeadBugDetector
from app.rendering.window_display import close_window, is_esc_pressed, show_frame
from configs.default import (
    EXTENSION_CONFIRM_FRAMES,
    MODEL_PATH,
    PREFER_GPU,
    PROCESS_EVERY_N_FRAMES,
    REP_ANNOUNCER_ENABLED,
    REP_ANNOUNCER_RATE,
    REP_ANNOUNCER_VOLUME,
    RESET_CONFIRM_FRAMES,
    VISIBILITY_THRESHOLD,
)


def _format_pose_debug(pose_result) -> str:
    metrics = pose_result.metrics
    if metrics is None:
        return "metrics=None"
    return (
        f"side={pose_result.side}, standard={pose_result.is_standard}, "
        f"angles(le={metrics.left_elbow_angle:.1f}, re={metrics.right_elbow_angle:.1f}, "
        f"lk={metrics.left_knee_angle:.1f}, rk={metrics.right_knee_angle:.1f}), "
        f"extended(la={metrics.left_arm_extended}, ra={metrics.right_arm_extended}, "
        f"ll={metrics.left_leg_extended}, rl={metrics.right_leg_extended})"
    )


class VideoReceiver:
    def __init__(self, track) -> None:
        self._track = track

    async def run(self) -> None:
        logger.info("Start receiving video frames, process_every_n={}", PROCESS_EVERY_N_FRAMES)

        frame_count = 0
        processed_count = 0
        detector = DeadBugDetector(
            model_path=MODEL_PATH,
            visibility_threshold=VISIBILITY_THRESHOLD,
            extension_confirm_frames=EXTENSION_CONFIRM_FRAMES,
            reset_confirm_frames=RESET_CONFIRM_FRAMES,
            prefer_gpu=PREFER_GPU,
        )
        announcer = RepAnnouncer(
            enabled=REP_ANNOUNCER_ENABLED,
            rate=REP_ANNOUNCER_RATE,
            volume=REP_ANNOUNCER_VOLUME,
        )
        last_announced_rep = 0
        last_pose_result = None
        last_annotated = None

        try:
            while True:
                frame = await self._track.recv()
                frame_count += 1
                raw_img = frame.to_ndarray(format="bgr24")
                timestamp_ms = int(frame.time * 1000) if frame.time is not None else frame_count * 33

                if frame_count % PROCESS_EVERY_N_FRAMES == 0 or last_pose_result is None:
                    processed_count += 1
                    last_annotated, last_pose_result = detector.process_frame(raw_img, timestamp_ms)
                    if last_pose_result.rep_count > last_announced_rep:
                        last_announced_rep = last_pose_result.rep_count
                        announcer.announce_count(last_announced_rep)

                display_img = last_annotated if last_annotated is not None else raw_img
                show_frame(display_img)

                if frame_count % 100 == 0:
                    logger.info(
                        "Received {} frames, processed={}, raw_shape={}, reps={}, phase={}, feedback={}, {}",
                        frame_count,
                        processed_count,
                        raw_img.shape,
                        last_pose_result.rep_count if last_pose_result is not None else 0,
                        last_pose_result.phase.value if last_pose_result is not None else "none",
                        " | ".join(last_pose_result.feedback) if last_pose_result is not None else "",
                        _format_pose_debug(last_pose_result) if last_pose_result is not None else "metrics=None",
                    )

                if is_esc_pressed():
                    logger.info("ESC pressed, closing display")
                    break

        except asyncio.CancelledError:
            logger.info("Video receive task cancelled")
        except MediaStreamError:
            logger.info("Video track ended")
        except Exception as e:
            logger.exception(f"Video receive error: {e!r}")
        finally:
            announcer.close()
            detector.close()
            close_window()