import cv2
import numpy as np
import time
from typing import List, Dict, Any

class InputLoader:
    def __init__(self, mode: str, config: Dict[str, Any] = None):
        """
        mode: 'rtsp', 'random', 'image'
        """
        self.mode = mode
        self.stream = None
        
        if mode == 'rtsp':
            from nxva.streaming import MultiStreaming
            self.stream = MultiStreaming(
                config=config.get('streaming_config'),
                reconnect_interval=30,
                verbose=False
            )
            self.stream.run()
            # Wait for first frames
            print("Waiting for RTSP stream to initialize...")
            for _ in range(100):
                frames = self.stream.get_frames()
                if isinstance(frames, tuple) and len(frames[0]) > 0:
                    break
                elif isinstance(frames, np.ndarray) and len(frames) > 0:
                    break
                time.sleep(0.1)

    def get_input(self, model_config: Dict[str, Any]):
        """
        Returns a list of frames (images) or a batched numpy array.
        """
        batch_size = model_config.get('batch_size', 1)
        h, w = model_config.get('input_size', (640, 640))
        input_source = model_config.get('input_source', 0) # Index for RTSP outputs

        if self.mode == 'rtsp':
            frames = self.stream.get_frames()
            # frames can be a single array or a tuple of arrays
            if isinstance(frames, tuple):
                return frames[input_source]
            return frames

        elif self.mode == 'random':
            return np.random.randint(0, 255, (batch_size, h, w, 3), dtype=np.uint8)

        elif self.mode == 'image':
            image_path = model_config.get('image_path', './src/test_imgs/test_pose.jpg')
            img = cv2.imread(image_path)
            if img is None:
                # Fallback to random if image not found
                return np.random.randint(0, 255, (batch_size, h, w, 3), dtype=np.uint8)
            img = cv2.resize(img, (w, h))
            batch = [img] * batch_size
            return batch

        return None

    def release(self):
        if self.stream:
            self.stream.release()
