import type { AssistantMessage } from "@earendil-works/pi-ai";
import {
  AssistantMessageComponent,
  type ExtensionAPI,
} from "@earendil-works/pi-coding-agent";
import { Markdown, truncateToWidth, type Component } from "@earendil-works/pi-tui";

const LINES = 8;
type AssistantInternals = {
  contentContainer: { children: Component[] };
  lastMessage?: AssistantMessage;
};
type ThinkingMarkdown = {
  defaultTextStyle?: { italic?: boolean };
  paddingX: number;
};
type Prototype = typeof AssistantMessageComponent.prototype & {
  setExpanded?: (expanded: boolean) => void;
};

export default function (pi: ExtensionAPI) {
  const prototype = AssistantMessageComponent.prototype as Prototype;
  const updateContent = prototype.updateContent;
  const setExpanded = prototype.setExpanded;
  const expanded = new WeakSet<AssistantMessageComponent>();

  prototype.updateContent = function (message) {
    updateContent.call(this, message);
    if (expanded.has(this)) return;

    const { children } = (this as unknown as AssistantInternals).contentContainer;
    for (const child of children) {
      if (!(child instanceof Markdown)) continue;
      const markdown = child as unknown as ThinkingMarkdown;
      if (!markdown.defaultTextStyle?.italic) continue;
      const render = child.render.bind(child);
      child.render = (width) => {
        const lines = render(width);
        if (lines.length <= LINES) return lines;
        const indicator = `(${lines.length - LINES} earlier lines hidden - ctrl+o to expand)`;
        const margin = " ".repeat(markdown.paddingX);
        return [
          ...lines.slice(-LINES),
          truncateToWidth(`${margin}${indicator}${margin}`, width, "", true),
        ];
      };
    }
  };

  prototype.setExpanded = function (value) {
    value ? expanded.add(this) : expanded.delete(this);
    setExpanded?.call(this, value);
    const message = (this as unknown as AssistantInternals).lastMessage;
    if (message) this.updateContent(message);
  };

  pi.on("session_shutdown", () => {
    prototype.updateContent = updateContent;
    if (setExpanded) prototype.setExpanded = setExpanded;
    else delete prototype.setExpanded;
  });
}