import os
import csv
import argparse
import importlib
import threading
from datetime import datetime
import traceback
from psychopy import event, core

from .utils import gui, eyetracker
from .utils.voice import VoiceRecorder
from .utils.custom_logger import CustomLogger
from .configs.custom_config import CustomConfig
from .datasets.custom_dataset import TextDataset, ImageDataset, TimeSeriesDataset


LOGGER = None


def main(config, loop_count, eyetracker_config_file,
         enable_eyetracker, enable_voice,
         raw_data, enable_psychopy):
    """
    Main execution function for the experiment or data recording session.

    This function handles both PsychoPy GUI-based trials and headless eye-tracker
    recording. Depending on the options, it can:
      - Display stimuli and buttons via PsychoPy
      - Record gaze data and/or full raw eye-tracker samples
      - Record voice input
      - Apply a custom model to predict outputs based on displayed stimuli

    The function ensures proper resource management and logging of all actions
    and gracefully handles keyboard interrupts and unexpected exceptions.

    Parameters
    ----------
    config : CustomConfig
        Configuration object loaded from a YAML file, containing display, model,
        output, and instruction settings.

    loop_count : int
        Maximum number of trials or data items to process. If higher than the
        dataset length, the dataset length is used instead.

    eyetracker_config_file : str
        Path to the YAML configuration file for the Tobii eye-tracker.

    enable_eyetracker : bool
        If True, launches the eye-tracker using `launchHubServer()` and records
        gaze or raw samples.

    enable_voice : bool
        If True, records voice input during each trial or in headless mode.

    raw_data : bool
        If True, records full Tobii raw samples instead of filtered gaze
        positions and pupil sizes.

    enable_psychopy : bool
        Controls whether the PsychoPy GUI is used:
          - True: Displays GUI windows, stimuli, and buttons, recording gaze
            and model outputs per trial.
          - False: Runs in headless mode, continuously recording raw gaze
            samples (if enabled) and/or voice until interrupted by the user.
            No screenshots, classification, or buttons are used in this mode.

    Returns
    -------
    None
        This function writes experiment or recording data to CSV files in the
        output folder. It does not return any Python objects.

    Files Created
    -------------
    - `data.csv` (PsychoPy mode):
        Contains trial-specific data with the following columns:
            'screenshot_file', 'input_data', 'classification', 'user_classification',
            'gaze_data', 'objects_bboxes', 'voice_file', 'voice_start_timestamp'
    - `raw_stream.csv` (Headless mode):
        Contains timestamped raw eye-tracker samples when `enable_psychopy=False`
        and `enable_eyetracker=True`.

    Exception Handling
    ------------------
    - KeyboardInterrupt: Gracefully stops trials or headless recording, ensuring
      voice threads and eye-tracker recording are stopped and resources are released.
    - Exception: Logs any unexpected error during:
        - CSV writing
        - Eye-tracker sample retrieval
        - Voice recording
        - Model processing
      Errors do **not** stop the experiment unless critical.

    Notes
    -----
    - When running headless (`enable_psychopy=False`), no screenshots, data, or
      classification columns are written—only raw eye-tracker samples and voice.
    - Voice recording is threaded to avoid blocking the trial loop.
    - Eye-tracker recording is started immediately after initialization and
      stopped at the end, ensuring all buffered samples are captured.

    Examples
    --------
    >>> config = CustomConfig("configs/config.yaml")
    >>> main(
    ...     config=config,
    ...     loop_count=10,
    ...     eyetracker_config_file="configs/eyetracker_config.yaml",
    ...     enable_eyetracker=True,
    ...     enable_model=False,
    ...     enable_voice=True,
    ...     raw_data=True,
    ...     enable_psychopy=False
    ... )
    Running headless will continuously record raw eye-tracker data and voice
    until the user presses Ctrl+C.
    """
    
    if config.dataset_type == "text":
        dataset = TextDataset(config)
    elif config.dataset_type == "time_series":
        dataset = TimeSeriesDataset(config)
    elif config.dataset_type == "image":
        dataset = ImageDataset(config)
    monitor, window, buttons = None, None, None

    # --- Optional PsychoPy GUI setup ---
    if enable_psychopy:
        monitor = gui.prepare_monitor(config)
        window = gui.prepare_window(config, monitor)
        buttons = gui.prepare_buttons(config, window, dataset)

    # --- Initialize Eye Tracker ---
    io, tracker = (None, None)
    buffer, last_event_id = {}, 0
    if enable_eyetracker:
        io, tracker = eyetracker.launch_hub_server(eyetracker_config_file, window if enable_psychopy else None)
 

    # --- Output setup ---
    output_config = config.get_output_config()
    output_folder = os.path.join(output_config["folder"], datetime.now().strftime("%Y%m%d_%H%M%S"))
    os.makedirs(output_folder, exist_ok=True)
    output_file = os.path.join(output_folder, "data.csv")

    if loop_count > len(dataset.data):
        loop_count = len(dataset.data)

    # --- Show instructions if PsychoPy enabled ---
    if enable_psychopy:
        intro_text = config.get_instructions_config()["intro"]
        gui.show_instructions(window, intro_text)
        window.winHandle.activate()

    # Detect if this is a MouseGaze tracker
    iohub_config = CustomConfig.read_config(eyetracker_config_file)
    is_mouse_tracker =  eyetracker.is_mouse_eyetracker(iohub_config)
    move_button_idx = eyetracker.get_mouse_move_button_idx(iohub_config)

    # --- Main try/finally block to ensure cleanup ---
    try:
        with open(output_file, "w", newline='', encoding='utf-8') as csvfile:
            fieldnames = [
                'screenshot_file', 'input_data', 'classification', 'user_classification',
                'gaze_data', 'objects_bboxes', 'voice_file', 'voice_start_timestamp'
            ]
            writer = csv.DictWriter(csvfile, fieldnames=fieldnames, quoting=csv.QUOTE_MINIMAL, delimiter=';')
            writer.writeheader()

            last_click_time = core.getTime()
            focus_time = 2.0
            debounce_time = 0.5

            # --- Headless mode (no PsychoPy GUI) ---
            if not enable_psychopy and enable_eyetracker:
                LOGGER.info("Running in headless mode (no GUI). Press Ctrl+C to stop.")
                try:
                    with open(os.path.join(output_folder, "raw_stream.csv"), "w", newline='', encoding="utf-8") as rawfile:
                        raw_writer = csv.writer(rawfile)
                        raw_writer.writerow(["timestamp", "sample_data"])

                        while True:
                            try:
                                buffer, last_event_id = eyetracker.poll_tracker_events(tracker, buffer, last_event_id)
                                sample = eyetracker.extract_full_raw_event(buffer)
                                if sample:
                                    raw_writer.writerow([datetime.now().isoformat(), sample])
                                core.wait(0.01)
                            except KeyboardInterrupt:
                                LOGGER.info("Recording stopped by user.")
                                break
                            except Exception as e:
                                LOGGER.error(f"Error during headless recording: {e}")
                finally:
                    if tracker:
                        tracker.setRecordingState(False)
                        io.quit()
                return

            # --- PsychoPy GUI loop ---
            for i, sample in enumerate(dataset.data):
                if i == loop_count:
                    break

                objects_bboxes = {}
                data = sample['data']
                classification = sample['class'].lower()
                screenshot_path, objects_bboxes = gui.draw_window(config, 
                                            window,
                                            sample,
                                            dataset,
                                            buttons,
                                            focus_time,
                                            output_folder
                                        )

                gaze_data = [] 
                next_data = False
                voice_thread, voice_stop_event, voice_filename, voice_start_time = None, None, None, None

                if enable_voice:
                    voice_filename = os.path.join(output_folder, f"voice_{i:03d}.wav")
                    voice_stop_event = threading.Event()
                    voice_start_time = core.getTime()
                    voice_thread = threading.Thread(
                        target=VoiceRecorder.record_voice,
                        args=(voice_filename,),
                        kwargs={'stop_event': voice_stop_event}
                    )
                    voice_thread.start()

                try:
                    while not next_data:
                        current_time = core.getTime()
                        mouse = event.Mouse(win=window)

                        if enable_eyetracker:
                            try:
                                # --- Poll buffer and avoid duplicates ---
                                buffer, last_event_id = eyetracker.poll_tracker_events(tracker, buffer, last_event_id)
                                # Determine if we should record gaze samples
                                should_record = True
                                if is_mouse_tracker:
                                    pressed = mouse.getPressed()
                                    should_record = pressed[move_button_idx]  # only record while configured button pressed

                                if should_record:
                                    current_time = core.getTime()
                                    if raw_data:
                                        full_sample = eyetracker.extract_full_raw_event(buffer, system_time=current_time)
                                        if full_sample:
                                            gaze_data.extend(full_sample)
                                    else:
                                        eye_sample = eyetracker.extract_eye_gaze_events(buffer, system_time=current_time, config=config)
                                        if eye_sample:
                                            gaze_data.extend(eye_sample)
                                elif enable_eyetracker and tracker:
                                    tracker.clearEvents()
                                    buffer.clear()
                                 

                            except Exception as e:
                                import traceback
                                LOGGER.error(f"Eye tracker error: {e}")
                                traceback.print_exc()

                        # --- Check button presses ---
                        for rect, text, label in buttons:
                            if mouse.isPressedIn(rect) and (current_time - last_click_time) > debounce_time:
                                user_classification = text.text

                                if enable_voice and voice_thread:
                                    voice_stop_event.set()
                                    voice_thread.join()

                                row_data = {
                                    'screenshot_file': screenshot_path,
                                    'input_data': data,
                                    'classification': classification,
                                    'user_classification': user_classification.lower(),
                                    'gaze_data': gaze_data,
                                    'objects_bboxes': objects_bboxes,
                                    'voice_file': voice_filename,
                                    'voice_start_timestamp': round(voice_start_time, 4) if voice_start_time else None
                                }
                                writer.writerow(row_data)

                                if label == "functional_quit":
                                    if enable_eyetracker and tracker:
                                        tracker.clearEvents()
                                    LOGGER.info("Exit button pressed. Quitting...")
                                    return

                                last_click_time = current_time
                                next_data = True

                        core.wait(0.01)

                except KeyboardInterrupt:
                    LOGGER.info("Keyboard interrupt received. Exiting current trial...")
                    if enable_voice and voice_thread:
                        voice_stop_event.set()
                        voice_thread.join()
                    raise
                except Exception as e:
                    LOGGER.error(f"Error during trial loop: {e}")

    except KeyboardInterrupt:
        LOGGER.info("Experiment interrupted by user.")
    except Exception as e:
        LOGGER.error(f"Failed to save the data: {e}")
        import traceback
        traceback.print_exc()
    finally:
        # --- Cleanup ---
        if enable_psychopy and window:
            outro_text = config.get_instructions_config()["outro"]
            gui.show_instructions(window, outro_text, key='escape')
            window.close()

        if enable_eyetracker and tracker:
            tracker.setRecordingState(False)
            io.quit()

        if enable_voice and voice_thread and voice_thread.is_alive():
            voice_stop_event.set()
            voice_thread.join()

    LOGGER.info(f"Data saved to {output_file}")


