import { useCallback, useState } from "react";
import { AgentAction } from "../components/AgentAction";
import { ChatInput } from "../components/ChatInput";
import { ChatMessages } from "../components/ChatMessages";
import { InterruptPrompt } from "../components/InterruptPrompt";
import { useWebSocket } from "../hooks/useWebSocket";
import type {
  ChatMessage,
  ConnectionStatus,
  InterruptMessage,
  ServerMessage,
  ToolAction,
} from "../types";

let msgCounter = 0;
function nextId(): string {
  msgCounter += 1;
  return `msg-${msgCounter}`;
}

export function ChatPage() {
  const [messages, setMessages] = useState<ChatMessage[]>([]);
  const [toolActions, setToolActions] = useState<ToolAction[]>([]);
  const [currentInterrupt, setCurrentInterrupt] = useState<InterruptMessage | null>(null);
  const [isWaiting, setIsWaiting] = useState(false);

  const handleServerMessage = useCallback((msg: ServerMessage) => {
    switch (msg.type) {
      case "token": {
        setMessages((prev) => {
          const last = prev[prev.length - 1];
          if (last && last.sender === "agent" && last.isStreaming) {
            return [
              ...prev.slice(0, -1),
              { ...last, content: last.content + msg.content },
            ];
          }
          return [
            ...prev,
            {
              id: nextId(),
              sender: "agent",
              agent: msg.agent,
              content: msg.content,
              timestamp: Date.now(),
              isStreaming: true,
            },
          ];
        });
        break;
      }
      case "tool_call": {
        setToolActions((prev) => [
          ...prev,
          {
            id: nextId(),
            agent: msg.agent,
            tool: msg.tool,
            args: msg.args,
            timestamp: Date.now(),
          },
        ]);
        break;
      }
      case "interrupt": {
        setCurrentInterrupt(msg);
        setIsWaiting(false);
        break;
      }
      case "message_complete": {
        setMessages((prev) => {
          const last = prev[prev.length - 1];
          if (last && last.isStreaming) {
            return [...prev.slice(0, -1), { ...last, isStreaming: false }];
          }
          return prev;
        });
        setIsWaiting(false);
        break;
      }
      case "error": {
        setMessages((prev) => [
          ...prev,
          {
            id: nextId(),
            sender: "agent",
            agent: "System",
            content: `Error: ${msg.message}`,
            timestamp: Date.now(),
          },
        ]);
        setIsWaiting(false);
        break;
      }
    }
  }, []);

  const { status, sendMessage, sendInterruptResponse } =
    useWebSocket(handleServerMessage);

  const handleSend = useCallback(
    (content: string) => {
      setMessages((prev) => [
        ...prev,
        {
          id: nextId(),
          sender: "user",
          content,
          timestamp: Date.now(),
        },
      ]);
      setIsWaiting(true);
      sendMessage(content);
    },
    [sendMessage]
  );

  const handleInterruptResponse = useCallback(
    (approved: boolean) => {
      sendInterruptResponse(approved);
      setCurrentInterrupt(null);
      setIsWaiting(true);
    },
    [sendInterruptResponse]
  );

  return (
    <div style={styles.page}>
      <div style={styles.header}>
        <h1 style={styles.title}>Smart Support</h1>
        <StatusIndicator status={status} />
      </div>
      <ChatMessages messages={messages} />
      {toolActions.length > 0 && (
        <div style={styles.actionsBar}>
          {toolActions.slice(-3).map((action) => (
            <AgentAction key={action.id} action={action} />
          ))}
        </div>
      )}
      {currentInterrupt && (
        <InterruptPrompt
          interrupt={currentInterrupt}
          onRespond={handleInterruptResponse}
        />
      )}
      <ChatInput onSend={handleSend} disabled={isWaiting || status !== "connected"} />
    </div>
  );
}

function StatusIndicator({ status }: { status: ConnectionStatus }) {
  const colors: Record<ConnectionStatus, string> = {
    connected: "#4caf50",
    connecting: "#ff9800",
    disconnected: "#f44336",
  };
  return (
    <div style={{ display: "flex", alignItems: "center", gap: "6px" }}>
      <div
        style={{
          width: "8px",
          height: "8px",
          borderRadius: "50%",
          background: colors[status],
        }}
      />
      <span style={{ fontSize: "12px", color: "#666" }}>{status}</span>
    </div>
  );
}

const styles: Record<string, React.CSSProperties> = {
  page: {
    height: "100vh",
    display: "flex",
    flexDirection: "column",
    background: "white",
    maxWidth: "800px",
    margin: "0 auto",
    boxShadow: "0 0 20px rgba(0,0,0,0.1)",
  },
  header: {
    display: "flex",
    justifyContent: "space-between",
    alignItems: "center",
    padding: "12px 16px",
    borderBottom: "1px solid #e0e0e0",
  },
  title: {
    fontSize: "18px",
    fontWeight: 700,
    margin: 0,
    color: "#333",
  },
  actionsBar: {
    borderTop: "1px solid #eee",
    paddingTop: "4px",
  },
};