import asyncio
import re
import shlex
import time
from pathlib import Path, PurePosixPath

from .asciinema_handler import AsciinemaHandler
from harbor.environments.base import BaseEnvironment
from harbor.utils.logger import logger


class TmuxSession:
    _ENTER_KEYS = {"Enter", "C-m", "KPEnter", "C-j", "^M", "^J"}
    _ENDS_WITH_NEWLINE_PATTERN = r"[\r\n]$"
    _NEWLINE_CHARS = "\r\n"
    _TMUX_COMPLETION_COMMAND = "; tmux wait -S done"
    # tmux silently drops commands exceeding its internal buffer (~16 KB since
    # tmux 1.9, see https://github.com/tmux/tmux/issues/254). Keep a
    # conservative margin below that hard ceiling.
    _TMUX_SEND_KEYS_MAX_COMMAND_LENGTH = 16_000
    GET_ASCIINEMA_TIMESTAMP_SCRIPT_CONTAINER_PATH = Path(
        "/tmp/get-asciinema-timestamp.sh"
    )
    _GET_ASCIINEMA_TIMESTAMP_SCRIPT_HOST_PATH = Path(__file__).parent / (
        "get-asciinema-timestamp.sh"
    )

    def __init__(
        self,
        session_name: str,
        environment: BaseEnvironment,
        logging_path: Path | PurePosixPath,
        local_asciinema_recording_path: Path | None,
        remote_asciinema_recording_path: Path | PurePosixPath | None,
        pane_width: int = 160,
        pane_height: int = 40,
        extra_env: dict[str, str] | None = None,
        user: str | int | None = None,
    ):
        try:
            self._pane_width = int(pane_width)
            self._pane_height = int(pane_height)
        except (ValueError, TypeError):
            raise ValueError("pane_width and pane_height must be valid integers.")
        if self._pane_width <= 0 or self._pane_height <= 0:
            raise ValueError("pane_width and pane_height must be positive integers.")
        self._logging_path = logging_path
        self._local_asciinema_recording_path = local_asciinema_recording_path
        self._remote_asciinema_recording_path = remote_asciinema_recording_path
        self._session_name = session_name
        self._logger = logger
        self._previous_buffer: str | None = None
        self._disable_recording = False
        self.environment = environment
        self._markers: list[tuple[float, str]] = []
        self._extra_env: dict[str, str] = extra_env or {}
        self._user = user

    # TODO: Add asciinema logging
    # @property
    # def logging_path(self) -> Path:
    #     return (
    #         Path(DockerComposeManager.CONTAINER_SESSION_LOGS_PATH)
    #         / f"{self._session_name}.log"
    #     )

    # @property
    # def _recording_path(self) -> Path | None:
    #     if self._disable_recording:
    #         return None
    #     return (
    #         Path(DockerComposeManager.CONTAINER_SESSION_LOGS_PATH)
    #         / f"{self._session_name}.cast"
    #     )

    async def _attempt_tmux_installation(self) -> None:
        """
        Install both tmux and asciinema in a single operation for efficiency.
        """
        # Check what's already installed
        tmux_result = await self.environment.exec(command="tmux -V", user="root")
        tmux_installed = tmux_result.return_code == 0

        needs_asciinema = self._remote_asciinema_recording_path is not None
        if needs_asciinema:
            asciinema_result = await self.environment.exec(
                command="asciinema --version", user="root"
            )
            asciinema_installed = asciinema_result.return_code == 0
        else:
            asciinema_installed = True

        if tmux_installed and asciinema_installed:
            self._logger.debug("Both tmux and asciinema are already installed")
            return

        tools_needed = []
        if not tmux_installed:
            tools_needed.append("tmux")
        if needs_asciinema and not asciinema_installed:
            tools_needed.append("asciinema")

        self._logger.debug(f"Installing: {', '.join(tools_needed)}")

        # Detect system and package manager
        system_info = await self._detect_system_info()

        if system_info["package_manager"]:
            install_command = self._get_combined_install_command(
                system_info, tools_needed
            )
            if install_command:
                self._logger.debug(
                    f"Installing tools using {system_info['package_manager']}: {install_command}"
                )
                result = await self.environment.exec(
                    command=install_command, user="root"
                )

                if result.return_code == 0:
                    # Verify installations
                    if not tmux_installed:
                        verify_tmux = await self.environment.exec(
                            command="tmux -V", user="root"
                        )
                        if verify_tmux.return_code != 0:
                            self._logger.warning(
                                "tmux installation verification failed"
                            )
                            await self._build_tmux_from_source()

                    if needs_asciinema and not asciinema_installed:
                        verify_asciinema = await self.environment.exec(
                            command="asciinema --version", user="root"
                        )
                        if verify_asciinema.return_code != 0:
                            self._logger.warning(
                                "asciinema installation verification failed"
                            )
                            await self._install_asciinema_with_pip()

                    return

        # Fallback to individual installations
        if not tmux_installed:
            self._logger.warning("Installing tmux from source...")
            await self._build_tmux_from_source()

        if needs_asciinema and not asciinema_installed:
            self._logger.warning("Installing asciinema via pip...")
            await self._install_asciinema_with_pip()

    async def _detect_system_info(self) -> dict:
        """
        Detect the operating system and available package managers.
        """
        system_info: dict[str, str | None] = {
            "os": None,
            "package_manager": None,
            "update_command": None,
        }

        # Check for OS release files
        os_release_result = await self.environment.exec(
            command="cat /etc/os-release 2>/dev/null || echo 'not found'",
            user="root",
        )

        # Check uname for system type
        uname_result = await self.environment.exec(command="uname -s", user="root")

        # Detect package managers by checking if they exist
        package_managers = [
            "apt-get",
            "dnf",
            "yum",
            "apk",
            "pacman",
            "brew",
            "pkg",
            "zypper",
        ]

        for pm_name in package_managers:
            check_result = await self.environment.exec(
                command=f"which {pm_name} >/dev/null 2>&1", user="root"
            )
            if check_result.return_code == 0:
                system_info["package_manager"] = pm_name
                break

        # Try to determine OS from available info
        if (
            os_release_result.return_code == 0
            and os_release_result.stdout
            and "not found" not in os_release_result.stdout
        ):
            stdout_lower = os_release_result.stdout.lower()
            if "ubuntu" in stdout_lower or "debian" in stdout_lower:
                system_info["os"] = "debian-based"
            elif "fedora" in stdout_lower or "rhel" in stdout_lower:
                system_info["os"] = "rhel-based"
            elif "alpine" in stdout_lower:
                system_info["os"] = "alpine"
            elif "arch" in stdout_lower:
                system_info["os"] = "arch"
        elif uname_result.return_code == 0 and uname_result.stdout:
            stdout_lower = uname_result.stdout.lower()
            if "darwin" in stdout_lower:
                system_info["os"] = "macos"
            elif "freebsd" in stdout_lower:
                system_info["os"] = "freebsd"

        self._logger.debug(f"Detected system: {system_info}")
        return system_info

    def _get_combined_install_command(self, system_info: dict, tools: list[str]) -> str:
        """
        Get the appropriate installation command for multiple tools based on system info.
        """
        package_manager = system_info.get("package_manager")

        if not package_manager or not isinstance(package_manager, str):
            return ""

        # Build the package list
        packages = " ".join(tools)

        # Package manager commands with non-interactive flags
        install_commands = {
            "apt-get": f"DEBIAN_FRONTEND=noninteractive apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y {packages}",
            "dnf": f"dnf install -y {packages}",
            "yum": f"yum install -y {packages}",
            "apk": f"apk add --no-cache {packages}",
            "pacman": f"pacman -S --noconfirm {packages}",
            "brew": f"brew install {packages}",
            "pkg": f"ASSUME_ALWAYS_YES=yes pkg install -y {packages}",
            "zypper": f"zypper install -y -n {packages}",
        }

        return install_commands.get(package_manager, "")

    async def _build_tmux_from_source(self) -> None:
        """
        Build tmux from source as a fallback option.
        """
        try:
            # Install build dependencies based on detected system - with non-interactive flags
            dep_commands = [
                "DEBIAN_FRONTEND=noninteractive apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y build-essential libevent-dev libncurses5-dev curl",
                "yum groupinstall -y 'Development Tools' && yum install -y libevent-devel ncurses-devel curl",
                "dnf groupinstall -y 'Development Tools' && dnf install -y libevent-devel ncurses-devel curl",
                "apk add --no-cache build-base libevent-dev ncurses-dev curl",
            ]

            # Try to install build dependencies
            for cmd in dep_commands:
                result = await self.environment.exec(command=cmd, user="root")
                if result.return_code == 0:
                    break

            # Download and build tmux
            build_cmd = (
                "cd /tmp && "
                "curl -L https://github.com/tmux/tmux/releases/download/3.4/tmux-3.4.tar.gz -o tmux.tar.gz && "
                "tar -xzf tmux.tar.gz && "
                "cd tmux-3.4 && "
                "./configure --prefix=/usr/local && "
                "make && "
                "make install"
            )

            result = await self.environment.exec(command=build_cmd, user="root")

            # Verify installation
            verify_result = await self.environment.exec(
                command="tmux -V || /usr/local/bin/tmux -V", user="root"
            )
            if verify_result.return_code == 0:
                self._logger.debug("tmux successfully built and installed from source")
            else:
                self._logger.error("Failed to install tmux from source")

        except Exception as e:
            self._logger.error(f"Failed to build tmux from source: {e}")

    async def _install_asciinema_with_pip(self) -> None:
        """
        Install asciinema using pip as a fallback.
        """
        try:
            # Try to install python3-pip first - with non-interactive flags
            pip_install_commands = [
                "DEBIAN_FRONTEND=noninteractive apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y python3-pip",
                "yum install -y python3-pip",
                "dnf install -y python3-pip",
                "apk add --no-cache python3 py3-pip",
            ]

            # Try to install pip
            for cmd in pip_install_commands:
                result = await self.environment.exec(command=cmd, user="root")
                if result.return_code == 0:
                    break

            # Install asciinema using pip
            pip_commands = ["pip3 install asciinema", "pip install asciinema"]

            for cmd in pip_commands:
                result = await self.environment.exec(command=cmd, user="root")
                if result.return_code == 0:
                    # Verify installation
                    verify_result = await self.environment.exec(
                        command="asciinema --version", user="root"
                    )
                    if verify_result.return_code == 0:
                        self._logger.debug("asciinema successfully installed using pip")
                        return

            self._logger.error("Failed to install asciinema using pip")

        except Exception as e:
            self._logger.error(f"Failed to install asciinema with pip: {e}")

    @property
    def _tmux_start_session(self) -> str:
        # Build environment variable options for tmux new-session -e KEY=value
        env_options = "".join(
            f"-e {shlex.quote(f'{key}={value}')} "
            for key, value in self._extra_env.items()
        )

        return (
            f"export TERM=xterm-256color && "
            f"export SHELL=/bin/bash && "
            f"tmux new-session {env_options}-x {self._pane_width} -y {self._pane_height} -d -s {self._session_name} 'bash --login' \\; "
            f"pipe-pane -t {self._session_name} "
            f"'cat > {self._logging_path}'"
        )

    def _tmux_send_keys(self, keys: list[str]) -> list[str]:
        """Build one or more ``tmux send-keys`` commands for *keys*.

        If the shell-escaped command would exceed the tmux command-length
        limit, the keys are spread across multiple commands so that each
        individual command stays within the limit.  Oversized single keys
        are split into sub-strings whose quoted form fits.
        """
        prefix = "tmux send-keys -t " + shlex.quote(self._session_name)
        max_len = self._TMUX_SEND_KEYS_MAX_COMMAND_LENGTH

        escaped_keys = [shlex.quote(key) for key in keys]
        single = prefix + " " + " ".join(escaped_keys)
        if len(single) <= max_len:
            return [single]

        commands: list[str] = []
        current_escaped: list[str] = []
        current_len = len(prefix)

        def _flush() -> None:
            nonlocal current_len
            if current_escaped:
                commands.append(prefix + " " + " ".join(current_escaped))
                current_escaped.clear()
                current_len = len(prefix)

        for key in keys:
            escaped = shlex.quote(key)
            addition = 1 + len(escaped)  # space + quoted key

            if current_len + addition <= max_len:
                current_escaped.append(escaped)
                current_len += addition
            elif len(prefix) + addition <= max_len:
                _flush()
                current_escaped.append(escaped)
                current_len = len(prefix) + addition
            else:
                _flush()
                max_escaped = max_len - len(prefix) - 1
                for chunk_escaped in self._split_key_for_tmux(key, max_escaped):
                    if current_len + 1 + len(chunk_escaped) <= max_len:
                        current_escaped.append(chunk_escaped)
                        current_len += 1 + len(chunk_escaped)
                    else:
                        _flush()
                        current_escaped.append(chunk_escaped)
                        current_len = len(prefix) + 1 + len(chunk_escaped)

        _flush()
        return commands

    @staticmethod
    def _split_key_for_tmux(key: str, max_escaped_len: int) -> list[str]:
        """Split *key* into ``shlex.quote``-d chunks each ≤ *max_escaped_len*."""
        chunks: list[str] = []
        remaining = key
        while remaining:
            lo, hi, best = 1, len(remaining), 1
            while lo <= hi:
                mid = (lo + hi) // 2
                if len(shlex.quote(remaining[:mid])) <= max_escaped_len:
                    best = mid
                    lo = mid + 1
                else:
                    hi = mid - 1
            chunks.append(shlex.quote(remaining[:best]))
            remaining = remaining[best:]
        return chunks

    def _tmux_capture_pane(self, capture_entire: bool = False) -> str:
        if capture_entire:
            extra_args = ["-S", "-"]
        else:
            extra_args = []

        return " ".join(
            [
                "tmux",
                "capture-pane",
                "-p",
                *extra_args,
                "-t",
                self._session_name,
            ]
        )

    async def start(self) -> None:
        await self._attempt_tmux_installation()
        start_session_result = await self.environment.exec(
            command=self._tmux_start_session, user=self._user
        )
        if start_session_result.return_code != 0:
            raise RuntimeError(
                f"Failed to start tmux session. Error: {start_session_result.stderr}"
            )

        history_limit = 10_000_000
        command = f"tmux set-option -g history-limit {history_limit}"
        set_history_result = await self.environment.exec(
            command=command, user=self._user
        )
        if set_history_result.return_code != 0:
            self._logger.warning(
                "Failed to increase tmux history-limit: %s",
                (set_history_result.stderr or "").strip(),
            )

        if self._remote_asciinema_recording_path:
            self._logger.debug("Starting recording.")
            await self.send_keys(
                keys=[
                    f"asciinema rec --stdin {self._remote_asciinema_recording_path}",
                    "Enter",
                ],
                min_timeout_sec=1.0,
            )
            await self.send_keys(
                keys=[
                    "clear",
                    "Enter",
                ],
            )

        if self._remote_asciinema_recording_path:
            await self.environment.upload_file(
                source_path=self._GET_ASCIINEMA_TIMESTAMP_SCRIPT_HOST_PATH,
                target_path=str(self.GET_ASCIINEMA_TIMESTAMP_SCRIPT_CONTAINER_PATH),
            )

    async def stop(self) -> None:
        if self._remote_asciinema_recording_path:
            self._logger.debug("Stopping recording.")
            await self.send_keys(
                keys=["C-d"],
                min_timeout_sec=0.1,
            )

            # Wait a moment for the recording to finish writing
            import asyncio

            await asyncio.sleep(0.5)

            if self._local_asciinema_recording_path:
                self._local_asciinema_recording_path.parent.mkdir(
                    parents=True, exist_ok=True
                )
                # Ensure recording exists locally before merging markers
                await self.environment.download_file(
                    source_path=str(self._remote_asciinema_recording_path),
                    target_path=self._local_asciinema_recording_path,
                )

            # Merge markers into the recording
            if self._markers and self._local_asciinema_recording_path:
                self._logger.debug(
                    f"Merging {len(self._markers)} markers into recording"
                )
                handler = AsciinemaHandler(
                    self._markers, self._local_asciinema_recording_path
                )
                handler.merge_markers()
                self._logger.debug(
                    f"Successfully merged markers into {self._local_asciinema_recording_path}"
                )
            else:
                self._logger.debug("No markers to merge")
        ...

    def _is_enter_key(self, key: str) -> bool:
        return key in self._ENTER_KEYS

    def _ends_with_newline(self, key: str) -> bool:
        result = re.search(self._ENDS_WITH_NEWLINE_PATTERN, key)
        return result is not None

    async def is_session_alive(self) -> bool:
        """Check if the tmux session is still alive."""
        result = await self.environment.exec(
            command="tmux has-session -t {}".format(self._session_name),
            user=self._user,
        )
        return result.return_code == 0

    def _is_executing_command(self, key: str) -> bool:
        return self._is_enter_key(key) or self._ends_with_newline(key)

    def _prevent_execution(self, keys: list[str]) -> list[str]:
        keys = keys.copy()
        while keys and self._is_executing_command(keys[-1]):
            if self._is_enter_key(keys[-1]):
                keys.pop()
            else:
                stripped_key = keys[-1].rstrip(self._NEWLINE_CHARS)

                if stripped_key:
                    keys[-1] = stripped_key
                else:
                    keys.pop()

        return keys

    def _prepare_keys(
        self,
        keys: str | list[str],
        block: bool,
    ) -> tuple[list[str], bool]:
        """
        Prepare keys for sending to the terminal.

        Args:
            keys (str | list[str]): The keys to send to the terminal.
            block (bool): Whether to wait for the command to complete.

        Returns:
            tuple[list[str], bool]: The keys to send to the terminal and whether the
                the keys are blocking.
        """
        if isinstance(keys, str):
            keys = [keys]

        if not block or not keys or not self._is_executing_command(keys[-1]):
            return keys, False

        keys = self._prevent_execution(keys)
        keys.extend([self._TMUX_COMPLETION_COMMAND, "Enter"])

        return keys, True

    async def _send_blocking_keys(
        self,
        keys: list[str],
        max_timeout_sec: float,
    ):
        start_time_sec = time.time()

        for command in self._tmux_send_keys(keys):
            result = await self.environment.exec(command=command, user=self._user)
            if result.return_code != 0:
                raise RuntimeError(
                    f"{self.environment.session_id}: failed to send blocking keys: {result.stderr}"
                )

        result = await self.environment.exec(
            f"timeout {max_timeout_sec}s tmux wait done", user=self._user
        )
        if result.return_code != 0:
            raise TimeoutError(f"Command timed out after {max_timeout_sec} seconds")

        elapsed_time_sec = time.time() - start_time_sec
        self._logger.debug(f"Blocking command completed in {elapsed_time_sec:.2f}s.")

    async def _send_non_blocking_keys(
        self,
        keys: list[str],
        min_timeout_sec: float,
    ):
        start_time_sec = time.time()

        for command in self._tmux_send_keys(keys):
            result = await self.environment.exec(command=command, user=self._user)
            if result.return_code != 0:
                raise RuntimeError(
                    f"{self.environment.session_id}: failed to send non-blocking keys: {result.stderr}"
                )

        elapsed_time_sec = time.time() - start_time_sec

        if elapsed_time_sec < min_timeout_sec:
            await asyncio.sleep(min_timeout_sec - elapsed_time_sec)

    async def send_keys(
        self,
        keys: str | list[str],
        block: bool = False,
        min_timeout_sec: float = 0.0,
        max_timeout_sec: float = 180.0,
    ):
        """
        Execute a command in the tmux session.

        Args:
            keys (str): The keys to send to the tmux session.
            block (bool): Whether to wait for the command to complete.
            min_timeout_sec (float): Minimum time in seconds to wait after executing.
                Defaults to 0.
            max_timeout_sec (float): Maximum time in seconds to wait for blocking
                commands. Defaults to 3 minutes.
        """
        if block and min_timeout_sec > 0.0:
            self._logger.debug("min_timeout_sec will be ignored because block is True.")

        prepared_keys, is_blocking = self._prepare_keys(
            keys=keys,
            block=block,
        )

        self._logger.debug(
            f"Sending keys: {prepared_keys}"
            f" min_timeout_sec: {min_timeout_sec}"
            f" max_timeout_sec: {max_timeout_sec}"
        )

        if is_blocking:
            await self._send_blocking_keys(
                keys=prepared_keys,
                max_timeout_sec=max_timeout_sec,
            )
        else:
            await self._send_non_blocking_keys(
                keys=prepared_keys,
                min_timeout_sec=min_timeout_sec,
            )

    async def capture_pane(self, capture_entire: bool = False) -> str:
        result = await self.environment.exec(
            self._tmux_capture_pane(capture_entire=capture_entire), user=self._user
        )
        return result.stdout or ""

    async def _get_visible_screen(self) -> str:
        return await self.capture_pane(capture_entire=False)

    async def _find_new_content(self, current_buffer: str) -> str | None:
        pb = "" if self._previous_buffer is None else self._previous_buffer.strip()
        if pb in current_buffer:
            idx = current_buffer.index(pb)
            # Find the end of the previous buffer content
            if "\n" in pb:
                idx = pb.rfind("\n")
            return current_buffer[idx:]
        return None

    async def get_incremental_output(self) -> str:
        """
        Get either new terminal output since last call, or current screen if
        unable to determine.

        This method tracks the previous buffer state and attempts to find new content
        by comparing against the current full buffer. This provides better handling for
        commands with large output that would overflow the visible screen.

        Returns:
            str: Formatted output with either "New Terminal Output:" or
                 "Current Terminal Screen:"
        """
        current_buffer = await self.capture_pane(capture_entire=True)

        # First capture - no previous state
        if self._previous_buffer is None:
            self._previous_buffer = current_buffer
            visible_screen = await self._get_visible_screen()
            return f"Current Terminal Screen:\n{visible_screen}"

        # Try to find new content
        new_content = await self._find_new_content(current_buffer)

        # Update state
        self._previous_buffer = current_buffer

        if new_content is not None:
            if new_content.strip():
                return f"New Terminal Output:\n{new_content}"
            else:
                # No new content, show current screen
                return f"Current Terminal Screen:\n{await self._get_visible_screen()}"
        else:
            # Couldn't reliably determine new content, fall back to current screen
            return f"Current Terminal Screen:\n{await self._get_visible_screen()}"
