src/app/components/ChatSurface.tsx

import { useEffect, useMemo, useRef, useState, type FormEvent, type KeyboardEvent } from 'react';
import type {
  ChatMessage,
  ComposerSubmission,
  ModelLoadProgress,
  ModelId,
  ModelOption,
  ShardRetryNotice,
  StreamState,
  ToolRun,
} from '../../lib/contracts';
import type { CodeTheme } from '../../lib/code-highlighter';
import { TOOL_CATALOG } from '../../lib/agent-tools';
import { segmentsForMessage } from '../../lib/message-segments';
import { ArtifactCard } from './ArtifactCard';
import { ComposerActivity } from './ComposerActivity';
import { MarkdownContent } from './MarkdownContent';

interface ChatSurfaceProps {
  activeModel: ModelOption;
  loadedModelId: ModelId | null;
  messages: ChatMessage[];
  streamState: StreamState;
  modelLoadProgress: ModelLoadProgress;
  toolsEnabled: boolean;
  runtimeStatus: string;
  error: string | null;
  shardRetry: ShardRetryNotice | null;
  theme: CodeTheme;
  onLoadModel: (modelId: ModelId) => void;
  onDismissError: () => void;
  onRetryShard: (modelId: ModelId) => void;
  onSend: (submission: ComposerSubmission) => void;
  onRegenerate: (messageId: string) => void;
  onEditUserMessage: (messageId: string, content: string) => void;
  onCancel: () => void;
  onToolsEnabledChange: (enabled: boolean) => void;
  onOpenConversationList: () => void;
  onOpenSettings: () => void;
}

