"""Streaming output abstractions for chat responses."""

from typing import Protocol

from .logger import ChatLogger


class Output(Protocol):
    """Protocol for streaming output handlers."""

    logger: ChatLogger | None

    def write(self, chunk: str) -> None: ...
    def end(self) -> None: ...


class ConsoleOutput(Output):
    """Prints streaming chunks to stdout and optionally logs them."""

    def __init__(
        self,
        begin_mark: str = "",
        end_mark: str = "",
        logger: ChatLogger | None = None,
    ) -> None:
        self.started = False
        self.begin_mark = begin_mark
        self.end_mark = end_mark
        self.logger = logger

    def write(self, chunk: str) -> None:
        if not self.started:
            print(self.begin_mark, end="", flush=True)
            self.started = True
        print(chunk, end="", flush=True)
        if self.logger:
            self.logger.add(chunk)

    def end(self) -> None:
        if self.started:
            print(self.end_mark, end="", flush=True)
            self.started = False
            if self.logger:
                self.logger.save()
                if self.logger.last_saved_file:
                    print(f"💾 Conversation saved to {self.logger.last_saved_file}")
                self.logger = None
