import hljs from "npm:highlight.js/lib/core";
import javascript from "npm:highlight.js/lib/languages/javascript";
import python from "npm:highlight.js/lib/languages/python";
import rust from "npm:highlight.js/lib/languages/rust";

hljs.registerLanguage("javascript", javascript);
hljs.registerLanguage("rust", rust);
hljs.registerLanguage("python", python);

export class ChatResponse {
  element = (<article className="assistant" />);

  currentLine: Text | undefined;

  codeBlockContext: Element | undefined;
  thinkingContext: HTMLDetailsElement | undefined;

  currentContext(): Element {
    if (this.codeBlockContext) return this.codeBlockContext;
    if (this.thinkingContext) return this.thinkingContext;
    return this.element;
  }

  finalizeLine() {
    if (!this.currentLine?.textContent) return;
    const line = this.currentLine.textContent;

    if (!this.thinkingContext && line === "<think>") {
      this.thinkingContext = (
        <details open>
          <summary>{`<think />`}</summary>
        </details>
      ) as HTMLDetailsElement;
      this.element.append(this.thinkingContext);
      this.currentLine.remove();
    } else if (this.thinkingContext && line === "</think>") {
      this.currentLine.remove();
      this.thinkingContext.open = false;
      this.thinkingContext = undefined;
    }

    if (line.startsWith("```")) {
      const remainder = line.substring(3).trimStart();
      this.currentLine.remove();

      if (!this.codeBlockContext) {
        const pre = (
          <pre>
            <code className={`language-${remainder}`} />
          </pre>
        );
        this.currentContext().append(pre);
        this.codeBlockContext = pre.querySelector("code")!;
      } else {
        hljs.highlightElement(this.codeBlockContext as HTMLElement);
        // TODO: highlight.js the guy

        this.codeBlockContext = undefined;
        this.currentContext().append(remainder);
      }
    }
  }

  append(text: string) {
    const tokens = text
      .split(/(\s+)/)
      .flatMap(s => s.split(/(\n)/))
      .filter(s => s);
    for (const token of tokens) {
      if (token === "\n") {
        if (!this.currentLine) continue;

        this.finalizeLine();
        this.currentLine.appendData(token);
        this.currentLine = undefined;
        continue;
      }

      if (this.currentLine === undefined) {
        this.currentLine = document.createTextNode(token);
        this.currentContext().append(this.currentLine);
      } else {
        this.currentLine.appendData(token);
      }
    }
  }

  finalize() {
    this.finalizeLine();

    // TODO: check if we have a dangling think section or code block and lift it back out
  }
}