#!/usr/bin/env python3
"""
Quick test for event-driven stream processing - runs for 20 seconds.
"""

import asyncio
import os
import logging
from dotenv import load_dotenv

from services import StreamConnectionManager, YOLOv8Utils, COCO_CLASSES

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


async def main():
    """Quick test with callback pattern"""
    logger.info("=== Quick Event-Driven Test (20 seconds) ===")

    # Load environment variables
    load_dotenv()
    camera_url = os.getenv('CAMERA_URL_1')
    if not camera_url:
        logger.error("CAMERA_URL_1 not found in .env file")
        return

    # Create manager
    manager = StreamConnectionManager(
        gpu_id=0,
        batch_size=16,
        force_timeout=0.05,  # 50ms
        poll_interval=0.01,  # 100 FPS
    )

    # Initialize with YOLOv8 model
    model_path = "models/yolov8n.trt"
    logger.info(f"Initializing with model: {model_path}")

    await manager.initialize(
        model_path=model_path,
        model_id="yolo",
        preprocess_fn=YOLOv8Utils.preprocess,
        postprocess_fn=YOLOv8Utils.postprocess,
    )

    result_count = 0

    # Define callback for tracking results
    def on_tracking_result(result):
        nonlocal result_count
        result_count += 1

        if result_count % 5 == 0:  # Log every 5th result
            logger.info(f"[{result.stream_id}] Frame {result.metadata.get('frame_number', 0)}")
            logger.info(f"  Tracked objects: {len(result.tracked_objects)}")

            for obj in result.tracked_objects[:3]:  # Show first 3
                class_name = COCO_CLASSES.get(obj.class_id, f"Class {obj.class_id}")
                logger.info(
                    f"    Track ID {obj.track_id}: {class_name}, "
                    f"conf={obj.confidence:.2f}"
                )

    def on_error(error):
        logger.error(f"Stream error: {error}")

    # Connect to stream
    logger.info(f"Connecting to stream...")
    connection = await manager.connect_stream(
        rtsp_url=camera_url,
        stream_id="test_camera",
        on_tracking_result=on_tracking_result,
        on_error=on_error,
    )

    # Monitor for 20 seconds with stats updates
    for i in range(4):  # 4 x 5 seconds = 20 seconds
        await asyncio.sleep(5)

        stats = manager.get_stats()
        model_stats = stats['model_controller']

        logger.info(f"\n=== Stats Update {i+1}/4 ===")
        logger.info(f"Results received: {result_count}")
        logger.info(f"Buffer A: {model_stats['buffer_a_size']} ({model_stats['buffer_a_state']})")
        logger.info(f"Buffer B: {model_stats['buffer_b_size']} ({model_stats['buffer_b_state']})")
        logger.info(f"Active buffer: {model_stats['active_buffer']}")
        logger.info(f"Total frames processed: {model_stats['total_frames_processed']}")
        logger.info(f"Total batches: {model_stats['total_batches_processed']}")
        logger.info(f"Avg batch size: {model_stats['avg_batch_size']:.2f}")

    # Final statistics
    stats = manager.get_stats()
    logger.info("\n=== Final Statistics ===")
    logger.info(f"Total results received: {result_count}")
    logger.info(f"Manager: {stats['manager']}")
    logger.info(f"Model Controller: {stats['model_controller']}")
    logger.info(f"Connection: {stats['connections']['test_camera']}")

    # Cleanup
    logger.info("\nShutting down...")
    await manager.shutdown()
    logger.info("Test complete!")


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logger.info("\nInterrupted by user")
    except Exception as e:
        logger.error(f"Error: {e}", exc_info=True)