import cv2
import numpy as np
import csv
import sys
import time
from pathlib import Path
import argparse
import pandas as pd

sys.path.append(str(Path(__file__).parent.parent))
from skeleton.mediapipe_detector import MediaPipeDetector
from ml_gestures_dynamic.feature_extractor import extract_sequence

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--label', required=True, help='Название жеста (например, wave_left)')
    parser.add_argument('--output', default='dynamic_data.csv', help='CSV файл для сохранения')
    parser.add_argument('--camera', type=int, default=0)
    parser.add_argument('--duration', type=float, default=3.0, help='Длительность записи (сек)')
    args = parser.parse_args()
    
    output_path = Path(args.output)
    
    # Создаём файл с заголовком, если его ещё нет
    if not output_path.exists():
        with open(output_path, 'w', newline='') as f:
            writer = csv.writer(f)
            header = ['label', 'sequence_id', 'frame_idx'] + [f'f{i}' for i in range(99)]
            writer.writerow(header)

    detector = MediaPipeDetector()
    cap = cv2.VideoCapture(args.camera)
    if not cap.isOpened():
        print("Камера не найдена")
        sys.exit(1)

    # Определяем следующий ID последовательности
    try:
        df = pd.read_csv(output_path)
        next_id = df['sequence_id'].max() + 1 if not df.empty else 0
    except:
        next_id = 0

    print(f"Запись жеста: {args.label}. Нажмите SPACE для начала, q для выхода.")
    recording = False
    start_time = 0
    sequence = []

    while True:
        ret, frame = cap.read()
        if not ret:
            break
        frame = cv2.flip(frame, 1)  # зеркало для удобства
        result = detector.detect(frame)
        if result['success']:
            landmarks = result['landmarks']
            vis = detector.draw_landmarks(frame, result['pose_landmarks'])
        else:
            vis = frame
            landmarks = None

        # Если мы в режиме записи
        if recording:
            if landmarks is not None:
                sequence.append(landmarks)
            elapsed = time.time() - start_time
            # Отображаем прогресс записи
            cv2.putText(vis, f"RECORDING... {elapsed:.1f}/{args.duration}", (10, 30),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0,0,255), 2)
            if elapsed >= args.duration:
                recording = False
                seq_features = extract_sequence(sequence)
                with open(output_path, 'a', newline='') as f:
                    writer = csv.writer(f)
                    for i, feat in enumerate(seq_features):
                        writer.writerow([args.label, next_id, i] + feat.tolist())
                print(f"Сохранено {len(sequence)} кадров для жеста {args.label}, ID={next_id}")
                sequence = []
                next_id += 1
        else:
            # Ожидание нажатия пробела
            cv2.putText(vis, f"Press SPACE to record '{args.label}'", (10, 30),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0,255,0), 2)

        cv2.imshow('Record dynamic gesture', vis)
        key = cv2.waitKey(1) & 0xFF
        if key == ord('q'):
            break
        if key == ord(' ') and not recording:
            # Обратный отсчёт на живом видео
            for i in range(3, 0, -1):
                # Получаем свежий кадр для отсчёта
                ret, frame = cap.read()
                if not ret:
                    break
                frame = cv2.flip(frame, 1)
                result = detector.detect(frame)
                if result['success']:
                    vis = detector.draw_landmarks(frame, result['pose_landmarks'])
                else:
                    vis = frame
                cv2.putText(vis, f"Starting in {i}...", (frame.shape[1]//2-100, frame.shape[0]//2),
                            cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0,0,255), 3)
                cv2.imshow('Record dynamic gesture', vis)
                cv2.waitKey(1000)
            # Показываем "GO!"
            ret, frame = cap.read()
            if ret:
                frame = cv2.flip(frame, 1)
                result = detector.detect(frame)
                if result['success']:
                    vis = detector.draw_landmarks(frame, result['pose_landmarks'])
                else:
                    vis = frame
                cv2.putText(vis, "GO!", (frame.shape[1]//2-50, frame.shape[0]//2),
                            cv2.FONT_HERSHEY_SIMPLEX, 2, (0,255,0), 3)
                cv2.imshow('Record dynamic gesture', vis)
                cv2.waitKey(500)
            # Начинаем запись
            recording = True
            start_time = time.time()
            sequence = []

    cap.release()
    cv2.destroyAllWindows()

if __name__ == '__main__':
    main()