"""
Copyright (C) 2024-2026 Mo Zhou <lumin@debian.org>

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Lesser General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.
"""
import argparse
import contextlib
import functools as ft
import json
import os
import select
import shlex
import sys
import threading
import time
import uuid
from typing import Dict, List, Optional, Union

from prompt_toolkit import PromptSession
from prompt_toolkit.completion import Completer, Completion
from prompt_toolkit.styles import Style
from rich.console import Console, Group
from rich.live import Live
from rich.markdown import Markdown
from rich.markup import escape
from rich.padding import Padding
from rich.style import Style as richStyle
from rich.text import Text

from . import defaults

console = defaults.console
console_stdout = Console()


class StreamInterrupted(Exception):
    """Raised when the user presses ESC during streaming.

    The partial response collected so far is saved into the session
    so the user can correct or continue the conversation.
    """

    pass


@contextlib.contextmanager
def _esc_watcher():
    """Context manager that watches for ESC keypress in a background thread.

    Yields a :class:`threading.Event` that is set when ESC is pressed.
    Falls back to a never-set event if stdin is not a tty or termios is
    unavailable (e.g. on Windows or when stdin is a pipe).
    """
    event = threading.Event()
    stop_ev = threading.Event()
    old_settings = None
    watcher = None
    fd = None

    if sys.stdin.isatty():
        try:
            import termios

            fd = sys.stdin.fileno()
            old_settings = termios.tcgetattr(fd)
            # Disable canonical mode and echo so single keypresses are
            # available immediately.  Keep ISIG so Ctrl-C still delivers
            # SIGINT as a fallback.
            new = termios.tcgetattr(fd)
            new[3] = new[3] & ~(termios.ECHO | termios.ICANON)
            new[6][termios.VMIN] = 0
            new[6][termios.VTIME] = 0
            termios.tcsetattr(fd, termios.TCSADRAIN, new)

            def _watch():
                try:
                    while not stop_ev.is_set():
                        r, _, _ = select.select([fd], [], [], 0.1)
                        if r:
                            ch = os.read(fd, 1)
                            if ch == b"\x1b":  # ESC
                                # Drain remaining escape-sequence bytes
                                # (e.g. arrow keys send \x1b[A)
                                while select.select([fd], [], [], 0.05)[0]:
                                    os.read(fd, 1024)
                                event.set()
                                return
                except Exception:
                    pass

            watcher = threading.Thread(target=_watch, daemon=True)
            watcher.start()
        except Exception:
            pass

    try:
        yield event
    finally:
        stop_ev.set()
        if old_settings is not None and fd is not None:
            # Drain any leftover bytes before restoring settings
            try:
                while select.select([fd], [], [], 0)[0]:
                    os.read(fd, 1024)
            except Exception:
                pass
            try:
                import termios

                termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
            except Exception:
                pass
        if watcher is not None:
            watcher.join(timeout=0.5)


def _tail_window(think_text, chunk_text, term_height):
    """Build a renderable showing the tail of streaming content.

    Keeps output within one terminal page so vertical_overflow='crop'
    always shows the latest tokens instead of freezing at the first page.
    """
    max_lines = max(term_height - 4, 10)
    parts = []
    if think_text and not chunk_text:
        lines = think_text.split("\n")
        if len(lines) > max_lines:
            think_text = "\n".join(lines[-max_lines:])
        part = Text(think_text)
        parts.append(Padding(part, (0, 2), style=richStyle(dim=True, italic=True)))
    elif think_text and chunk_text:
        indicator = Text(f"... ({len(think_text.split(chr(10)))} lines of thinking)")
        parts.append(Padding(indicator, (0, 2), style=richStyle(dim=True, italic=True)))
        remaining = max_lines - 2
        lines = chunk_text.split("\n")
        if len(lines) > remaining:
            chunk_text = "\n".join(lines[-remaining:])
        parts.append(Markdown(chunk_text))
    elif chunk_text:
        lines = chunk_text.split("\n")
        if len(lines) > max_lines:
            chunk_text = "\n".join(lines[-max_lines:])
        parts.append(Markdown(chunk_text))
    else:
        parts.append(Markdown(""))
    return Group(*parts) if len(parts) > 1 else parts[0]


def _check(messages: List[Dict]):
    """
    communitation protocol.
    both huggingface transformers and openapi api use this
    """
    assert isinstance(messages, list)
    assert all(isinstance(x, dict) for x in messages)
    assert all("role" in x.keys() for x in messages)
    assert all("content" in x.keys() for x in messages)
    assert all(isinstance(x["role"], str) for x in messages)
    assert all(isinstance(x["content"], str) for x in messages)
    assert all(x["role"] in ("system", "user", "assistant", "tool") for x in messages)


def retry_ratelimit(func: callable, exception: Exception, retry_interval: int = 15):
    """
    A decorator to retry the function call when the specified exception occurs.

    OpenAI API doc provides some other methods to retry:
    https://platform.openai.com/docs/guides/rate-limits/error-mitigation
    """

    @ft.wraps(func)
    def wrapper(*args, **kwargs):
        while True:
            try:
                result = func(*args, **kwargs)
                break
            except exception:
                console.log(
                    f"Rate limit reached. Will retry after {retry_interval} seconds."
                )
                time.sleep(retry_interval)
        return result

    return wrapper


