import { render, Box, Text } from "ink";
import {
  AssistantRuntimeProvider,
  useLocalRuntime,
  ThreadPrimitive,
  ComposerPrimitive,
  useAuiState,
  type ChatModelAdapter,
} from "@assistant-ui/react-ink";

const adapter: ChatModelAdapter = {
  async *run({ messages }) {
    yield { content: [{ type: "text", text: "Hello from AI!" }] };
  },
};

const Message = () => {
  const message = useAuiState((s) => s.message);
  const text = message.content.find((p) => p.type === "text")?.text ?? "";
  return (
    <Box marginBottom={1}>
      <Text color={message.role === "user" ? "green" : "blue"}>
        {message.role === "user" ? "You: " : "AI: "}
      </Text>
      <Text>{text}</Text>
    </Box>
  );
};

const App = () => {
  const runtime = useLocalRuntime(adapter);
  return (
    <AssistantRuntimeProvider runtime={runtime}>
      <ThreadPrimitive.Root>
        <ThreadPrimitive.Messages>{() => <Message />}</ThreadPrimitive.Messages>
        <Box borderStyle="round" paddingX={1}>
          <Text>{"> "}</Text>
          <ComposerPrimitive.Input submitOnEnter placeholder="Message..." autoFocus />
        </Box>
      </ThreadPrimitive.Root>
    </AssistantRuntimeProvider>
  );
};

render(<App />);