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 { Message, ProviderConfig } from "../types";

export function Chat() {
  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((models) => {
        if (models.length > 0) {
          setAvailableModels(models);
          // If we have models and current one isn't valid, switch to first
          if (!models.includes(model)) {
            setModel(models[0]);
          }
        }
      })
      .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%",
        maxWidth: "800px",
        margin: "0 auto",
      }}
    >
      {/* Settings Bar */}
      <div
        style={{
          padding: "10px",
          borderBottom: "1px solid #ddd",
          display: "flex",
          gap: "10px",
          alignItems: "center",
        }}
      >
        <label>Ollama Model:</label>
        {availableModels.length > 0 ? (
          <select
            value={model}
            onChange={(e) => setModel(e.target.value)}
            style={{ padding: "5px" }}
          >
            {availableModels.map((m) => (
              <option key={m} value={m}>
                {m}
              </option>
            ))}
          </select>
        ) : (
          <input
            value={model}
            onChange={(e) => setModel(e.target.value)}
            placeholder="e.g. llama3, mistral"
            style={{ padding: "5px" }}
          />
        )}
        <label
          style={{
            display: "flex",
            alignItems: "center",
            gap: "5px",
            marginLeft: "10px",
          }}
        >
          <input
            type="checkbox"
            checked={enableTools}
            onChange={(e) => setEnableTools(e.target.checked)}
          />
          Enable Tools
        </label>
      </div>

      {/* Messages Area */}
      <div
        style={{
          flex: 1,
          overflowY: "auto",
          padding: "20px",
          display: "flex",
          flexDirection: "column",
          gap: "15px",
        }}
      >
        {messages.map((msg, idx) => (
          <div
            key={idx}
            className={`message ${msg.role}`}
            style={{
              alignSelf: msg.role === "user" ? "flex-end" : "flex-start",
              maxWidth: "80%",
              padding: "10px 15px",
              borderRadius: "10px",
              background:
                msg.role === "user"
                  ? "#007AFF"
                  : msg.role === "tool"
                    ? "#f0f0f0"
                    : "#E5E5EA",
              color: msg.role === "user" ? "white" : "black",
              border: msg.role === "tool" ? "1px solid #ccc" : "none",
              fontFamily: msg.role === "tool" ? "monospace" : "inherit",
              fontSize: msg.role === "tool" ? "0.9em" : "1em",
              whiteSpace: msg.role === "tool" ? "pre-wrap" : "normal",
            }}
          >
            <strong>
              {msg.role === "user"
                ? "You"
                : msg.role === "tool"
                  ? "Tool Output"
                  : "Agent"}
            </strong>
            {msg.role === "tool" ? (
              <div style={{ maxHeight: "200px", overflow: "auto" }}>
                {msg.content}
              </div>
            ) : (
              <Markdown>{msg.content}</Markdown>
            )}

            {/* Show Tool Calls if present */}
            {msg.tool_calls && (
              <div
                style={{ marginTop: "10px", fontSize: "0.85em", color: "#666" }}
              >
                {msg.tool_calls.map((tc, i) => (
                  <div
                    key={i}
                    style={{
                      background: "rgba(0,0,0,0.05)",
                      padding: "5px",
                      borderRadius: "4px",
                    }}
                  >
                    🛠{" "}
                    <code>
                      {tc.function.name}({tc.function.arguments})
                    </code>
                  </div>
                ))}
              </div>
            )}
          </div>
        ))}
        {loading && (
          <div style={{ alignSelf: "flex-start", color: "#888" }}>
            Thinking...
          </div>
        )}
        <div ref={messagesEndRef} />
      </div>

      {/* Input Area */}
      <div
        style={{
          padding: "20px",
          borderTop: "1px solid #ddd",
          display: "flex",
          gap: "10px",
        }}
      >
        <input
          value={input}
          onChange={(e) => setInput(e.target.value)}
          onKeyDown={(e) => e.key === "Enter" && sendMessage()}
          placeholder="Ask the agent to do something..."
          style={{
            flex: 1,
            padding: "10px",
            borderRadius: "4px",
            border: "1px solid #ccc",
          }}
        />
        <button
          onClick={sendMessage}
          disabled={loading}
          style={{ padding: "10px 20px" }}
        >
          Send
        </button>
      </div>
    </div>
  );
}