def _accumulate_tool_calls(delta, tool_calls):
    if hasattr(delta, "tool_calls") and delta.tool_calls:
        for tc in delta.tool_calls:
            if tc.index not in tool_calls:
                tool_calls[tc.index] = {
                    "id": tc.id or "",
                    "type": "function",
                    "function": {
                        "name": tc.function.name or "",
                        "arguments": tc.function.arguments or "",
                    },
                }
            else:
                if getattr(tc, "id", None):
                    tool_calls[tc.index]["id"] += tc.id
                if getattr(tc.function, "name", None):
                    tool_calls[tc.index]["function"]["name"] += tc.function.name
                if getattr(tc.function, "arguments", None):
                    tool_calls[tc.index]["function"][
                        "arguments"
                    ] += tc.function.arguments


class AbstractFrontend:
    """
    The frontend instance holds the whole chat session. The context is the whole
    session for the next LLM query. Historical chats is also a part of the
    context for following up questions. You may feel LLMs smart when they
    get information from the historical chat in the same session.
    """

    NAME = "AbstractFrontend"

    def __init__(self, args):
        self.uuid = uuid.uuid4()
        self.session = []
        self.debgpt_home = args.debgpt_home
        if getattr(args, "color", "auto") == "auto":
            self.monochrome = not sys.stdout.isatty()
        else:
            self.monochrome = getattr(args, "color", "auto") == "no"

        if self.monochrome:
            console_stdout.no_color = True
            console.no_color = True

        self.multiline = getattr(args, "multiline", False)
        if getattr(args, "render", "auto") == "auto":
            self.render_markdown = sys.stdout.isatty()
        else:
            self.render_markdown = getattr(args, "render", "auto") == "yes"
        self.vertical_overflow = args.vertical_overflow
        self.verbose = args.verbose
        self.approve_all = getattr(args, "approve", False)
        self.no_approve = getattr(args, "no_approve", False)
        self.default_context_window = getattr(args, "context_window_size", 1000000)
        self.context_window = None
        console.log(f"{self.NAME}> Starting conversation {self.uuid}")

    def reset(self):
        """
        clear the context. No need to change UUID I think.
        """
        self.session = []

    def oneshot(self, message: str) -> str:
        """
        Generate response text from the given question, without history.
        And do not print anything. Just return the response text silently.

        Args:
            message: a string, the question.
        Returns:
            a string, the response text.
        """
        raise NotImplementedError("please override AbstractFrontend.oneshot()")

    def query(self, messages: List[Dict]) -> str:
        """
        Generate response text from the given chat history. This function
        will also handle printing and rendering.

        Args:
            messages: a list of dict, each dict contains a message.
        Returns:
            a string, the response text.
        the messages format can be found in _check(...) function above.
        """
        raise NotImplementedError("please override AbstractFrontend.query()")

    def update_session(self, messages: Union[List, Dict, str]) -> None:
        if isinstance(messages, list):
            # reset the chat with provided message list
            self.session = messages
        elif isinstance(messages, dict):
            # just append a new dict
            self.session.append(messages)
        elif isinstance(messages, str):
            # just append a new user dict
            self.session.append({"role": "user", "content": messages})
        else:
            raise TypeError(type(messages))
        _check(self.session)

    def __call__(self, *args, **kwargs):
        try:
            res = self.query(*args, **kwargs)
            return res
        except StreamInterrupted:
            raise  # not an error; partial response already saved
        except Exception as e:
            # this will only appear in dumped session files
            self.update_session({"role": "system", "content": str(e)})
            raise e

    def dump(self):
        fpath = os.path.join(self.debgpt_home, str(self.uuid) + ".json")
        with open(fpath, "wt") as f:
            json.dump(self.session, f, indent=2)
        # Write the LAST pointer so that `debgpt resume` can reload this session.
        last_path = os.path.join(self.debgpt_home, "LAST")
        with open(last_path, "wt") as f:
            f.write(os.path.basename(fpath) + "\n")
        console.log(f"{self.NAME}> Conversation saved at {fpath}")

    def _report_context_remaining(self, total_tokens: int) -> None:
        """Report context remaining percentage if context_window is known."""
        if self.context_window and total_tokens > 0:
            remaining = self.context_window - total_tokens
            pct = remaining / self.context_window * 100
            if self.verbose:
                console.log(
                    f"{self.NAME}> context window {self.context_window}, "
                    f"remaining {remaining} ({pct:.1f}%)."
                )

    def _maybe_summarize_history(self, total_tokens: int) -> None:
        """Summarize old messages if context window is getting full (>= 90%)."""
        if not self.context_window or total_tokens < self.context_window * 0.9:
            return

        # Need enough messages to make summarization worthwhile
        if len([m for m in self.session if m["role"] != "system"]) <= 4:
            return

        console.log(
            f"{self.NAME}> Context at {(total_tokens/self.context_window)*100:.1f}% capacity. Summarizing old messages to free up space..."
        )

        system_msgs = [m for m in self.session if m["role"] == "system"]
        # Keep the most recent 2 messages (usually user + assistant)
        recent_msgs = self.session[-2:]
        old_msgs = [
            m for m in self.session if m not in system_msgs and m not in recent_msgs
        ]

        if not old_msgs:
            return

        summary_prompt = "Summarize the following conversation history briefly but retain key facts, user preferences, tool results, and important context:\n\n"
        for m in old_msgs:
            # truncate very long content in old messages to avoid blowing up the oneshot prompt
            content = str(m.get("content", ""))
            if len(content) > 2000:
                content = content[:2000] + "... [truncated]"
            summary_prompt += f"Role: {m['role'].upper()}\nContent: {content}\n\n"

        try:
            summary = self.oneshot(summary_prompt)
            new_session = (
                system_msgs
                + [
                    {
                        "role": "system",
                        "content": f"Summary of previous conversation:\n{summary}",
                    }
                ]
                + recent_msgs
            )
            self.session = new_session
            console.log(f"{self.NAME}> Summarization complete. Session truncated.")
        except Exception as e:
            console.log(f"{self.NAME}> Failed to summarize history: {e}")

    def __len__(self):
        """
        Calculate the number of messages from user and assistant in the session,
        excluding system message.
        """
        return len([x for x in self.session if x["role"] != "system"])