export function ChatSurface({
  activeModel,
  loadedModelId,
  messages,
  streamState,
  modelLoadProgress,
  toolsEnabled,
  runtimeStatus,
  error,
  shardRetry,
  theme,
  onLoadModel,
  onDismissError,
  onRetryShard,
  onSend,
  onRegenerate,
  onEditUserMessage,
  onCancel,
  onToolsEnabledChange,
  onOpenConversationList,
  onOpenSettings,
}: ChatSurfaceProps) {
  const [prompt, setPrompt] = useState('');
  const isLoading = streamState === 'loading';
  const isStreaming = streamState === 'streaming';
  const isBusy = isStreaming || isLoading;
  const isLoaded = loadedModelId === activeModel.id;
  const runtimeContextLimit = activeModel.id === 'bonsai-27b'
    ? activeModel.defaultContext
    : activeModel.contextLimit;
  const logDate = useMemo(() => new Intl.DateTimeFormat('en-GB', {
    day: '2-digit', month: 'long', year: 'numeric',
  }).format(new Date()), []);

  const submit = () => {
    const trimmed = prompt.trim();
    if (!trimmed || isBusy || !isLoaded) return;
    onSend({ prompt: trimmed, toolsEnabled });
    setPrompt('');
  };

  const handleSubmit = (event: FormEvent<HTMLFormElement>) => {
    event.preventDefault();
    submit();
  };

  const handleKeyDown = (event: KeyboardEvent<HTMLTextAreaElement>) => {
    if (event.key === 'Enter' && !event.shiftKey) {
      event.preventDefault();
      submit();
    }
  };

  return (
    <main className="chat-surface">
      <header className="chat-header">
        <div className="chat-header-main">
          <button className="chat-list-button" type="button" disabled={isBusy} onClick={onOpenConversationList} aria-label="Open chat list">
            <BackIcon /> Chats
          </button>
          <button className="chat-settings-button" type="button" disabled={isBusy} onClick={onOpenSettings}>Chat settings</button>
          <div className="chat-title">
            <h2>{activeModel.label}</h2>
            <span className="architecture-subtitle">{activeModel.architectureLabel}</span>
          </div>
        </div>
        <div className="chat-header-meta">
          <span>{activeModel.footprint} download</span>
          <span>{Math.round(runtimeContextLimit / 1024)}K max context</span>
          <span className={`live-label${isLoaded ? '' : ' is-dormant'}`}><i /> {isLoaded ? 'loaded' : 'not loaded'}</span>
        </div>
      </header>

      {shardRetry && (
        <div className="runtime-error download-retry" role="alert">
          <div>
            <strong>Shard {shardRetry.shardIndex + 1}/{shardRetry.shardCount} failed</strong>
            <span>{shardRetry.shardPath} · {shardRetry.error}</span>
          </div>
          <button type="button" disabled={isBusy} onClick={() => onRetryShard(shardRetry.modelId)}>
            Retry shard from byte zero
          </button>
        </div>
      )}

      {error && !shardRetry && (
        <div className="runtime-error" role="alert">
          <div><strong>Runtime stopped safely</strong><span>{error}</span></div>
          <button type="button" onClick={onDismissError}>Dismiss</button>
        </div>
      )}

      <section className="transcript" aria-label="Conversation transcript" aria-live="polite">
        <div className="transcript-date"><span>{logDate}</span></div>
        {messages.length === 0 && (
          <div className="transcript-empty">
            <h3>Start a chat</h3>
            <p>Choose a model, keep the backend on Auto, and send a prompt. The model is downloaded once and runs inside this browser.</p>
            <button className="system-instructions-button" type="button" onClick={onOpenSettings}>Set system instructions</button>
          </div>
        )}
        {messages.map((message) => (
          <Message
            key={message.id}
            message={message}
            theme={theme}
            regenerateDisabled={isBusy || !isLoaded}
            onRegenerate={onRegenerate}
            onEditUserMessage={onEditUserMessage}
          />
        ))}
      </section>

      <form className="composer" onSubmit={handleSubmit}>
        <ComposerActivity
          modelId={activeModel.id}
          streamState={streamState}
          progress={modelLoadProgress}
        />
        <div className="composer-meta">
          <div className="tool-controls">
            <button
              className={`tool-toggle${toolsEnabled ? ' is-on' : ''}`}
              type="button"
              aria-pressed={toolsEnabled}
              onClick={() => onToolsEnabledChange(!toolsEnabled)}
            >
              <ToolIcon /> Tools {toolsEnabled ? 'on' : 'off'}
            </button>
            <div className="tool-info">
              <button type="button" aria-label="Show available tools" aria-describedby="tool-reference" title="Available tools"><InfoIcon /></button>
              <div className="tool-popover" id="tool-reference" role="tooltip">
                <header>
                  <strong>Available tools</strong>
                  <span>{TOOL_CATALOG.length} capabilities</span>
                </header>
                <ul>
                  {TOOL_CATALOG.map((tool) => (
                    <li key={tool.name}>
                      <div><code>{tool.name}</code><small>{tool.scope}</small></div>
                      <p>{tool.description}</p>
                    </li>
                  ))}
                </ul>
                <footer>Inference and local tools stay in this browser. Network tools ask before sending a query.</footer>
              </div>
            </div>
          </div>
          <span>{runtimeStatus}</span>
        </div>
        <div className="composer-field">
          <label htmlFor="bonsai-prompt">Prompt</label>
          <textarea
            id="bonsai-prompt"
            aria-label="Ask Bonsai"
            value={prompt}
            onChange={(event) => setPrompt(event.target.value)}
            onKeyDown={handleKeyDown}
            placeholder={isLoaded ? 'Ask anything, or request an HTML artifact…' : 'Write while the model loads, then run locally…'}
            rows={2}
            disabled={isStreaming}
          />
          {isBusy ? (
            <button className="send-button stop" type="button" onClick={onCancel}>
              <span aria-hidden="true" /> {isLoading ? 'Stop load' : 'Stop'}
            </button>
          ) : !isLoaded ? (
            <button className="send-button load" type="button" onClick={() => onLoadModel(activeModel.id)}>
              Load model <DownloadIcon />
            </button>
          ) : (
            <button className="send-button" type="submit" disabled={!prompt.trim()}>
              Run <ArrowIcon />
            </button>
          )}
        </div>
      </form>
    </main>
  );
}

interface MessageProps {
  message: ChatMessage;
  theme: CodeTheme;
  regenerateDisabled: boolean;
  onRegenerate: (messageId: string) => void;
  onEditUserMessage: (messageId: string, content: string) => void;
}

