"""
title: Live Progress Status
author: owuitools
version: 0.2.0
description: Shows what the model is doing right now - processing the prompt, thinking, and each tool call as it is written.
"""

# WHY THIS EXISTS
#
# Until the first token arrives, Open WebUI shows nothing at all. On this
# deployment the prompt is ~9.4k tokens of system prompt plus 20 tool
# schemas, and a forced re-prefill takes many seconds, so "nothing at all"
# is the most common thing on screen. Worse, the model spends ~70% of its
# generated characters thinking, and tool calls are written token by token
# into a block that is collapsed by default - so even once it IS producing
# output, there is no visible signal of which phase it is in.
#
# This filter emits Open WebUI status events on phase CHANGES only:
#
#     Processing prompt          <- inlet, before the model is called
#     Thinking                   <- first reasoning delta
#     Thought: <opening of it>   <- when a reasoning block ends
#     Calling web_search         <- a tool call is being written
#     Writing response           <- first content delta
#
# Together those read as a narrative of the whole turn, which is what makes
# the collapsed tool blocks acceptable: the chain of thought and the tool
# sequence stay visible without any tool input or output on screen.
#
# WHERE THE NUMBERS COME FROM
#
# The status line used to claim "prompt processed in N s" for a figure that
# was really inlet-to-first-token wall clock: Open WebUI middleware, queue
# wait and model load included. Graded runs put it 2-3x either side of
# Ollama's own prefill for the same turn, so the one number on screen was
# the one number that was wrong; 35 of 40 rubrics flagged it.
#
# Ollama's real figures are available, and this filter now reads them
# instead of guessing. Open WebUI's Ollama->OpenAI conversion builds a
# usage object on the final chunk of each round:
#
#   utils/response.py  convert_ollama_usage_to_openai()
#     'prompt_eval_count', 'prompt_eval_duration',   (nanoseconds)
#     'eval_count', 'eval_duration', 'total_duration', 'load_duration'
#   utils/misc.py  openai_chat_chunk_message_template()
#     template['usage'] = usage        <- TOP LEVEL, beside 'choices'
#
# and stream filters are handed that chunk verbatim, before any
# normalisation (middleware.py, filter_type='stream'). The same numbers
# reach outlet on the stored message's 'usage', so outlet can fall back to
# it - with one caveat that is worth knowing when the two disagree:
# merge_usage() sums only input_tokens/output_tokens and keeps the LATEST
# value of every provider-specific key, so a fallback prefill is the last
# round's warm re-prefill, while the stream path deliberately keeps the
# FIRST one, which is the cold prefill the user actually waited through.
#
# WHY THIS CANNOT AFFECT THE MODEL'S OUTPUT
#
# Status events travel over the websocket and are stored on a SIBLING key
# of the message ('statusHistory'), never in 'content'. The payload sent to
# the model is rebuilt from scratch by an explicit allowlist:
#
#   utils/payload.py  convert_messages_openai_to_ollama()
#     new_message = {'role': message['role']}
#     # then only: thinking, content, tool_calls, tool_call_id, images
#
# 'statusHistory' is not in that list, and the backend never references it
# outside persistence. The token sequence is therefore byte-identical with
# this filter on or off - no prompt change, no prefix-cache perturbation.
# The one thing that does reach the chat record is the empty-answer note
# below, and only on a turn that produced no answer to disturb.
#
# TWO SHARP EDGES IN THE STREAM HOOK, both load-bearing:
#
#   1. The return value REPLACES the chunk, and a FALSY return DROPS it:
#        data = wrap_item(JSONCodec.dumps(event)) if event else None
#      So stream() must return `event` on every single path. Returning None
#      would silently delete the model's output.
#   2. Exceptions from a stream filter are re-raised by the caller
#      (utils/filter.py logs then `raise`), which would break the response
#      mid-flight. So the whole body is wrapped and always falls through to
#      `return event`.
#
# The module object is cached across requests (load_from_db=False for
# stream), so per-request state must be keyed by message id rather than
# held on self - six users share this instance.

import time

from pydantic import BaseModel, Field

# Per-message progress state, keyed by message id. Bounded below so a
# request whose outlet never runs (client disconnect, cancelled generation)
# cannot leak entries forever.
_STATE: dict[str, dict] = {}
_MAX_TRACKED = 64