class EchoFrontend(AbstractFrontend):
    """
    A frontend that echoes the input text. Don't worry, this is just for
    running unit tests.
    """

    NAME = "EchoFrontend"
    lossy_mode: bool = False
    lossy_rate: int = 2

    def __init__(self, args: Optional[object] = None):
        # do not call super().__init__(args) here.
        self.session = []
        self.monochrome = False
        self.multiline = False
        self.render_markdown = False

    def oneshot(self, message: str) -> str:
        if self.lossy_mode:
            return message[:: self.lossy_rate]
        else:
            return message

    def query(self, messages: Union[List, Dict, str]) -> list:
        self.update_session(messages)
        new_input = self.session[-1]["content"]
        if self.lossy_mode:
            response = new_input[:: self.lossy_rate]
        else:
            response = new_input
        console_stdout.print(response)
        new_message = {"role": "assistant", "content": response}
        self.update_session(new_message)
        return self.session[-1]["content"]

    def dump(self):
        pass


class OpenAIFrontend(AbstractFrontend):
    """
    https://platform.openai.com/docs/quickstart?context=python
    """

    NAME: str = "OpenAIFrontend"
    debug: bool = False

    def __init__(self, args):
        super().__init__(args)
        try:
            from openai import OpenAI
        except ImportError:
            console.log('please install OpenAI package: "pip install openai"')
            exit(1)
        self.client = OpenAI(api_key=args.openai_api_key, base_url=args.openai_base_url)
        self.model = args.openai_model
        # XXX: some models do not support system messages yet. nor temperature.
        if self.model not in ("o1-mini", "o1-preview", "o3-mini"):
            self.session.append({"role": "system", "content": args.system_message})
            self.kwargs = {"temperature": args.temperature, "top_p": args.top_p}
        else:
            self.kwargs = {}
        self._query_context_window()
        if args.verbose:
            console.log(
                f"{self.NAME}> model={repr(self.model)}, "
                + f"temperature={args.temperature}, top_p={args.top_p}."
            )

    def _query_context_window(self):
        """Query the model context window from the API."""
        try:
            model_info = self.client.models.retrieve(self.model)
            # OpenAI models expose context_window (or fall back to None)
            self.context_window = getattr(model_info, "context_window", None)
        except Exception:
            self.context_window = None

        if self.context_window is None:
            self.context_window = self.default_context_window

    def oneshot(self, message: str) -> str:
        def _func() -> str:
            _callable = self.client.chat.completions.create
            completions = _callable(
                model=self.model,
                messages=[{"role": "user", "content": message}],
                **self.kwargs,
            )
            return completions.choices[0].message.content

        from openai import RateLimitError

        return retry_ratelimit(_func, RateLimitError)()

    def query(self, messages: Union[List, Dict, str]) -> list:
        from .tools import get_openai_tools, get_tool_by_name

        # add the message into the session
        self.update_session(messages)
        if self.debug:
            console.log("send:", self.session[-1])

        while True:
            kwargs = self.kwargs.copy()
            tools = get_openai_tools()
            if tools:
                kwargs["tools"] = tools

            kwargs["stream_options"] = {"include_usage": True}

            completion = self.client.chat.completions.create(
                model=self.model, messages=self.session, stream=True, **kwargs
            )

            tool_calls = {}  # to accumulate streaming tool calls
            generated_text = ""

            n_tokens: int = 0
            time_start_end: List[float] = [0.0, 0.0]
            think, chunks = [], []
            cursor = chunks
            usage_info = None

            interrupted = False
            time_start_end[0] = time.time()
            with _esc_watcher() as esc_pressed:
                if self.render_markdown:
                    with Live(
                        Markdown(""),
                        console=console_stdout,
                        vertical_overflow="crop",
                        transient=True,
                    ) as live:
                        try:
                            for chunk in completion:
                                if esc_pressed.is_set():
                                    interrupted = True
                                    try:
                                        completion.close()
                                    except Exception:
                                        pass
                                    break
                                if hasattr(chunk, "usage") and chunk.usage:
                                    usage_info = chunk.usage
                                if not chunk.choices:
                                    continue

                                delta = chunk.choices[0].delta

                                # Handle tool calls in stream
                                _accumulate_tool_calls(delta, tool_calls)

                                has_update = False
                                if (
                                    hasattr(delta, "reasoning_content")
                                    and delta.reasoning_content
                                ):
                                    think.append(delta.reasoning_content)
                                    has_update = True
                                if getattr(delta, "content", None):
                                    piece = delta.content
                                    n_tokens += 1
                                    if piece == "</think>" and len(think) > 0:
                                        cursor = chunks
                                    elif piece == "<think>":
                                        cursor = think
                                    else:
                                        cursor.append(piece)
                                    has_update = True
                                if not has_update:
                                    continue
                                live.update(
                                    _tail_window(
                                        "".join(think),
                                        "".join(chunks),
                                        console_stdout.size.height,
                                    ),
                                    refresh=True,
                                )
                        except KeyboardInterrupt:
                            interrupted = True
                            try:
                                completion.close()
                            except Exception:
                                pass
                    time_start_end[1] = time.time()
                    if think:
                        part1 = Text("".join(think))
                        part1 = Padding(
                            part1, (0, 2), style=richStyle(dim=True, italic=True)
                        )
                        console_stdout.print(part1)
                    if chunks:
                        console_stdout.print(Markdown("".join(chunks)))
                else:
                    try:
                        for chunk in completion:
                            if esc_pressed.is_set():
                                interrupted = True
                                try:
                                    completion.close()
                                except Exception:
                                    pass
                                break
                            if hasattr(chunk, "usage") and chunk.usage:
                                usage_info = chunk.usage
                            if not chunk.choices:
                                continue

                            delta = chunk.choices[0].delta
                            _accumulate_tool_calls(delta, tool_calls)

                            if (
                                hasattr(delta, "reasoning_content")
                                and delta.reasoning_content
                            ):
                                piece = delta.reasoning_content
                                think.append(piece)
                                print(piece, end="", flush=True)
                            if getattr(delta, "content", None):
                                piece = delta.content
                                n_tokens += 1
                                chunks.append(piece)
                                print(piece, end="", flush=True)
                    except KeyboardInterrupt:
                        interrupted = True
                        try:
                            completion.close()
                        except Exception:
                            pass
                    time_start_end[1] = time.time()

            generated_text = "".join(chunks)
            if not generated_text.endswith("\n") and generated_text:
                print()
                sys.stdout.flush()
            if interrupted:
                console.print(
                    "\n[bold yellow]⚠ Streaming interrupted by user.[/bold yellow]"
                )
                if generated_text:
                    partial = generated_text + "\n\n[interrupted by user]"
                    self.update_session({"role": "assistant", "content": partial})
                raise StreamInterrupted()
            if self.verbose:
                _gtps = n_tokens / max((time_start_end[1] - time_start_end[0]), 0.001)
                if n_tokens > 0:
                    console.log(
                        f"{self.NAME}({self.model})> {_gtps:.2f} generation tokens per second."
                    )
            if usage_info:
                if self.verbose:
                    console.log(
                        f"{self.NAME}> tokens consumed: input {usage_info.prompt_tokens}, "
                        f"output {usage_info.completion_tokens}, total {usage_info.total_tokens}."
                    )
                self._report_context_remaining(usage_info.total_tokens)
                self._maybe_summarize_history(usage_info.total_tokens)

            # convert tool_calls dict to list
            tool_calls_list = [tool_calls[k] for k in sorted(tool_calls.keys())]

            if tool_calls_list:
                # Add assistant message with tool calls
                new_message = {"role": "assistant", "content": generated_text or ""}
                new_message["tool_calls"] = tool_calls_list
                self.update_session(new_message)

                # Execute tools
                for tc in tool_calls_list:
                    func_name = tc["function"]["name"]
                    func_args = tc["function"]["arguments"]
                    tool_instance = get_tool_by_name(func_name)
                    if tool_instance:
                        try:
                            args_dict = json.loads(func_args)
                            if self.verbose:
                                console.log(
                                    f"Calling tool: {func_name} with {func_args}"
                                )
                            else:
                                console.print(
                                    f"[dim]⚙ ToolCall>[/dim] [cyan]{escape(tool_instance.summary(args_dict))}[/cyan]"
                                )
                            # Check if the tool requires user confirmation
                            confirm, reason = tool_instance.needs_confirmation(
                                args_dict
                            )
                            if self.no_approve and not confirm:
                                confirm = True
                                reason = f"Tool '{func_name}' requires approval (--no-approve mode)."
                            if confirm and self.approve_all:
                                console.print(
                                    f"\n[bold yellow]⚠ {reason}[/bold yellow]"
                                )
                                console.print(
                                    "[green]Auto-approved (--approve).[/green]"
                                )
                                confirm = False
                            if confirm:
                                console.print(
                                    f"\n[bold yellow]⚠ {reason}[/bold yellow]"
                                )
                                ctx = tool_instance.confirmation_context(args_dict)
                                if ctx:
                                    # ctx may contain ANSI escape codes from
                                    # Pygments, so print() it raw rather than
                                    # going through Rich markup.
                                    print(ctx)
                                answer = console.input(
                                    "[bold]Accept this operation? (y/n/comment): [/bold]"
                                ).strip()
                                if answer.lower() in ("y", "yes"):
                                    result = tool_instance(**args_dict)
                                    console.print("[green]Accepted.[/green]")
                                elif answer.lower() in ("n", "no", ""):
                                    result = "Operation rejected by the user."
                                    console.print("[red]Rejected.[/red]")
                                else:
                                    result = (
                                        "Operation rejected by the user with"
                                        f" comment: {answer}"
                                    )
                                    console.print(
                                        f"[red]Rejected with comment:[/red] {answer}"
                                    )
                            else:
                                result = tool_instance(**args_dict)
                        except Exception as e:
                            result = str(e)
                    else:
                        result = f"Tool {func_name} not found."

                    self.update_session(
                        {
                            "role": "tool",
                            "tool_call_id": tc["id"],
                            "name": func_name,
                            "content": str(result),
                        }
                    )
                # loop back to send tool results to LLM
                continue
            else:
                new_message = {"role": "assistant", "content": generated_text}
                self.update_session(new_message)
                if self.debug:
                    console.log("recv:", self.session[-1])
                return self.session[-1]["content"]