function Message({ message, theme, regenerateDisabled, onRegenerate, onEditUserMessage }: MessageProps) {
  const streaming = message.state === 'streaming';
  const [copyState, setCopyState] = useState<'idle' | 'copied' | 'error'>('idle');
  const [editing, setEditing] = useState(false);
  const [editText, setEditText] = useState(message.content);
  const copyResetRef = useRef<number | null>(null);

  useEffect(() => () => {
    if (copyResetRef.current !== null) window.clearTimeout(copyResetRef.current);
  }, []);

  const scheduleCopyReset = () => {
    if (copyResetRef.current !== null) window.clearTimeout(copyResetRef.current);
    copyResetRef.current = window.setTimeout(() => setCopyState('idle'), 1_600);
  };

  const copyMessage = () => {
    if (!message.content || !navigator.clipboard?.writeText) {
      setCopyState('error');
      scheduleCopyReset();
      return;
    }
    void navigator.clipboard.writeText(message.content)
      .then(() => setCopyState('copied'))
      .catch(() => setCopyState('error'))
      .finally(scheduleCopyReset);
  };

  const regenerateLabel = message.role === 'user'
    ? 'Regenerate from this message'
    : 'Regenerate this response';
  const segments = message.role === 'assistant' ? segmentsForMessage(message) : [];
  const tools = new Map((message.tools ?? []).map((tool) => [tool.id, tool]));
  const artifacts = new Map((message.artifacts ?? []).map((artifact) => [artifact.id, artifact]));
  const lastToolForArtifact = new Map<string, number>();
  segments.forEach((segment, index) => {
    if (segment.kind !== 'tool') return;
    const artifactId = tools.get(segment.toolCallId)?.artifactId;
    if (artifactId && artifacts.has(artifactId)) lastToolForArtifact.set(artifactId, index);
  });
  const lastReasoningId = [...segments].reverse().find((segment) => segment.kind === 'reasoning')?.id;
  const saveEdit = () => {
    const normalized = editText.trim();
    if (!normalized || normalized === message.content.trim() || regenerateDisabled) return;
    setEditing(false);
    onEditUserMessage(message.id, normalized);
  };

  return (
    <article className={`message message-${message.role}${streaming ? ' is-streaming' : ''}${message.error ? ' has-error' : ''}${message.artifacts?.length ? ' has-artifacts' : ''}`}>
      <header>
        <strong>{message.role === 'assistant' ? 'Bonsai' : message.role === 'system' ? 'System' : 'You'}</strong>
        {message.editedAt && <span className="message-edited">edited</span>}
        <time>{message.timestamp}</time>
      </header>
      <div className="message-content">
        {message.role === 'user' && editing ? (
          <form className="message-edit-form" onSubmit={(event) => { event.preventDefault(); saveEdit(); }}>
            <label htmlFor={`edit-${message.id}`}>Edit message</label>
            <textarea
              id={`edit-${message.id}`}
              aria-label="Edit message"
              value={editText}
              rows={Math.min(10, Math.max(3, editText.split('\n').length + 1))}
              autoFocus
              onChange={(event) => setEditText(event.target.value)}
              onKeyDown={(event) => {
                if (event.key === 'Escape') {
                  event.preventDefault();
                  setEditText(message.content);
                  setEditing(false);
                }
                if (event.key === 'Enter' && (event.metaKey || event.ctrlKey)) {
                  event.preventDefault();
                  saveEdit();
                }
              }}
            />
            <div>
              <button type="button" onClick={() => { setEditText(message.content); setEditing(false); }}>Cancel</button>
              <button type="submit" disabled={!editText.trim() || editText.trim() === message.content.trim()}>Save &amp; regenerate</button>
            </div>
          </form>
        ) : message.role === 'assistant' ? segments.map((segment, index) => {
          if (segment.kind === 'reasoning') {
            const active = message.isThinking && segment.id === lastReasoningId;
            return (
              <div className="message-segment message-segment-reasoning" key={segment.id}>
                <details className="reasoning-block" open={active}>
                  <summary>{active ? 'Reasoning in progress…' : 'Reasoning trace'}</summary>
                  <MarkdownContent content={segment.content} theme={theme} streaming={streaming && active} />
                </details>
              </div>
            );
          }
          if (segment.kind === 'text') {
            return <div className="message-segment message-segment-text" key={segment.id}><MarkdownContent content={segment.content} theme={theme} streaming={streaming} /></div>;
          }
          if (segment.kind === 'artifact') {
            const artifact = artifacts.get(segment.artifactId);
            return artifact ? <div className="message-segment message-segment-artifact" key={segment.id}><ArtifactCard artifact={artifact} /></div> : null;
          }
          const tool = tools.get(segment.toolCallId);
          if (!tool) return null;
          const artifact = tool.artifactId && lastToolForArtifact.get(tool.artifactId) === index
            ? artifacts.get(tool.artifactId)
            : undefined;
          return (
            <div className="message-segment message-segment-tool" key={segment.id}>
              <ToolCall tool={tool} />
              {artifact && <ArtifactCard artifact={artifact} />}
            </div>
          );
        }) : message.content ? <MarkdownContent content={message.content} theme={theme} streaming={streaming} /> : null}
        {streaming && message.toolActivity && (
          <p className="tool-generation" aria-label={`Generating ${message.toolActivity.name} tool call`}>
            <span>Calling</span>
            <code>{message.toolActivity.name}</code>
            <span>· {message.toolActivity.argumentCharacters.toLocaleString('en-US')} chars</span>
            <span className="stream-cursor" aria-hidden="true" />
          </p>
        )}
        {segments.length === 0 && !message.content && streaming && !message.toolActivity && (
          <p className="generation-placeholder">
            Generating… <span className="stream-cursor" aria-label="Streaming response" />
          </p>
        )}
        {(segments.length > 0 || message.content) && streaming && !message.toolActivity && (
          <span className="stream-cursor" aria-label="Streaming response" />
        )}
        {message.error && <p className="message-error">{message.error}</p>}
      </div>
      {message.metrics && (
        <footer>
          <span>{message.metrics.tokens} tok</span>
          {message.metrics.promptTokensPerSecond !== undefined && (
            <span>Prefill {message.metrics.promptTokensPerSecond.toFixed(1)} tok/s</span>
          )}
          <span>Decode {message.metrics.tokensPerSecond.toFixed(1)} tok/s</span>
          <span>TTFT {message.metrics.timeToFirstTokenMs} ms</span>
        </footer>
      )}
      {!editing && <div className="message-actions" aria-label="Message actions">
        <button
          type="button"
          disabled={!message.content}
          onClick={copyMessage}
          aria-label={copyState === 'copied' ? 'Message copied' : 'Copy message'}
        >
          <CopyIcon />
          {copyState === 'copied' ? 'Copied' : copyState === 'error' ? 'Copy failed' : 'Copy'}
        </button>
        {message.role === 'user' && (
          <button
            type="button"
            disabled={regenerateDisabled}
            onClick={() => { setEditText(message.content); setEditing(true); }}
            aria-label="Edit message and regenerate response"
            title="Edit this message, delete later turns, and generate a new response"
          >
            <EditIcon /> Edit
          </button>
        )}
        {message.role !== 'system' && (
          <button
            type="button"
            disabled={regenerateDisabled}
            onClick={() => onRegenerate(message.id)}
            aria-label={regenerateLabel}
            title={message.role === 'user'
              ? 'Delete later turns and answer this message again'
              : 'Delete this response and later turns, then generate it again'}
          >
            <RegenerateIcon /> Regenerate
          </button>
        )}
      </div>}
    </article>
  );
}

