import multiprocessing as mp
import shutil
import string
from time import sleep

from tabulate import tabulate
from termcolor import colored

from yera.events.models import TableData
from yera.events.stream import EventStream, OutputEvent, push_text_input


def pretty_print_table(data: TableData):
    print(tabulate(data.rows, headers=data.columns, tablefmt="fancy_grid"))


class PyRuntimeExecutor:
    """Execute an agent in a separate process using an EventStream."""

    def __init__(self, agent, args):
        """Initialise the runtime executor.

        Args:
            agent: The agent to execute.
            args: Positional arguments to pass to the agent's ``invoke`` method.
        """
        self.process: mp.Process | None = None
        self.agent = agent
        self.args = args
        self.stream: EventStream | None = None

    @staticmethod
    def _agent_process(agent, stream: EventStream, *args):
        """Target function for the subprocess.

        Sets the current EventStream, marks the agent as top-level, and invokes it.
        """
        stream.set_current()
        agent.top_level = True
        agent.invoke(*args)

    def start(self) -> None:
        """Start the agent execution in a new process with an internally created EventStream."""
        if self.process is not None and self.process.is_alive():
            # Already running; no-op
            return

        ctx = mp.get_context("spawn")
        self.stream = EventStream.build()
        self.process = ctx.Process(
            target=self._agent_process,
            args=(self.agent, self.stream, *self.args),
        )
        self.process.start()

    def stop(self) -> None:
        """Stop the agent process if it is running and clean up resources."""
        # Capture process locally to avoid race conditions
        process = self.process
        if process is None:
            return

        process.join(timeout=2)
        if process.is_alive():
            process.terminate()
            process.join()

        self.process = None
        self.stream = None

    def is_running(self) -> bool:
        """Check if the executor process is running.

        Returns:
            True if process is alive and stream is available
        """
        return (
            self.process is not None
            and self.process.is_alive()
            and self.stream is not None
        )

    def __enter__(self):
        """Start execution with a fresh EventStream and return self."""
        self.start()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Ensure the subprocess is terminated when leaving the context."""
        self.stop()
        return False


class TokenWrapper:
    def __init__(self, max_width=100):
        terminal_width = shutil.get_terminal_size().columns
        self.width = min(max_width, terminal_width)  # Cap at max_width
        self.current_line_length = 0

    def print_token(self, token):
        # Handle newlines in token
        if "\n" in token:
            print(token, end="", flush=True)
            self.current_line_length = len(token.split("\n")[-1])
            return

        token_length = len(token)

        # Check if token is only punctuation
        is_punctuation = token.strip() and all(
            c in string.punctuation for c in token.strip()
        )

        # Would this token overflow the line?
        longer = self.current_line_length + token_length > self.width
        if longer and not is_punctuation and self.current_line_length > 0:
            # Break to new line
            print("\n", end="", flush=True)
            token = token.lstrip(" ")  # Remove leading spaces
            token_length = len(token)
            self.current_line_length = 0

        print(token, end="", flush=True)
        self.current_line_length += token_length

    def reset(self, new_lines: int = 1):
        """Reset line counter (call when starting new paragraph/section)"""
        self.current_line_length = 0
        print("\n" * (new_lines - 1))


def stream_handler(metadata) -> OutputEvent:
    name = colored(metadata.name, "cyan")
    print(f"Starting {name}:")
    print(metadata.pretty_print())
    stream = EventStream.build()

    exit_event = None

    current_agent_id = None
    current_block_id = None

    wrapper = TokenWrapper(max_width=100)  # Cap at 100 columns

    while exit_event is None:
        for event in stream.iter_output():
            if event.block_type == "exit":
                exit_event = event
                break

            is_new_agent = event.agent_instance.agent_id != current_agent_id
            if is_new_agent:
                wrapper.reset(new_lines=2)
                label = colored(event.agent_instance.agent_id.split(".")[-1], "yellow")
                print(f"[{label}]", end="")

            is_new_block = event.block_id != current_block_id
            if is_new_block:
                wrapper.reset()
                label = colored(event.block_type.upper(), "green")
                print(f"[{label}]")

            # current_block_type = event.block_type
            current_agent_id = event.agent_instance.agent_id
            current_block_id = event.block_id

            if event.block_type == "input_request":
                print()
                res = input(event.data.message)
                push_text_input(res)

            elif event.block_type in ["input_prompt", "system_prompt", "markdown"]:
                wrapper.print_token(event.data.content)
            elif event.block_type == "table":
                pretty_print_table(event.data)
            else:
                print(event.data.model_dump_json(indent=2))

        sleep(0.1)

    exit_data = exit_event.data
    colour = "cyan" if exit_data.exit_code == 0 else "red"
    exit_msg = colored(f"Yera exit code {exit_data.exit_code}".upper(), colour)
    print(f"\n\n[{exit_msg}]")
    print(exit_data.model_dump_json(indent=2))
    list(stream.iter_output())
    return exit_event