class AnthropicFrontend(AbstractFrontend):
    """
    https://docs.anthropic.com/en/api/getting-started
    But we are currently using OpenAI API.

    The max_token limit for each model can be found here:
    https://docs.anthropic.com/en/docs/about-claude/models
    """

    NAME = "AnthropicFrontend"
    debug: bool = False
    max_tokens: int = 4096

    def __init__(self, args):
        super().__init__(args)
        try:
            from anthropic import Anthropic
        except ImportError:
            console.log('please install Anthropic package: "pip install anthropic"')
            exit(1)
        self.client = Anthropic(
            api_key=args.anthropic_api_key, base_url=args.anthropic_base_url
        )
        self.model = args.anthropic_model
        self.kwargs = {
            "temperature": args.temperature,
            "top_p": args.top_p,
            "system": args.system_message,
        }
        self._query_context_window()
        if args.verbose:
            console.log(
                f"{self.NAME}> model={repr(self.model)}, "
                + f"temperature={args.temperature}, top_p={args.top_p}."
            )

    def _query_context_window(self):
        """Query the model context window from the Anthropic API."""
        try:
            model_info = self.client.models.retrieve(model_id=self.model)
            self.context_window = getattr(model_info, "max_input_tokens", None)
        except Exception:
            self.context_window = None

        if self.context_window is None:
            self.context_window = self.default_context_window

    def oneshot(self, message: str) -> str:
        def _func():
            _callable = self.client.messages.create
            completion = _callable(
                model=self.model,
                messages=[{"role": "user", "content": message}],
                max_tokens=self.max_tokens,
                **self.kwargs,
            )
            return completion.content[0].text

        from anthropic import RateLimitError

        return retry_ratelimit(_func, RateLimitError)()

    def query(self, messages: Union[List, Dict, str]) -> list:
        # add the message into the session
        self.update_session(messages)
        if self.debug:
            console.log("send:", self.session[-1])
        think, chunks = [], []
        current_block_type = None
        usage_input_tokens = 0
        usage_output_tokens = 0
        interrupted = False
        with self.client.messages.stream(
            model=self.model,
            messages=self.session,
            max_tokens=self.max_tokens,
            **self.kwargs,
        ) as stream:
            with _esc_watcher() as esc_pressed:
                if self.render_markdown:
                    with Live(
                        Markdown(""),
                        console=console_stdout,
                        vertical_overflow="crop",
                        transient=True,
                    ) as live:
                        try:
                            for event in stream:
                                if esc_pressed.is_set():
                                    interrupted = True
                                    break
                                if event.type == "message_start":
                                    usage_input_tokens = getattr(
                                        event.message.usage, "input_tokens", 0
                                    )
                                elif event.type == "message_delta":
                                    usage_output_tokens = getattr(
                                        event.usage, "output_tokens", 0
                                    )
                                elif event.type == "content_block_start":
                                    current_block_type = event.content_block.type
                                elif event.type == "content_block_delta":
                                    if event.delta.type == "thinking_delta":
                                        think.append(event.delta.thinking)
                                    elif event.delta.type == "text_delta":
                                        chunks.append(event.delta.text)
                                elif event.type == "content_block_stop":
                                    current_block_type = None
                                else:
                                    continue
                                live.update(
                                    _tail_window(
                                        "".join(think),
                                        "".join(chunks),
                                        console_stdout.size.height,
                                    ),
                                    refresh=True,
                                )
                        except KeyboardInterrupt:
                            interrupted = True
                    if think:
                        part1 = Text("".join(think))
                        part1 = Padding(
                            part1, (0, 2), style=richStyle(dim=True, italic=True)
                        )
                        console_stdout.print(part1)
                    console_stdout.print(Markdown("".join(chunks)))
                else:
                    try:
                        for event in stream:
                            if esc_pressed.is_set():
                                interrupted = True
                                break
                            if event.type == "message_start":
                                usage_input_tokens = getattr(
                                    event.message.usage, "input_tokens", 0
                                )
                            elif event.type == "message_delta":
                                usage_output_tokens = getattr(
                                    event.usage, "output_tokens", 0
                                )
                            elif event.type == "content_block_start":
                                current_block_type = event.content_block.type
                            elif event.type == "content_block_delta":
                                if event.delta.type == "thinking_delta":
                                    think.append(event.delta.thinking)
                                    print(event.delta.thinking, end="", flush=True)
                                elif event.delta.type == "text_delta":
                                    chunks.append(event.delta.text)
                                    print(event.delta.text, end="", flush=True)
                            elif event.type == "content_block_stop":
                                if current_block_type == "thinking":
                                    print()  # newline after thinking block
                                current_block_type = None
                    except KeyboardInterrupt:
                        interrupted = True
        generated_text = "".join(chunks)
        if not generated_text.endswith("\n"):
            print()
            sys.stdout.flush()
        if interrupted:
            console.print(
                "\n[bold yellow]⚠ Streaming interrupted by user.[/bold yellow]"
            )
            if generated_text:
                partial = generated_text + "\n\n[interrupted by user]"
                self.update_session({"role": "assistant", "content": partial})
            raise StreamInterrupted()
        _total = usage_input_tokens + usage_output_tokens
        if self.verbose:
            console.log(
                f"{self.NAME}> tokens consumed: input {usage_input_tokens}, "
                f"output {usage_output_tokens}, total {_total}."
            )
        self._report_context_remaining(_total)
        self._maybe_summarize_history(_total)
        new_message = {"role": "assistant", "content": generated_text}
        self.update_session(new_message)
        if self.debug:
            console.log("recv:", self.session[-1])
        return self.session[-1]["content"]


