
import threading
import time
from typing import Any, Dict, Optional

from .InferencePipeline import InferencePipeline, PipelineData
from .camera_source import CameraSource
from .video_source import VideoFileSource
from .result_handler import FileOutputManager
# Import other data sources as they are created

class WorkflowOrchestrator:
    """
    Coordinates the entire data flow from input source to the inference pipeline
    and handles the results.
    """
    def __init__(self, pipeline: InferencePipeline, input_config: Dict[str, Any], output_config: Dict[str, Any]):
        """
        Initializes the WorkflowOrchestrator.

        Args:
            pipeline (InferencePipeline): The configured inference pipeline.
            input_config (Dict[str, Any]): The configuration for the input source.
            output_config (Dict[str, Any]): The configuration for the output.
        """
        self.pipeline = pipeline
        self.input_config = input_config
        self.output_config = output_config
        self.data_source = None
        self.result_handler = None
        self.running = False
        self._stop_event = threading.Event()
        self.frame_callback = None
        self.result_callback = None

    def start(self):
        """
        Starts the workflow, including the data source and the pipeline.
        """
        if self.running:
            print("Workflow is already running.")
            return

        print("Starting workflow orchestrator...")
        self.running = True
        self._stop_event.clear()

        # Create the result handler
        self.result_handler = self._create_result_handler()

        # Create and start the data source
        self.data_source = self._create_data_source()
        if not self.data_source:
            print("Error: Could not create data source. Aborting workflow.")
            self.running = False
            return

        # Set the pipeline's put_data method as the callback
        self.data_source.set_data_callback(self.pipeline.put_data)

        # Set the result callback on the pipeline - always set this regardless of result_handler
        # The handle_result method will decide what to do with the results
        self.pipeline.set_result_callback(self.handle_result)

        # Start the pipeline
        self.pipeline.initialize()
        self.pipeline.start()

        # Start the data source
        self.data_source.start()

        print("🚀 Workflow orchestrator started successfully.")
        print(f"📊 Pipeline: {self.pipeline.pipeline_name}")
        print(f"🎥 Input: {self.input_config.get('source_type', 'Unknown')} source")
        print(f"💾 Output: {self.output_config.get('output_type', 'Unknown')} destination")
        print("🔄 Inference pipeline is now processing data...")
        print("📡 Inference results will appear below:")
        print("="*60)

    def stop(self):
        """
        Stops the workflow gracefully.
        """
        if not self.running:
            return

        print("🛑 Stopping workflow orchestrator...")
        self.running = False
        self._stop_event.set()

        if self.data_source:
            self.data_source.stop()
            print("📹 Data source stopped")
        
        if self.pipeline:
            self.pipeline.stop()
            print("⚙️  Inference pipeline stopped")

        print("✅ Workflow orchestrator stopped successfully.")
        print("="*60)

    def set_frame_callback(self, callback):
        """
        Sets the callback function for frame updates.
        """
        self.frame_callback = callback
    
    def set_result_callback(self, callback):
        """
        Sets the callback function for inference results.
        """
        self.result_callback = callback

    def _create_data_source(self) -> Optional[Any]:
        """
        Creates the appropriate data source based on the input configuration.
        """
        source_type = self.input_config.get('source_type', '').lower()
        print(f"Creating data source of type: {source_type}")

        if source_type == 'camera':
            return CameraSource(
                camera_index=self.input_config.get('device_id', 0),
                resolution=self._parse_resolution(self.input_config.get('resolution')),
                fps=self.input_config.get('fps', 30),
                data_callback=self.pipeline.put_data,
                frame_callback=self.frame_callback
            )
        elif source_type == 'file':
            # Assuming 'file' means video file for now
            return VideoFileSource(
                file_path=self.input_config.get('source_path', ''),
                loop=True,  # Or get from config if available
                data_callback=self.pipeline.put_data,
                frame_callback=self.frame_callback
            )
        # Add other source types here (e.g., 'rtsp stream', 'image file')
        else:
            print(f"Error: Unsupported source type '{source_type}'")
            return None

    def _create_result_handler(self) -> Optional[Any]:
        """
        Creates the appropriate result handler based on the output configuration.
        """
        output_type = self.output_config.get('output_type', '').lower()
        print(f"Creating result handler of type: {output_type}")

        if output_type == 'file':
            return FileOutputManager(
                base_path=self.output_config.get('destination', './output')
            )
        # Add other result handlers here
        else:
            print(f"Warning: Unsupported output type '{output_type}'. No results will be saved.")
            return None

    def handle_result(self, result_data: PipelineData):
        """
        Callback function to handle results from the pipeline.
        """
        try:
            # Convert PipelineData to a dictionary for serialization
            result_dict = {
                "pipeline_id": result_data.pipeline_id,
                "timestamp": result_data.timestamp,
                "metadata": result_data.metadata,
                "stage_results": result_data.stage_results
            }
            
            # Save to file if result handler exists
            if self.result_handler:
                self.result_handler.save_result(
                    result_dict,
                    self.pipeline.pipeline_name,
                    format=self.output_config.get('format', 'json').lower()
                )
            
            # Always call the result callback if set (for GUI updates)
            if self.result_callback:
                self.result_callback(result_dict)
                
        except Exception as e:
            print(f"❌ Error handling result: {e}")

    def _parse_resolution(self, resolution_str: Optional[str]) -> Optional[tuple[int, int]]:
        """
        Parses a resolution string (e.g., '1920x1080') into a tuple.
        """
        if not resolution_str:
            return None
        try:
            width, height = map(int, resolution_str.lower().split('x'))
            return (width, height)
        except ValueError:
            print(f"Warning: Invalid resolution format '{resolution_str}'. Using default.")
            return None
