"use client";

import { FormEvent, useState } from "react";

import {
  formatProbabilityPercent,
  getConfidenceBand,
  type Lab1ConfidenceMessage,
  type Lab1ConfidenceResponse,
  type Lab1ResponseToken,
} from "~/lib/lab1-confidence";

type UserTurn = {
  content: string;
  id: string;
  role: "user";
};

type AssistantTurn = Lab1ConfidenceResponse & {
  error?: string;
  id: string;
};

type ChatTurn = AssistantTurn | UserTurn;

const starterPrompts = [
  "The quick brown fox",
  "Write one sentence explaining what a firewall does.",
  "List three words that describe a phishing email.",
] as const;

function buildTurnId() {
  return `lab1-turn-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
}

function toConversation(messages: ChatTurn[]) {
  return messages.map(({ content, role }) => ({ content, role }));
}

function renderTooltip(token: Lab1ResponseToken) {
  return (
    <span className="lab1-confidence__tooltip">
      <strong>{formatProbabilityPercent(token.probability)}</strong>
      {token.topAlternatives.length > 0 ? (
        <span className="lab1-confidence__tooltip-list">
          {token.topAlternatives.map((candidate) => (
            <span key={`${token.token}-${candidate.token}`}>
              {formatProbabilityPercent(candidate.probability)}:{" "}
              <code>{candidate.token}</code>
            </span>
          ))}
        </span>
      ) : (
        <span className="lab1-confidence__tooltip-list">
          <span>No alternate tokens returned for this position.</span>
        </span>
      )}
    </span>
  );
}

export function Lab1ConfidenceChat() {
  const [draft, setDraft] = useState<string>(starterPrompts[0]);
  const [messages, setMessages] = useState<ChatTurn[]>([]);
  const [error, setError] = useState<string | null>(null);
  const [isSubmitting, setIsSubmitting] = useState(false);

  async function handleSubmit(event: FormEvent<HTMLFormElement>) {
    event.preventDefault();

    const prompt = draft.trim();
    if (!prompt) {
      setError("Enter a prompt to inspect the model output.");
      return;
    }

    const nextUserTurn: UserTurn = {
      content: prompt,
      id: buildTurnId(),
      role: "user",
    };
    const nextConversation = [...messages, nextUserTurn];

    setMessages(nextConversation);
    setDraft("");
    setError(null);
    setIsSubmitting(true);

    try {
      const response = await fetch("/api/lab1/chat", {
        body: JSON.stringify({
          messages: toConversation(nextConversation),
        }),
        headers: {
          "Content-Type": "application/json",
        },
        method: "POST",
      });

      const payload = (await response.json()) as Lab1ConfidenceResponse & {
        error?: string;
      };

      if (!response.ok) {
        throw new Error(payload.error || "The local Ollama request failed.");
      }

      setMessages((currentMessages) => [
        ...currentMessages,
        {
          ...payload,
          id: buildTurnId(),
        },
      ]);
    } catch (caughtError) {
      setError(
        caughtError instanceof Error
          ? caughtError.message
          : "The local Ollama request failed.",
      );
    } finally {
      setIsSubmitting(false);
    }
  }

  return (
    <section className="lab1-confidence" data-widget-enhanced="true">
      <div className="lab1-confidence__header">
        <p className="lab1-confidence__eyebrow">Lab 1 Confidence View</p>
        <h3>Visualize token confidence locally</h3>
        <p className="lab1-confidence__lede">
          This widget uses the preloaded Gemma 4 E2B Q4 model. Hover over any
          output token to inspect its probability and the strongest alternate
          predictions returned for that position.
        </p>
      </div>

      <div className="lab1-confidence__prompt-row">
        {starterPrompts.map((prompt) => (
          <button
            className="lab1-confidence__prompt-chip"
            key={prompt}
            onClick={() => setDraft(prompt)}
            type="button"
          >
            {prompt}
          </button>
        ))}
      </div>

      <div className="lab1-confidence__transcript" aria-live="polite">
        {messages.length === 0 ? (
          <div className="lab1-confidence__empty">
            <strong>Try a short prompt first.</strong>
            <p>
              Start with one of the suggested prompts, then hover across the
              model output to compare high-confidence and low-confidence tokens.
            </p>
          </div>
        ) : (
          messages.map((message) => {
            if (message.role === "user") {
              return (
                <article
                  className="lab1-confidence__message lab1-confidence__message--user"
                  key={message.id}
                >
                  <div className="lab1-confidence__message-meta">
                    <span>You</span>
                  </div>
                  <pre className="lab1-confidence__message-body">
                    <code>{message.content}</code>
                  </pre>
                </article>
              );
            }

            return (
              <article className="lab1-confidence__message" key={message.id}>
                <div className="lab1-confidence__message-meta">
                  <span>Assistant</span>
                  <code>{message.model}</code>
                </div>

                <div className="lab1-confidence__token-stream" role="list">
                  {message.tokens.map((token, index) => (
                    <span
                      aria-label={`${token.token} ${formatProbabilityPercent(
                        token.probability,
                      )}`}
                      className={`lab1-confidence__token lab1-confidence__token--${getConfidenceBand(
                        token.probability,
                      )}`}
                      key={`${message.id}-${index}-${token.token}`}
                      role="listitem"
                    >
                      {token.token}
                      {renderTooltip(token)}
                    </span>
                  ))}
                </div>

                {message.error ? (
                  <p className="lab1-confidence__message-warning">
                    {message.error}
                  </p>
                ) : null}
              </article>
            );
          })
        )}
      </div>

      <form className="lab1-confidence__composer" onSubmit={handleSubmit}>
        <label
          className="lab1-confidence__composer-label"
          htmlFor="lab1-confidence-draft"
        >
          Prompt
        </label>
        <textarea
          id="lab1-confidence-draft"
          name="draft"
          onChange={(event) => setDraft(event.target.value)}
          placeholder="Ask a question or start a phrase and inspect the output."
          rows={5}
          value={draft}
        />

        <div className="lab1-confidence__composer-actions">
          <div className="lab1-confidence__composer-state">
            <span>Inference target</span>
            <strong>Gemma 4 E2B Q4</strong>
          </div>
          <button disabled={isSubmitting} type="submit">
            {isSubmitting ? "Generating..." : "Generate Output"}
          </button>
        </div>

        {error ? <p className="lab1-confidence__error">{error}</p> : null}
      </form>
    </section>
  );
}