function ToolCall({ tool }: { tool: ToolRun }) {
  const output = parseToolOutput(tool.output);
  return (
    <section className={`tool-call tool-call-${tool.state}`}>
      <header>
        <span className={`tool-state tool-state-${tool.state}`} />
        <code>{tool.name}</code>
        <span>{tool.state}</span>
      </header>
      <ToolResult tool={tool} output={output} />
      <details className="tool-call-raw">
        <summary>Arguments and raw result</summary>
        <div>
          <strong>Arguments</strong>
          <pre>{prettyJson(tool.input)}</pre>
          {tool.output && <><strong>Result</strong><pre>{prettyJson(tool.output)}</pre></>}
        </div>
      </details>
    </section>
  );
}

function ToolResult({ tool, output }: { tool: ToolRun; output: Record<string, unknown> | null }) {
  if (tool.state === 'running' || tool.state === 'queued') {
    return <p className="tool-result-status">Running locally…</p>;
  }
  if (!output) return tool.output ? <p className="tool-result-status">Result returned</p> : null;
  if (output.ok === false) {
    return <p className="tool-result-status is-error">{String(output.error ?? 'Tool failed')}</p>;
  }
  if (tool.name === 'js_eval' && 'value' in output) {
    const logs = Array.isArray(output.logs) ? output.logs.map(String) : [];
    return <div className="tool-value-result"><span>Result</span><code>{String(output.value)}</code>{logs.length > 0 && <small>{logs.join('\n')}</small>}</div>;
  }
  if (tool.name === 'web_search' && Array.isArray(output.results)) {
    return (
      <ol className="tool-search-results">
        {output.results.slice(0, 8).map((item, index) => {
          const result = isRecord(item) ? item : {};
          const url = safeHttpUrl(result.url);
          return (
            <li key={`${String(result.title ?? 'Result')}-${index}`}>
              {url ? <a href={url} target="_blank" rel="noreferrer">{String(result.title ?? url)}</a> : <strong>{String(result.title ?? 'Result')}</strong>}
              {Boolean(result.snippet) && <p>{String(result.snippet)}</p>}
            </li>
          );
        })}
      </ol>
    );
  }
  const evaluation = isRecord(output.evaluation) ? output.evaluation : null;
  if (evaluation) {
    const errors = Array.isArray(evaluation.errors) ? evaluation.errors.length : 0;
    return <p className={`tool-result-status${evaluation.status === 'passed' ? ' is-passed' : ' is-error'}`}>Artifact test {String(evaluation.status ?? 'complete')}{errors > 0 ? ` · ${errors} ${errors === 1 ? 'error' : 'errors'}` : ''}</p>;
  }
  if (typeof output.action === 'string') {
    return <p className="tool-result-status is-passed">Memory {output.action} complete</p>;
  }
  return <p className="tool-result-status is-passed">Result returned</p>;
}

