"""Forkable conversation thread with checkpoint/rewind support."""

from __future__ import annotations

import copy
from typing import TYPE_CHECKING, Any, Literal

from koan.memory.types import ThreadMessage

if TYPE_CHECKING:
    from collections.abc import Callable


class ForkableThread:
    """In-memory conversation thread with fork and rewind capabilities.

    Maintains an ordered list of messages with named checkpoints.
    Supports forking (branching) from any checkpoint and rewinding.

    An optional ``on_fork`` callback is invoked whenever ``fork_at()``,
    ``fork()``, or ``rewind_to()`` is called.  This lets external systems
    (e.g. the decision graph) record the event as an implicit signal.

    The callback signature is ``(thread_id, event_type, checkpoint_name)``.

    Usage::

        thread = ForkableThread(thread_id="t1")
        thread.append(role="user", content="Hello")
        thread.checkpoint("start")
        thread.append(role="assistant", content="Hi there!")

        # Fork from checkpoint
        branch = thread.fork_at("start")
        branch.append(role="assistant", content="Different response")

        # Rewind
        thread.rewind_to("start")
        assert len(thread.messages) == 1
    """

    def __init__(
        self,
        thread_id: str = "default",
        *,
        on_fork: Callable[[str, str, str], None] | None = None,
    ) -> None:
        self.thread_id = thread_id
        self._messages: list[ThreadMessage] = []
        self._checkpoints: dict[str, int] = {}
        self._seq: int = 0
        self._on_fork = on_fork

    @property
    def messages(self) -> list[ThreadMessage]:
        """Return a copy of the message list."""
        return list(self._messages)

    def append(
        self,
        *,
        role: Literal["system", "user", "assistant", "tool"],
        content: str,
        name: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> ThreadMessage:
        """Append a message to the thread."""
        self._seq += 1
        msg = ThreadMessage(
            id=f"{self.thread_id}-{self._seq}",
            thread_id=self.thread_id,
            role=role,
            content=content,
            name=name,
            metadata=metadata or {},
        )
        self._messages.append(msg)
        return msg

    def checkpoint(self, name: str) -> None:
        """Save the current position as a named checkpoint."""
        self._checkpoints[name] = len(self._messages)

    def rewind_to(self, checkpoint_name: str) -> None:
        """Rewind the thread to a named checkpoint, removing later messages."""
        if checkpoint_name not in self._checkpoints:
            raise KeyError(f"Checkpoint not found: '{checkpoint_name}'")
        pos = self._checkpoints[checkpoint_name]
        self._messages = self._messages[:pos]
        # Remove checkpoints that are after this position
        self._checkpoints = {k: v for k, v in self._checkpoints.items() if v <= pos}
        if self._on_fork:
            self._on_fork(self.thread_id, "rewind", checkpoint_name)

    def fork_at(self, checkpoint_name: str) -> ForkableThread:
        """Create an independent branch from a checkpoint.

        The forked thread has its own copy of messages up to the checkpoint.
        Modifications to the fork do not affect the original.
        """
        if checkpoint_name not in self._checkpoints:
            raise KeyError(f"Checkpoint not found: '{checkpoint_name}'")
        pos = self._checkpoints[checkpoint_name]

        fork = ForkableThread(thread_id=f"{self.thread_id}-fork-{checkpoint_name}")
        fork._messages = copy.deepcopy(self._messages[:pos])
        fork._seq = self._seq
        fork._checkpoints = {k: v for k, v in self._checkpoints.items() if v <= pos}
        if self._on_fork:
            self._on_fork(self.thread_id, "fork", checkpoint_name)
        return fork

    def fork(self) -> ForkableThread:
        """Create an independent branch from the current position."""
        fork = ForkableThread(thread_id=f"{self.thread_id}-fork")
        fork._messages = copy.deepcopy(self._messages)
        fork._seq = self._seq
        fork._checkpoints = dict(self._checkpoints)
        if self._on_fork:
            self._on_fork(self.thread_id, "fork", "_current")
        return fork

    def __len__(self) -> int:
        return len(self._messages)