class GoogleFrontend(AbstractFrontend):
    """
    https://ai.google.dev/gemini-api/docs
    """

    NAME = "GoogleFrontend"
    debug: bool = False

    def __init__(self, args):
        super().__init__(args)
        try:
            import google.generativeai as genai
        except ImportError:
            console.log(
                'please install gemini package: "pip install google-generativeai"'
            )
            exit(1)
        genai.configure(api_key=args.google_api_key)
        self.client = genai.GenerativeModel(
            args.google_model, system_instruction=args.system_message
        )
        self.model_name = args.google_model
        self.chat = self.client.start_chat()
        self.kwargs = genai.types.GenerationConfig(
            temperature=args.temperature, top_p=args.top_p
        )
        self._query_context_window(genai)
        if args.verbose:
            console.log(
                f"{self.NAME}> model={repr(args.google_model)}, "
                + f"temperature={args.temperature}, top_p={args.top_p}."
            )

    def _query_context_window(self, genai):
        """Query the model context window from the Google Generative AI API."""
        try:
            model_info = genai.get_model(f"models/{self.model_name}")
            self.context_window = getattr(model_info, "input_token_limit", None)
        except Exception:
            self.context_window = None

        if self.context_window is None:
            self.context_window = self.default_context_window

    def oneshot(self, message: str, *, retry: bool = True) -> str:
        def _func():
            _callable = self.client.generate_content
            result = _callable(message, generation_config=self.kwargs)
            return result.text

        from google.api_core.exceptions import ResourceExhausted

        return retry_ratelimit(_func, ResourceExhausted)()

    def query(self, messages: Union[List, Dict, str]) -> list:
        # add the message into the session
        self.update_session(messages)
        if self.debug:
            console.log("send:", self.session[-1])
        chunks = []
        response = self.chat.send_message(
            self.session[-1]["content"], stream=True, generation_config=self.kwargs
        )
        interrupted = False
        with _esc_watcher() as esc_pressed:
            if self.render_markdown:
                with Live(
                    Markdown(""),
                    console=console_stdout,
                    vertical_overflow="crop",
                    transient=True,
                ) as live:
                    try:
                        for chunk in response:
                            if esc_pressed.is_set():
                                interrupted = True
                                break
                            chunks.append(chunk.text)
                            live.update(
                                _tail_window(
                                    "", "".join(chunks), console_stdout.size.height
                                ),
                                refresh=True,
                            )
                    except KeyboardInterrupt:
                        interrupted = True
                console_stdout.print(Markdown("".join(chunks)))
            else:
                try:
                    for chunk in response:
                        if esc_pressed.is_set():
                            interrupted = True
                            break
                        chunks.append(chunk.text)
                        print(chunk.text, end="", flush=True)
                except KeyboardInterrupt:
                    interrupted = True
        generated_text = "".join(chunks)
        if not generated_text.endswith("\n") and generated_text:
            print()
            sys.stdout.flush()
        if interrupted:
            console.print(
                "\n[bold yellow]⚠ Streaming interrupted by user.[/bold yellow]"
            )
            if generated_text:
                partial = generated_text + "\n\n[interrupted by user]"
                self.update_session({"role": "assistant", "content": partial})
            raise StreamInterrupted()
        try:
            usage = response.usage_metadata
            if self.verbose:
                console.log(
                    f"{self.NAME}> tokens consumed: input {usage.prompt_token_count}, "
                    f"output {usage.candidates_token_count}, total {usage.total_token_count}."
                )
            self._report_context_remaining(usage.total_token_count)
            self._maybe_summarize_history(usage.total_token_count)
        except Exception:
            pass
        new_message = {"role": "assistant", "content": generated_text}
        self.update_session(new_message)
        if self.debug:
            console.log("recv:", self.session[-1])
        return self.session[-1]["content"]


