storkit/src/components/Chat.tsx

import { useState, useRef, useEffect } from "react";
import { invoke } from "@tauri-apps/api/core";
import { listen } from "@tauri-apps/api/event";
import Markdown from "react-markdown";
import { Prism as SyntaxHighlighter } from "react-syntax-highlighter";
import { oneDark } from "react-syntax-highlighter/dist/esm/styles/prism";
import { Message, ProviderConfig } from "../types";

interface ChatProps {
  projectPath: string;
  onCloseProject: () => void;
}

export function Chat({ projectPath, onCloseProject }: ChatProps) {
  const [messages, setMessages] = useState<Message[]>([]);
  const [input, setInput] = useState("");
  const [loading, setLoading] = useState(false);
  const [model, setModel] = useState("llama3.1"); // Default local model
  const [enableTools, setEnableTools] = useState(true);
  const [availableModels, setAvailableModels] = useState<string[]>([]);
  const messagesEndRef = useRef<HTMLDivElement>(null);

  useEffect(() => {
    invoke<string[]>("get_ollama_models")
      .then(async (models) => {
        if (models.length > 0) {
          setAvailableModels(models);

          // Check backend store for saved model
          try {
            const savedModel = await invoke<string | null>(
              "get_model_preference",
            );
            if (savedModel && models.includes(savedModel)) {
              setModel(savedModel);
            } else if (!models.includes(model)) {
              setModel(models[0]);
            }
          } catch (e) {
            console.error(e);
          }
        }
      })
      .catch((err) => console.error(err));
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, []);

  useEffect(() => {
    const unlistenPromise = listen<Message[]>("chat:update", (event) => {
      setMessages(event.payload);
    });

    return () => {
      unlistenPromise.then((unlisten) => unlisten());
    };
  }, []);

  const scrollToBottom = () => {
    messagesEndRef.current?.scrollIntoView({ behavior: "smooth" });
  };

  useEffect(scrollToBottom, [messages]);

  const sendMessage = async () => {
    if (!input.trim() || loading) return;

    const userMsg: Message = { role: "user", content: input };
    const newHistory = [...messages, userMsg];

    setMessages(newHistory);
    setInput("");
    setLoading(true);

    try {
      const config: ProviderConfig = {
        provider: "ollama",
        model: model,
        base_url: "http://localhost:11434",
        enable_tools: enableTools,
      };

      // Invoke backend chat command
      // We rely on 'chat:update' events to update the state in real-time
      await invoke("chat", {
        messages: newHistory,
        config: config,
      });
    } catch (e) {
      console.error(e);
      setMessages((prev) => [
        ...prev,
        { role: "assistant", content: `**Error:** ${e}` },
      ]);
    } finally {
      setLoading(false);
    }
  };

  return (
    <div
      className="chat-container"
      style={{
        display: "flex",
        flexDirection: "column",
        height: "100%",
        backgroundColor: "#171717",
        color: "#ececec",
      }}
    >
      {/* Sticky Header */}
      <div
        style={{
          padding: "12px 24px",
          borderBottom: "1px solid #333",
          display: "flex",
          alignItems: "center",
          justifyContent: "space-between",
          background: "#171717",
          flexShrink: 0,
          fontSize: "0.9rem",
          color: "#ececec",
        }}
      >
        {/* Project Info */}
        <div
          style={{
            display: "flex",
            alignItems: "center",
            gap: "12px",
            overflow: "hidden",
            flex: 1,
            marginRight: "20px",
          }}
        >
          <div
            title={projectPath}
            style={{
              whiteSpace: "nowrap",
              overflow: "hidden",
              textOverflow: "ellipsis",
              fontWeight: "500",
              color: "#aaa",
              direction: "rtl",
              textAlign: "left",
              fontFamily: "monospace",
              fontSize: "0.85em",
            }}
          >
            {projectPath}
          </div>
          <button
            onClick={onCloseProject}
            style={{
              background: "transparent",
              border: "none",
              cursor: "pointer",
              color: "#999",
              fontSize: "0.8em",
              padding: "4px 8px",
              borderRadius: "4px",
            }}
            onMouseOver={(e) => (e.currentTarget.style.background = "#333")}
            onMouseOut={(e) =>
              (e.currentTarget.style.background = "transparent")
            }
          >
            ✕
          </button>
        </div>

        {/* Model Controls */}
        <div style={{ display: "flex", alignItems: "center", gap: "16px" }}>
          {availableModels.length > 0 ? (
            <select
              value={model}
              onChange={(e) => {
                const newModel = e.target.value;
                setModel(newModel);
                invoke("set_model_preference", { model: newModel }).catch(
                  console.error,
                );
              }}
              style={{
                padding: "6px 32px 6px 16px",
                borderRadius: "99px",
                border: "none",
                fontSize: "0.9em",
                backgroundColor: "#2f2f2f",
                color: "#ececec",
                cursor: "pointer",
                outline: "none",
                appearance: "none",
                WebkitAppearance: "none",
                backgroundImage: `url("data:image/svg+xml;charset=US-ASCII,%3Csvg%20xmlns%3D%22http%3A%2F%2Fwww.w3.org%2F2000%2Fsvg%22%20width%3D%22292.4%22%20height%3D%22292.4%22%3E%3Cpath%20fill%3D%22%23ececec%22%20d%3D%22M287%2069.4a17.6%2017.6%200%200%200-13-5.4H18.4c-5%200-9.3%201.8-12.9%205.4A17.6%2017.6%200%200%200%200%2082.2c0%205%201.8%209.3%205.4%2012.9l128%20127.9c3.6%203.6%207.8%205.4%2012.8%205.4s9.2-1.8%2012.8-5.4L287%2095c3.5-3.5%205.4-7.8%205.4-12.8%200-5-1.9-9.2-5.5-12.8z%22%2F%3E%3C%2Fsvg%3E")`,
                backgroundRepeat: "no-repeat",
                backgroundPosition: "right 12px center",
                backgroundSize: "10px",
              }}
            >
              {availableModels.map((m) => (
                <option key={m} value={m}>
                  {m}
                </option>
              ))}
            </select>
          ) : (
            <input
              value={model}
              onChange={(e) => {
                const newModel = e.target.value;
                setModel(newModel);
                invoke("set_model_preference", { model: newModel }).catch(
                  console.error,
                );
              }}
              placeholder="Model"
              style={{
                padding: "6px 12px",
                borderRadius: "99px",
                border: "none",
                fontSize: "0.9em",
                background: "#2f2f2f",
                color: "#ececec",
                outline: "none",
              }}
            />
          )}
          <label
            style={{
              display: "flex",
              alignItems: "center",
              gap: "6px",
              cursor: "pointer",
              fontSize: "0.9em",
              color: "#aaa",
            }}
            title="Allow the Agent to read/write files"
          >
            <input
              type="checkbox"
              checked={enableTools}
              onChange={(e) => setEnableTools(e.target.checked)}
              style={{ accentColor: "#000" }}
            />
            <span>Tools</span>
          </label>
        </div>
      </div>

      {/* Messages Area */}
      <div
        style={{
          flex: 1,
          overflowY: "auto",
          padding: "20px 0",
          display: "flex",
          flexDirection: "column",
          gap: "24px",
        }}
      >
        <div
          style={{
            maxWidth: "768px",
            margin: "0 auto",
            width: "100%",
            padding: "0 24px",
            display: "flex",
            flexDirection: "column",
            gap: "24px",
          }}
        >
          {messages.map((msg, idx) => (
            <div
              key={idx}
              style={{
                display: "flex",
                flexDirection: "column",
                alignItems: msg.role === "user" ? "flex-end" : "flex-start",
              }}
            >
              <div
                style={{
                  maxWidth: "100%",
                  padding: msg.role === "user" ? "10px 16px" : "0",
                  borderRadius: msg.role === "user" ? "20px" : "0",
                  background:
                    msg.role === "user"
                      ? "#2f2f2f"
                      : msg.role === "tool"
                        ? "#222"
                        : "transparent",
                  color: "#ececec",
                  border: msg.role === "tool" ? "1px solid #333" : "none",
                  fontFamily: msg.role === "tool" ? "monospace" : "inherit",
                  fontSize: msg.role === "tool" ? "0.85em" : "1em",
                  fontWeight: "500",
                  whiteSpace: msg.role === "tool" ? "pre-wrap" : "normal",
                  lineHeight: "1.6",
                }}
              >
                {msg.role === "user" ? (
                  msg.content
                ) : msg.role === "tool" ? (
                  <details style={{ cursor: "pointer" }}>
                    <summary
                      style={{
                        color: "#aaa",
                        fontSize: "0.9em",
                        marginBottom: "8px",
                        listStyle: "none",
                        display: "flex",
                        alignItems: "center",
                        gap: "6px",
                      }}
                    >
                      <span style={{ fontSize: "0.8em" }}>▶</span>
                      <span>
                        Tool Output
                        {msg.tool_call_id && ` (${msg.tool_call_id})`}
                      </span>
                    </summary>
                    <pre
                      style={{
                        maxHeight: "300px",
                        overflow: "auto",
                        margin: 0,
                        padding: "8px",
                        background: "#1a1a1a",
                        borderRadius: "4px",
                        fontSize: "0.85em",
                        whiteSpace: "pre-wrap",
                        wordBreak: "break-word",
                      }}
                    >
                      {msg.content}
                    </pre>
                  </details>
                ) : (
                  <div className="markdown-body">
                    <Markdown
                      components={{
                        code: ({ className, children, ...props }: any) => {
                          const match = /language-(\w+)/.exec(className || "");
                          const isInline = !className;
                          return !isInline && match ? (
                            <SyntaxHighlighter
                              style={oneDark as any}
                              language={match[1]}
                              PreTag="div"
                            >
                              {String(children).replace(/\n$/, "")}
                            </SyntaxHighlighter>
                          ) : (
                            <code className={className} {...props}>
                              {children}
                            </code>
                          );
                        },
                      }}
                    >
                      {msg.content}
                    </Markdown>
                  </div>
                )}

                {/* Show Tool Calls if present */}
                {msg.tool_calls && (
                  <div
                    style={{
                      marginTop: "12px",
                      fontSize: "0.85em",
                      color: "#aaa",
                      display: "flex",
                      flexDirection: "column",
                      gap: "8px",
                    }}
                  >
                    {msg.tool_calls.map((tc, i) => {
                      // Parse arguments to extract key info
                      let argsSummary = "";
                      try {
                        const args = JSON.parse(tc.function.arguments);
                        const firstKey = Object.keys(args)[0];
                        if (firstKey && args[firstKey]) {
                          argsSummary = String(args[firstKey]);
                          // Truncate if too long
                          if (argsSummary.length > 50) {
                            argsSummary = argsSummary.substring(0, 47) + "...";
                          }
                        }
                      } catch (e) {
                        // If parsing fails, just show empty
                      }

                      return (
                        <div
                          key={i}
                          style={{
                            display: "flex",
                            alignItems: "center",
                            gap: "8px",
                            fontFamily: "monospace",
                          }}
                        >
                          <span style={{ color: "#888" }}>▶</span>
                          <span
                            style={{
                              background: "#333",
                              padding: "2px 6px",
                              borderRadius: "4px",
                            }}
                          >
                            {tc.function.name}
                            {argsSummary && `(${argsSummary})`}
                          </span>
                        </div>
                      );
                    })}
                  </div>
                )}
              </div>
            </div>
          ))}
          {loading && (
            <div
              style={{
                alignSelf: "flex-start",
                color: "#888",
                fontSize: "0.9em",
                marginTop: "10px",
              }}
            >
              <span className="pulse">Thinking...</span>
            </div>
          )}
          <div ref={messagesEndRef} />
        </div>
      </div>

      {/* Input Area */}
      <div
        style={{
          padding: "24px",
          background: "#171717",
          display: "flex",
          justifyContent: "center",
        }}
      >
        <div
          style={{
            maxWidth: "768px",
            width: "100%",
            position: "relative",
          }}
        >
          <input
            value={input}
            onChange={(e) => setInput(e.target.value)}
            onKeyDown={(e) => e.key === "Enter" && sendMessage()}
            placeholder="Send a message..."
            style={{
              width: "100%",
              padding: "14px 20px",
              paddingRight: "50px", // space for button
              borderRadius: "24px",
              border: "1px solid #333",
              outline: "none",
              fontSize: "1rem",
              fontWeight: "500",
              background: "#2f2f2f",
              color: "#ececec",
              boxShadow: "0 2px 6px rgba(0,0,0,0.02)",
            }}
          />
          <button
            onClick={sendMessage}
            disabled={loading}
            style={{
              position: "absolute",
              right: "8px",
              top: "50%",
              transform: "translateY(-50%)",
              background: "#ececec",
              color: "black",
              border: "none",
              borderRadius: "50%",
              width: "32px",
              height: "32px",
              display: "flex",
              alignItems: "center",
              justifyContent: "center",
              cursor: loading ? "not-allowed" : "pointer",
              opacity: loading ? 0.5 : 1,
            }}
          >
            ↑
          </button>
        </div>
      </div>
    </div>
  );
}