import glob
import logging
import os
import subprocess
from pathlib import Path
from typing import TYPE_CHECKING, Any, Dict, Optional, Self

from aind_behavior_services.utils import run_bonsai_process
from typing_extensions import overload, override

from aind_behavior_experiment_launcher.ui import DefaultUIHelper, UiHelper

from ._base import App

logger = logging.getLogger(__name__)

VISUALIZERS_DIR = "VisualizerLayouts"

if TYPE_CHECKING:
    from aind_behavior_experiment_launcher.launcher.behavior_launcher import BehaviorLauncher
else:
    BehaviorLauncher = Any


class BonsaiApp(App):
    """
    A class to manage the execution of Bonsai workflows.

    Attributes:
        executable (os.PathLike): Path to the Bonsai executable.
        workflow (os.PathLike): Path to the Bonsai workflow file.
        is_editor_mode (bool): Whether to run Bonsai in editor mode.
        is_start_flag (bool): Whether to use the start flag when running Bonsai.
        layout (Optional[os.PathLike | str]): Path to the visualizer layout file.
        layout_directory (Optional[os.PathLike]): Directory containing visualizer layouts.
        additional_properties (Optional[Dict[str, str]]): Additional properties to pass to Bonsai.
        cwd (Optional[os.PathLike]): Working directory for the Bonsai process.
        timeout (Optional[float]): Timeout for the Bonsai process.
        print_cmd (bool): Whether to print the command being executed.
        ui_helper (UiHelper): Helper for user interface interactions.
        _result (Optional[subprocess.CompletedProcess]): Result of the Bonsai process execution.
    """

    def __init__(
        self,
        /,
        workflow: os.PathLike,
        executable: os.PathLike = Path("./bonsai/bonsai.exe"),
        is_editor_mode: bool = True,
        is_start_flag: bool = True,
        layout: Optional[os.PathLike] = None,
        layout_dir: Optional[os.PathLike] = None,
        additional_properties: Optional[Dict[str, str]] = None,
        cwd: Optional[os.PathLike] = None,
        timeout: Optional[float] = None,
        print_cmd: bool = False,
        ui_helper: Optional[UiHelper] = None,
        **kwargs,
    ) -> None:
        """
        Initializes the BonsaiApp instance.

        Args:
            workflow (os.PathLike): Path to the Bonsai workflow file.
            executable (os.PathLike): Path to the Bonsai executable. Defaults to "./bonsai/bonsai.exe".
            is_editor_mode (bool): Whether to run Bonsai in editor mode. Defaults to True.
            is_start_flag (bool): Whether to use the start flag. Defaults to True.
            layout (Optional[os.PathLike]): Path to the visualizer layout file. Defaults to None.
            layout_dir (Optional[os.PathLike]): Directory containing visualizer layouts. Defaults to None.
            additional_properties (Optional[Dict[str, str]]): Additional properties for Bonsai. Defaults to None.
            cwd (Optional[os.PathLike]): Working directory for the Bonsai process. Defaults to None.
            timeout (Optional[float]): Timeout for the Bonsai process. Defaults to None.
            print_cmd (bool): Whether to print the command being executed. Defaults to False.
            ui_helper (Optional[UiHelper]): UI helper instance. Defaults to DefaultUIHelper.
        """
        self.executable = Path(executable).resolve()
        self.workflow = Path(workflow).resolve()
        self.is_editor_mode = is_editor_mode
        self.is_start_flag = is_start_flag if not self.is_editor_mode else True
        self.layout = layout
        self.layout_directory = layout_dir
        self.additional_properties = dict(additional_properties) if additional_properties else {}
        self.cwd = cwd
        self.timeout = timeout
        self.print_cmd = print_cmd
        self._result: Optional[subprocess.CompletedProcess] = None
        self.ui_helper = ui_helper if ui_helper is not None else DefaultUIHelper()

    @property
    def result(self) -> subprocess.CompletedProcess:
        """
        Returns the result of the Bonsai process execution.

        Returns:
            subprocess.CompletedProcess: The result of the Bonsai process.

        Raises:
            RuntimeError: If the app has not been run yet.
        """
        if self._result is None:
            raise RuntimeError("The app has not been run yet.")
        return self._result

    @override
    def add_app_settings(self, **kwargs):
        """
        Adds application-specific settings to the additional properties.

        Args:
            **kwargs: Additional keyword arguments.

        Returns:
            Self: The updated instance of BonsaiApp.
        """

        if self.additional_properties is not None:
            self.additional_properties.update(**kwargs)
        else:
            self.additional_properties = kwargs
        return self

    def validate(self, *args, **kwargs) -> bool:
        """
        Validates the existence of required files and directories.

        Args:
            *args: Additional positional arguments.
            **kwargs: Additional keyword arguments.

        Returns:
            bool: True if validation is successful.

        Raises:
            FileNotFoundError: If any required file or directory is missing.
        """
        if not Path(self.executable).exists():
            raise FileNotFoundError(f"Executable not found: {self.executable}")
        if not Path(self.workflow).exists():
            raise FileNotFoundError(f"Workflow file not found: {self.workflow}")
        if self.layout and not Path(self.layout).exists():
            raise FileNotFoundError(f"Layout file not found: {self.layout}")
        if self.layout_directory and not Path(self.layout_directory).exists():
            raise FileNotFoundError(f"Layout directory not found: {self.layout_directory}")
        return True

    @override
    def run(self) -> subprocess.CompletedProcess:
        """
        Runs the Bonsai process.

        Returns:
            subprocess.CompletedProcess: The result of the Bonsai process execution.

        Raises:
            FileNotFoundError: If validation fails.
        """
        self.validate()
        self.prompt_input()

        if self.is_editor_mode:
            logger.warning("Bonsai is running in editor mode. Cannot assert successful completion.")
        logger.info("Bonsai process running...")
        proc = run_bonsai_process(
            workflow_file=self.workflow,
            bonsai_exe=self.executable,
            is_editor_mode=self.is_editor_mode,
            is_start_flag=self.is_start_flag,
            layout=self.layout,
            additional_properties=self.additional_properties,
            cwd=self.cwd,
            timeout=self.timeout,
            print_cmd=self.print_cmd,
        )
        self._result = proc
        logger.info("Bonsai process completed.")
        return proc

    @override
    def output_from_result(self, allow_stderr: Optional[bool]) -> Self:
        """
        Processes the output from the Bonsai process result.

        Args:
            allow_stderr (Optional[bool]): Whether to allow stderr output.

        Returns:
            Self: The updated instance of BonsaiApp.

        Raises:
            subprocess.CalledProcessError: If the process exits with an error.
        """
        proc = self.result
        try:
            proc.check_returncode()
        except subprocess.CalledProcessError as e:
            self._log_process_std_output("Bonsai", proc)
            raise e
        else:
            logger.info("Result from bonsai process is valid.")
            self._log_process_std_output("Bonsai", proc)

            if len(proc.stdout) > 0:
                logger.error("Bonsai process finished with errors.")
                if allow_stderr is None:
                    allow_stderr = self.ui_helper.prompt_yes_no_question("Would you like to see the error message?")
                if allow_stderr is False:
                    raise subprocess.CalledProcessError(1, proc.args)
        return self

    def prompt_visualizer_layout_input(
        self,
        directory: Optional[os.PathLike] = None,
    ) -> Optional[str | os.PathLike]:
        """
        Prompts the user to select a visualizer layout.

        Args:
            directory (Optional[os.PathLike]): Directory containing visualizer layouts.

        Returns:
            Optional[str | os.PathLike]: The selected layout file path.
        """
        if directory is None:
            directory = self.layout_directory
        else:
            directory = Path(os.path.join(directory, VISUALIZERS_DIR, os.environ["COMPUTERNAME"]))

        layout_schemas_path = directory if directory is not None else self.layout_directory
        available_layouts = glob.glob(os.path.join(str(layout_schemas_path), "*.bonsai.layout"))
        picked: Optional[str | os.PathLike] = None
        has_pick = False
        while has_pick is False:
            try:
                picked = self.ui_helper.prompt_pick_from_list(
                    value=available_layouts, prompt="Pick a visualizer layout:"
                )
                picked = picked if picked else ""
                has_pick = True
            except ValueError as e:
                logger.error("Invalid choice. Try again. %s", e)
        self.layout = Path(picked) if picked else None
        return self.layout

    def _log_process_std_output(self, process_name: str, proc: subprocess.CompletedProcess) -> None:
        """
        Logs the standard output and error of a process.

        Args:
            process_name (str): Name of the process.
            proc (subprocess.CompletedProcess): The process result.
        """
        if len(proc.stdout) > 0:
            logger.info("%s full stdout dump: \n%s", process_name, proc.stdout)
        if len(proc.stderr) > 0:
            logger.error("%s full stderr dump: \n%s", process_name, proc.stderr)

    @override
    def prompt_input(self, *args, **kwargs):
        """
        Prompts the user for input if required.

        Args:
            *args: Additional positional arguments.
            **kwargs: Additional keyword arguments.

        Returns:
            Self: The updated instance of BonsaiApp.
        """
        layout_dir = kwargs.pop("layout_directory", None)
        if self.layout is None:
            self.layout = self.prompt_visualizer_layout_input(layout_dir if layout_dir else self.layout_directory)
        return self


class AindBehaviorServicesBonsaiApp(BonsaiApp):
    @overload
    def add_app_settings(self, *, launcher: Optional[BehaviorLauncher] = None, **kwargs) -> Self: ...

    def add_app_settings(self, **kwargs) -> Self:
        launcher: BehaviorLauncher = kwargs.pop("launcher", None)
        if launcher is None:
            raise ValueError("Missing required argument 'launcher'.")

        settings = {
            "TaskLogicPath": os.path.abspath(
                launcher.save_temp_model(model=launcher.task_logic_schema, directory=launcher.temp_dir)
            ),
            "SessionPath": os.path.abspath(
                launcher.save_temp_model(model=launcher.session_schema, directory=launcher.temp_dir)
            ),
            "RigPath": os.path.abspath(
                launcher.save_temp_model(model=launcher.rig_schema, directory=launcher.temp_dir)
            ),
        }
        return super().add_app_settings(**settings)