function parseToolOutput(value: string | undefined): Record<string, unknown> | null {
  if (!value) return null;
  try {
    const parsed: unknown = JSON.parse(value);
    return isRecord(parsed) ? parsed : null;
  } catch {
    return null;
  }
}

function prettyJson(value: string): string {
  try { return JSON.stringify(JSON.parse(value), null, 2); } catch { return value; }
}

function isRecord(value: unknown): value is Record<string, unknown> {
  return typeof value === 'object' && value !== null && !Array.isArray(value);
}

function safeHttpUrl(value: unknown): string | undefined {
  if (typeof value !== 'string') return undefined;
  try {
    const url = new URL(value);
    return url.protocol === 'http:' || url.protocol === 'https:' ? url.href : undefined;
  } catch {
    return undefined;
  }
}

function ToolIcon() {
  return (
    <svg viewBox="0 0 20 20" aria-hidden="true">
      <path d="m12.7 3.2-2.4 2.4 4.1 4.1 2.4-2.4a5 5 0 0 1-6.4 6.4l-5 5-3.1-3.1 5-5a5 5 0 0 1 5.4-7.4Z" />
    </svg>
  );
}

function InfoIcon() {
  return (
    <svg viewBox="0 0 20 20" aria-hidden="true">
      <circle cx="10" cy="10" r="7" />
      <path d="M10 8.4v5M10 5.8h.01" />
    </svg>
  );
}

function CopyIcon() {
  return (
    <svg viewBox="0 0 20 20" aria-hidden="true">
      <rect x="6.5" y="6.5" width="9" height="9" rx="1.5" />
      <path d="M13.5 6.5V5A1.5 1.5 0 0 0 12 3.5H5A1.5 1.5 0 0 0 3.5 5v7A1.5 1.5 0 0 0 5 13.5h1.5" />
    </svg>
  );
}

function RegenerateIcon() {
  return (
    <svg viewBox="0 0 20 20" aria-hidden="true">
      <path d="M15.7 7.2A6.3 6.3 0 1 0 16 12M15.7 7.2V3.8m0 3.4h-3.4" />
    </svg>
  );
}

function EditIcon() {
  return (
    <svg viewBox="0 0 20 20" aria-hidden="true">
      <path d="m4 14.8.5-3.4 7.8-7.8a1.5 1.5 0 0 1 2.1 0l2 2a1.5 1.5 0 0 1 0 2.1l-7.8 7.8-3.4.5L4 14.8Z" />
      <path d="m11.3 4.6 4.1 4.1" />
    </svg>
  );
}

function BackIcon() {
  return <svg viewBox="0 0 20 20" aria-hidden="true"><path d="m11.5 5-5 5 5 5M7 10h7" /></svg>;
}

function ArrowIcon() {
  return (
    <svg viewBox="0 0 20 20" aria-hidden="true">
      <path d="M3 10h13M11 5l5 5-5 5" />
    </svg>
  );
}

function DownloadIcon() {
  return (
    <svg viewBox="0 0 20 20" aria-hidden="true">
      <path d="M10 2v11m-4-4 4 4 4-4M3 17h14" />
    </svg>
  );
}