class DeepSeekFrontend(OpenAIFrontend):
    """
    https://api-docs.deepseek.com/
    """

    NAME = "DeepSeekFrontend"

    def __init__(self, args):
        AbstractFrontend.__init__(self, args)
        from openai import OpenAI

        self.client = OpenAI(
            api_key=args.deepseek_api_key, base_url=args.deepseek_base_url
        )
        if args.deepseek_model not in ("deepseek-reasoner"):
            # see the usage recommendations at
            # https://huggingface.co/deepseek-ai/DeepSeek-R1
            self.session.append({"role": "system", "content": args.system_message})
        self.model = args.deepseek_model
        self.kwargs = {"temperature": args.temperature, "top_p": args.top_p}
        self._query_context_window()
        if args.verbose:
            console.log(
                f"{self.NAME}> model={repr(self.model)}, "
                + f"temperature={args.temperature}, top_p={args.top_p}."
            )


class vLLMFrontend(OpenAIFrontend):
    """
    https://docs.vllm.ai/en/stable/serving/openai_compatible_server.html
    """

    NAME = "vLLMFrontend"

    def __init__(self, args):
        AbstractFrontend.__init__(self, args)
        from openai import OpenAI

        self.client = OpenAI(api_key="your-vllm-api-key", base_url=args.vllm_base_url)
        self.session.append({"role": "system", "content": args.system_message})
        self.model = args.vllm_model
        self.kwargs = {"temperature": args.temperature, "top_p": args.top_p}
        self._query_context_window()
        if args.verbose:
            console.log(
                f"{self.NAME}> model={repr(self.model)}, "
                + f"temperature={args.temperature}, top_p={args.top_p}."
            )