def cli():
    parser = argparse.ArgumentParser()
    parser = argparse.ArgumentParser()
    parser.add_argument('--config_file', type=str, default='configs/config.yaml', help='Path to YAML script config file')
    parser.add_argument('--eyetracker_config_file', type=str, default='configs/eyetracker_config.yaml', help='Path to YAML eyetracker config file')

    # Boolean flags
    parser.add_argument('--enable_eyetracker', action='store_true', help='Launch script with launchHubServer (needs connected eyetracker if set)')
    parser.add_argument('--enable_voice', action='store_true', help='Enable voice recording and processing')
    parser.add_argument('--raw_data', action='store_true', help='Record full Tobii raw samples instead of filtered gaze position')
    parser.add_argument('--disable_psychopy', action='store_true', help='Run headless (no GUI) and continuously record gaze + voice until stopped')

    # Other arguments
    parser.add_argument('--loop_count', type=int, default=10, help='Number of times that different data will be displayed before the script exits')
    parser.add_argument('--log_level', type=str, default='info', help='Main logger level ("info", "debug", "warning", "error", "critical")')

    args = parser.parse_args()

    global LOGGER
    LOGGER = CustomLogger(args.log_level, __name__).logger
    config = CustomConfig(args.config_file)

    main(
        config=config,
        loop_count=args.loop_count,
        eyetracker_config_file=args.eyetracker_config_file,
        enable_eyetracker=args.enable_eyetracker,
        enable_voice=args.enable_voice,
        raw_data=args.raw_data,
        enable_psychopy=not args.disable_psychopy
    )


if __name__ == "__main__":
    cli()



    