# Shown in the message itself when the turn ended with no answer text. One
# graded run spent 140 s and 8,201 tokens repeating a single sentence inside
# its reasoning block, emitted nothing, and the indicator said "Done" - so
# the failure looked exactly like a success that had scrolled off screen.
# 0.11.3's filter API is inlet/stream/request/outlet only, with no hook that
# re-enters the chat pipeline, so telling the user to regenerate is the most
# this filter can do.
_EMPTY_NOTE = (
    "_No answer was produced: the model ended its turn while still reasoning. "
    "Regenerate to try again._"
)


def _new_state() -> dict:
    """The per-message progress record. One factory, used by both creation
    sites - they were duplicated literals once and silently drifted."""
    return {
        "started": time.monotonic(),
        "phase": None,
        "announced": set(),
        "round": 0,
        "prompt_reported": False,
        "tools": 0,
        "reasoning": [],
        # Ollama's own figures, filled in from the usage chunk. prefill is
        # first-wins (the cold one); generated tokens accumulate over every
        # round of a tool-calling turn.
        "prefill_s": None,
        "prompt_tokens": None,
        "gen_tokens": 0,
    }


def _prune() -> None:
    while len(_STATE) > _MAX_TRACKED:
        # dicts are insertion-ordered, so this drops the oldest.
        _STATE.pop(next(iter(_STATE)), None)


def _number(value) -> float | None:
    """Usage values arrive as whatever the provider sent. Anything that is
    not a real number (None, 'N/A', a string) is treated as absent."""
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        return None
    return float(value)