class ZMQFrontend(AbstractFrontend):
    """
    ZMQ frontend communicates with a self-hosted ZMQ backend.
    """

    NAME = "ZMQFrontend"
    debug: bool = False

    def __init__(self, args):
        import zmq

        super().__init__(args)
        self.zmq_backend = args.zmq_backend
        self.socket = zmq.Context().socket(zmq.REQ)
        self.socket.connect(self.zmq_backend)
        console.log(f"{self.NAME}> Connected to ZMQ backend {self.zmq_backend}.")
        self.temperature = getattr(args, "temperature", None)
        self.top_p = getattr(args, "top_p", None)

    def query(self, content: Union[List, Dict, str]) -> list:
        if isinstance(content, list):
            self.session = content
        elif isinstance(content, dict):
            self.session.append(content)
        elif isinstance(content, str):
            self.session.append({"role": "user", "content": content})
        _check(self.session)

        payload = {"session": self.session}
        if self.temperature is not None:
            payload["temperature"] = self.temperature
        if self.top_p is not None:
            payload["top_p"] = self.top_p
        msg_json = json.dumps(payload)
        if self.debug:
            console.log("send:", msg_json)
        self.socket.send_string(msg_json)

        chunks = []
        interrupted = False

        def _drain_zmq():
            """Drain remaining backend messages to keep REQ/REP in sync."""
            try:
                while True:
                    if self.socket.poll(500):
                        msg = self.socket.recv_json()
                        if msg.get("status") == "chunk":
                            self.socket.send_json({"status": "ack"})
                        else:
                            break
                    else:
                        break
            except Exception:
                pass

        with _esc_watcher() as esc_pressed:
            if self.render_markdown:
                with Live(
                    Markdown(""),
                    console=console_stdout,
                    vertical_overflow="crop",
                    transient=True,
                ) as live:
                    try:
                        while True:
                            msg = self.socket.recv_json()
                            if esc_pressed.is_set():
                                interrupted = True
                                if msg.get("status") == "chunk":
                                    self.socket.send_json({"status": "ack"})
                                _drain_zmq()
                                break
                            if msg.get("status") == "chunk":
                                chunk = msg.get("content", "")
                                chunks.append(chunk)
                                live.update(
                                    _tail_window(
                                        "", "".join(chunks), console_stdout.size.height
                                    ),
                                    refresh=True,
                                )
                                self.socket.send_json({"status": "ack"})
                            elif msg.get("status") == "done":
                                self.session = msg.get("session")
                                break
                            elif msg.get("status") == "error":
                                raise RuntimeError(
                                    f"Backend error: {msg.get('content')}"
                                )
                            else:
                                break
                    except KeyboardInterrupt:
                        interrupted = True
                        _drain_zmq()
                console_stdout.print(Markdown("".join(chunks)))
            else:
                try:
                    while True:
                        msg = self.socket.recv_json()
                        if esc_pressed.is_set():
                            interrupted = True
                            if msg.get("status") == "chunk":
                                self.socket.send_json({"status": "ack"})
                            _drain_zmq()
                            break
                        if msg.get("status") == "chunk":
                            chunk = msg.get("content", "")
                            chunks.append(chunk)
                            print(chunk, end="", flush=True)
                            self.socket.send_json({"status": "ack"})
                        elif msg.get("status") == "done":
                            self.session = msg.get("session")
                            break
                        elif msg.get("status") == "error":
                            raise RuntimeError(f"Backend error: {msg.get('content')}")
                        else:
                            break
                except KeyboardInterrupt:
                    interrupted = True
                    _drain_zmq()

        generated_text = "".join(chunks)
        if not generated_text.endswith("\n") and generated_text:
            print()
            sys.stdout.flush()
        if interrupted:
            console.print(
                "\n[bold yellow]⚠ Streaming interrupted by user.[/bold yellow]"
            )
            if generated_text:
                partial = generated_text + "\n\n[interrupted by user]"
                self.session.append({"role": "assistant", "content": partial})
            raise StreamInterrupted()

        _check(self.session)
        if self.debug:
            console.log("recv:", self.session[-1])
        return self.session[-1]["content"]