class Filter:
    class Valves(BaseModel):
        enabled: bool = Field(
            default=True,
            description="Turn the whole filter off without deleting it.",
        )
        show_prompt_processing: bool = Field(
            default=True,
            description="Show 'Processing prompt' before the first token. This is the phase "
            "that currently shows nothing at all, and on a cold prompt it is the longest.",
        )
        show_thinking: bool = Field(
            default=True,
            description="Show 'Thinking' once the model starts producing reasoning tokens.",
        )
        show_tool_calls: bool = Field(
            default=True,
            description="Show 'Calling <tool>' as each tool call is written. Tool calls are "
            "still counted for the final line when this is off.",
        )
        report_prompt_time: bool = Field(
            default=True,
            description="Append timing to the status line: how long the first token took to "
            "arrive, and - on the final line, from Ollama's own counters - how long the "
            "prefill took and how many tokens went in and came out.",
        )
        thought_excerpt_chars: int = Field(
            default=160,
            description="When a block of reasoning ends, post an excerpt of it as a status so "
            "the timeline reads as a narrative of what the model was doing. 0 disables it. "
            "The renderer clamps a status to one line, so values much above ~180 are wasted.",
        )
        clear_when_done: bool = Field(
            default=False,
            description="Hide the status line entirely once the response finishes. Off by "
            "default so the finished message keeps a short record of what it did, which "
            "matches how Open WebUI's own built-in tools behave.",
        )

    def __init__(self):
        self.valves = self.Valves()

    # ---------------------------------------------------------------- utils

    async def _emit(self, emitter, description: str, *, done: bool = False, hidden: bool = False) -> None:
        if emitter is None:
            return
        try:
            await emitter(
                {
                    "type": "status",
                    "data": {"description": description, "done": done, "hidden": hidden},
                }
            )
        except Exception:
            # A status line is a nicety; never let it disturb the response.
            pass

    def _message_id(self, metadata) -> str | None:
        if isinstance(metadata, dict):
            return metadata.get("message_id") or metadata.get("chat_id")
        return None

    def _state(self, message_id: str | None) -> dict | None:
        if not message_id:
            return None
        state = _STATE.get(message_id)
        if state is None:
            # stream() can legitimately run without inlet having recorded
            # anything, so build a usable state rather than giving up.
            state = _new_state()
            _STATE[message_id] = state
            _prune()
        return state

    def _record_usage(self, state: dict, usage: dict) -> None:
        """Take Ollama's own prefill and token counts off a usage chunk.

        Prefill is first-wins because a tool-calling turn prefills once per
        round and only the first one is the wait the user sat through; every
        later round is served warm off the prefix cache. Generated tokens
        accumulate, because all of them were generated on this turn.
        """
        if state["prefill_s"] is None:
            duration = _number(usage.get("prompt_eval_duration"))
            if duration and duration > 0:
                state["prefill_s"] = duration / 1e9
        if state["prompt_tokens"] is None:
            count = _number(usage.get("prompt_eval_count"))
            if count and count > 0:
                state["prompt_tokens"] = int(count)
        generated = _number(usage.get("eval_count"))
        if generated and generated > 0:
            state["gen_tokens"] += int(generated)

    def _final_usage(self, state: dict | None, message: dict | None) -> tuple:
        """(prefill_s, prompt_tokens, gen_tokens) for the final line.

        The stream numbers are preferred; the stored message's usage is the
        fallback for the paths where stream() never saw a usage chunk - a
        non-streaming response, or a filter loaded after the stream began.
        """
        prefill_s = state["prefill_s"] if state else None
        prompt_tokens = state["prompt_tokens"] if state else None
        gen_tokens = state["gen_tokens"] if state else 0

        usage = message.get("usage") if isinstance(message, dict) else None
        if isinstance(usage, dict):
            if prefill_s is None:
                duration = _number(usage.get("prompt_eval_duration"))
                if duration and duration > 0:
                    prefill_s = duration / 1e9
            if prompt_tokens is None:
                count = _number(usage.get("prompt_eval_count")) or _number(usage.get("prompt_tokens"))
                if count and count > 0:
                    prompt_tokens = int(count)
            if not gen_tokens:
                # output_tokens is the one Open WebUI sums across rounds;
                # the Ollama-native keys hold the last round only.
                generated = (
                    _number(usage.get("output_tokens"))
                    or _number(usage.get("completion_tokens"))
                    or _number(usage.get("eval_count"))
                )
                if generated and generated > 0:
                    gen_tokens = int(generated)

        return prefill_s, prompt_tokens, gen_tokens

    def _usage_clause(self, prefill_s, prompt_tokens, gen_tokens) -> str:
        """The measured half of the final line. Empty when nothing was
        measured - an absent clause is honest, an invented one is not."""
        if not self.valves.report_prompt_time:
            return ""
        parts = []
        if prefill_s is not None:
            if prompt_tokens:
                parts.append(f"prefill {prefill_s:.1f} s for {prompt_tokens:,} tokens")
            else:
                parts.append(f"prefill {prefill_s:.1f} s")
        if gen_tokens:
            parts.append(f"{gen_tokens:,} generated")
        return "".join(f" · {part}" for part in parts)

    async def _flush_thought(self, state: dict, emitter) -> None:
        """Post an excerpt of the reasoning block that just ended, then drop it.

        This is what makes the collapsed view usable: the tool payloads stay
        shut, but the status timeline still reads as a narrative of what the
        model was thinking between calls. The buffer is cleared either way, so
        a disabled excerpt cannot accumulate text for the life of the request.

        The excerpt shows both ends of the block, not just the opening. Head
        truncation once surfaced "…I should verify with actual execution since
        I have tools. Actually, I" - cutting the sentence exactly where the
        model reversed itself, so a user watching live would have believed it
        was about to run code it had in fact just decided not to run. The
        conclusion of a thought lives at its end.
        """
        buffered = state["reasoning"]
        state["reasoning"] = []
        limit = self.valves.thought_excerpt_chars
        if limit <= 0 or not buffered:
            return
        # Collapse to a single line - the renderer clamps to one line anyway,
        # and newlines in a status read as ragged whitespace.
        text = " ".join("".join(buffered).split())
        if not text:
            return
        if len(text) > limit:
            head_budget = max(1, int(limit * 0.55))
            tail_budget = max(1, int(limit * 0.40))
            # Prefer breaking at a word boundary near each edge.
            cut = text.rfind(" ", int(head_budget * 0.6), head_budget)
            head = text[: cut if cut != -1 else head_budget].rstrip()
            start = len(text) - tail_budget
            space = text.find(" ", start)
            if space != -1 and space + 1 < len(text):
                start = space + 1
            text = f"{head} … {text[start:].lstrip()}"
        await self._emit(emitter, f"Thought: {text}", done=True)

    def _suffix(self, state: dict) -> str:
        """Report, once, how long the first token took to arrive.

        Deliberately not called prefill: this is inlet-to-first-token wall
        clock and includes Open WebUI's own work, the queue and any model
        load. The real prefill is not known until the usage chunk at the end
        of the round, and it goes on the final line instead.
        """
        if state["prompt_reported"] or not self.valves.report_prompt_time:
            return ""
        state["prompt_reported"] = True
        elapsed = time.monotonic() - state["started"]
        if elapsed < 0.5:
            return ""
        return f" · first token after {elapsed:.1f} s"

    def _mark_empty(self, message: dict) -> None:
        """Leave a visible note on a turn that produced no answer text.

        Both places, as inline_tool_images.py does it for its images: Open
        WebUI renders `output` when it is present and falls back to `content`
        only when it is not (misc.py get_output_text), so writing content
        alone stores the note and shows nothing.
        """
        content = message.get("content") or ""
        message["content"] = (content.rstrip() + "\n\n" + _EMPTY_NOTE).lstrip()

        output_items = message.get("output")
        if not isinstance(output_items, list):
            return
        message_items = [
            item for item in output_items if isinstance(item, dict) and item.get("type") == "message"
        ]
        if message_items:
            last = message_items[-1]
            parts = last.get("content")
            if not isinstance(parts, list):
                parts = []
            texts = [p for p in parts if isinstance(p, dict) and p.get("type") == "output_text"]
            if texts:
                texts[-1]["text"] = (texts[-1].get("text") or "").rstrip() + "\n\n" + _EMPTY_NOTE
            else:
                parts.append({"type": "output_text", "text": _EMPTY_NOTE})
            last["content"] = parts
        else:
            # An output list of nothing but reasoning and tool traffic renders
            # as an empty message, and the note in `content` would never be
            # reached. Give it the one item the renderer reads.
            output_items.append(
                {
                    "type": "message",
                    "role": "assistant",
                    "content": [{"type": "output_text", "text": _EMPTY_NOTE}],
                }
            )

    # ---------------------------------------------------------------- hooks

    async def inlet(self, body: dict, __event_emitter__=None, __metadata__=None) -> dict:
        """Mark the start of prompt processing. `body` is returned untouched."""
        if not self.valves.enabled:
            return body
        try:
            message_id = self._message_id(__metadata__)
            continued = False
            if message_id:
                state = _STATE.get(message_id)
                # inlet runs a SECOND time for the same assistant message when
                # the turn contains an interactive tool call - the native
                # ask_user widget re-enters the pipeline. Wiping the state
                # there zeroed the tool counter mid-turn, and a turn with
                # three real calls reported "Done · 2 tool calls". Keep what
                # has already been counted; reset only the per-round fields.
                if state and (state["tools"] or state["phase"]):
                    continued = True
                    state["round"] += 1
                    state["phase"] = None
                    state["started"] = time.monotonic()
                    state["prompt_reported"] = False
                else:
                    _STATE[message_id] = _new_state()
                _prune()
            if self.valves.show_prompt_processing:
                await self._emit(
                    __event_emitter__,
                    "Processing prompt (continued)" if continued else "Processing prompt",
                    done=False,
                )
        except Exception:
            pass
        # Never modified. This filter must not touch the prompt.
        return body

    async def stream(self, event: dict, __event_emitter__=None, __metadata__=None) -> dict:
        """Announce phase changes. MUST return `event` on every path."""
        if not self.valves.enabled:
            return event
        try:
            if not isinstance(event, dict):
                return event

            state = self._state(self._message_id(__metadata__))
            if state is None:
                return event

            # Read before the `choices` checks below: the usage chunk carries
            # usage at the top level and an empty delta, and some providers
            # send it with no choices at all.
            usage = event.get("usage")
            if isinstance(usage, dict):
                self._record_usage(state, usage)

            choices = event.get("choices")
            if not isinstance(choices, list) or not choices:
                return event
            first = choices[0]
            if not isinstance(first, dict):
                return event
            delta = first.get("delta")
            if not isinstance(delta, dict):
                return event

            tool_calls = delta.get("tool_calls")
            if tool_calls:
                for call in tool_calls if isinstance(tool_calls, list) else []:
                    if not isinstance(call, dict):
                        continue
                    function = call.get("function")
                    name = function.get("name") if isinstance(function, dict) else None
                    if not name:
                        continue
                    # A tool id is the reliable key; index plus round is the
                    # fallback, so the same tool called twice still counts twice.
                    key = call.get("id") or f"{state['round']}:{call.get('index', 0)}"
                    if key in state["announced"]:
                        continue
                    state["announced"].add(key)
                    # Counting, de-duplication and the phase change all sit
                    # OUTSIDE the display valve: they used to be inside it, so
                    # turning the announcements off silently zeroed the final
                    # tool count as well, and stalled the round counter that
                    # the de-duplication key depends on. Only the announcement
                    # itself is optional.
                    state["tools"] += 1
                    # The reasoning that led to this call belongs before it.
                    await self._flush_thought(state, __event_emitter__)
                    state["phase"] = "tool"
                    if self.valves.show_tool_calls:
                        await self._emit(
                            __event_emitter__,
                            f"Calling {name}{self._suffix(state)}",
                            done=False,
                        )
                return event

            reasoning = delta.get("reasoning_content")
            if reasoning:
                if isinstance(reasoning, str) and self.valves.thought_excerpt_chars > 0:
                    # Only keep what an excerpt could possibly need, so a long
                    # think cannot grow this without bound.
                    if sum(len(p) for p in state["reasoning"]) < 4096:
                        state["reasoning"].append(reasoning)
                if self.valves.show_thinking and state["phase"] != "thinking":
                    if state["phase"] == "tool":
                        state["round"] += 1
                    state["phase"] = "thinking"
                    await self._emit(
                        __event_emitter__, f"Thinking{self._suffix(state)}", done=False
                    )
            elif delta.get("content"):
                if state["phase"] != "writing":
                    if state["phase"] == "tool":
                        state["round"] += 1
                    await self._flush_thought(state, __event_emitter__)
                    state["phase"] = "writing"
                    await self._emit(
                        __event_emitter__,
                        f"Writing response{self._suffix(state)}",
                        done=False,
                    )
        except Exception:
            pass
        return event

    async def outlet(self, body: dict, __event_emitter__=None, __metadata__=None) -> dict:
        """Settle the status line and drop the per-message state."""
        if not self.valves.enabled:
            return body
        try:
            message_id = self._message_id(__metadata__)
            state = _STATE.pop(message_id, None) if message_id else None

            # A response that ended while still reasoning still has something
            # worth showing.
            if state and not self.valves.clear_when_done:
                await self._flush_thought(state, __event_emitter__)

            messages = body.get("messages")
            message = None
            if isinstance(messages, list) and messages and isinstance(messages[-1], dict):
                message = messages[-1]

            prefill_s, prompt_tokens, gen_tokens = self._final_usage(state, message)

            # Open WebUI builds outlet's `content` as content OR the rendered
            # text of `output`, so blank here means the turn really produced
            # no answer - not that the answer lives somewhere else.
            empty = (
                message is not None
                and message.get("role", "assistant") == "assistant"
                and not (message.get("content") or "").strip()
            )
            if empty:
                self._mark_empty(message)

            if self.valves.clear_when_done and not empty:
                # StatusHistory hides the whole block when the LATEST status
                # is hidden, so this clears the line entirely.
                await self._emit(__event_emitter__, "", done=True, hidden=True)
            else:
                if empty:
                    # Never hidden, even with clear_when_done: a turn that
                    # produced nothing is the one turn whose status matters.
                    # The token count is dropped rather than printed as 0 when
                    # no usage arrived - "(0 tokens)" would be a measurement
                    # nobody took.
                    spent = f" ({gen_tokens:,} tokens)" if gen_tokens else ""
                    summary = (
                        "Done · no answer produced — the model stopped inside its "
                        f"reasoning{spent}. Press Regenerate."
                    )
                else:
                    tools = state["tools"] if state else 0
                    if tools == 1:
                        summary = "Done · 1 tool call"
                    elif tools > 1:
                        summary = f"Done · {tools} tool calls"
                    else:
                        summary = "Done"
                    summary += self._usage_clause(prefill_s, prompt_tokens, gen_tokens)
                await self._emit(__event_emitter__, summary, done=True)
        except Exception:
            pass
        return body