def get_username():
    try:
        import getpass

        return getpass.getuser()
    except Exception:
        pass
    try:
        import pwd

        return pwd.getpwuid(os.getuid())[0]
    except Exception:
        pass
    try:
        return os.getlogin()
    except Exception:
        pass
    # common shell env
    for env in ("USER", "USERNAME", "LOGNAME"):
        if u := os.environ.get(env):
            return u
    # final fallback
    return "user"


def create_frontend(args):
    if args.frontend == "zmq":
        frontend = ZMQFrontend(args)
    elif args.frontend == "openai":
        frontend = OpenAIFrontend(args)
    elif args.frontend == "anthropic":
        frontend = AnthropicFrontend(args)
    elif args.frontend == "google":
        frontend = GoogleFrontend(args)
    elif args.frontend == "deepseek":
        frontend = DeepSeekFrontend(args)
    elif args.frontend == "vllm":
        frontend = vLLMFrontend(args)
    elif args.frontend == "dryrun":
        frontend = None
    elif args.frontend == "echo":
        frontend = EchoFrontend(args)
    else:
        raise NotImplementedError
    return frontend


def interact_once(f: AbstractFrontend, text: str) -> None:
    """
    we have prepared text -- let frontend send it to LLM, and this function
    will print the LLM reply.

    f: any frontend instance from the current source file.
    text: the text to be sent to LLM.
    """
    end = "" if not f.render_markdown else "\n"
    esc_hint = " (press ESC to interrupt)"
    if f.monochrome:
        lprompt = escape(f"LLM[{2+len(f)}]> ") + esc_hint
        console.print(lprompt, end=end, highlight=False, markup=False)
    else:
        lprompt = f"[bold green]LLM[{2+len(f)}]>[/bold green] [dim]{esc_hint}[/dim]"
        console.print(lprompt, end=end)
    try:
        _ = f(text)
    except StreamInterrupted:
        pass  # partial response already saved; caller can continue


def interact_with(f: AbstractFrontend) -> None:
    # create prompt_toolkit style
    if f.monochrome:
        prompt_style = Style([("prompt", "bold")])
    else:
        prompt_style = Style(
            [("prompt", "bold fg:ansibrightcyan"), ("", "bold ansiwhite")]
        )

    # Completer with several keywords keywords to be completed
    class CustomCompleter(Completer):
        def get_completions(self, document, complete_event):
            # Get the current text before the cursor
            text_before_cursor = document.text_before_cursor

            # Check if the text starts with '/'
            if text_before_cursor.startswith("/"):
                # Define the available keywords
                keywords = ["/quit", "/save", "/reset"]

                # Generate completions for each keyword
                for keyword in keywords:
                    if keyword.startswith(text_before_cursor):
                        yield Completion(keyword, -len(text_before_cursor))

    # start prompt session
    prompt_session = PromptSession(
        style=prompt_style, multiline=f.multiline, completer=CustomCompleter()
    )

    # if multiline is enabled, print additional help message
    if f.multiline:
        console.print(
            "In multiline mode, please press [Meta+Enter], or [Esc] followed by [Enter] to send the message."
        )

    # loop
    user = get_username()
    try:
        while text := prompt_session.prompt(f"{user}[{1+len(f)}]> "):
            # parse escaped interaction commands
            if text.startswith("/"):
                cmd = shlex.split(text)
                if cmd[0] == "/save":
                    # save the last LLM reply to a file
                    if len(cmd) != 2:
                        console.print("syntax error: /save <path>")
                        continue
                    path = cmd[-1]
                    with open(path, "wt") as fp:
                        fp.write(f.session[-1]["content"])
                    console.log(f"The last LLM response is saved at {path}")
                elif cmd[0] == "/reset":
                    if len(cmd) != 1:
                        console.print("syntax error: /reset")
                        continue
                    f.reset()
                elif cmd[0] == "/quit":
                    if len(cmd) != 1:
                        console.print("syntax error: /quit")
                        continue
                    break
                else:
                    console.print(f"unknown command: {cmd[0]}")
            else:
                interact_once(f, text)
    except EOFError:
        pass
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    ag = argparse.ArgumentParser()
    ag.add_argument("--zmq_backend", "-B", default="tcp://localhost:11177")
    ag.add_argument(
        "--frontend",
        "-F",
        default="zmq",
        choices=("dryrun", "zmq", "openai", "anthropic", "google", "vllm"),
    )
    ag.add_argument("--debgpt_home", default=defaults.HOME)
    ag = ag.parse_args()
    console.print(ag)

    frontend = create_frontend(ag)
    f = frontend
    import IPython

    IPython.embed(colors="neutral")
