import { useState, useEffect, useCallback, PropsWithChildren, useMemo, useRef } from 'react';
import {
  useExternalStoreRuntime,
  AssistantRuntimeProvider,
  CompositeAttachmentAdapter,
  SimpleImageAttachmentAdapter,
  SimpleTextAttachmentAdapter,
  type ExternalStoreAdapter,
  type ThreadMessageLike,
} from '@assistant-ui/react';
import { SimpleFileAttachmentAdapter } from './runtime/file-attachment-adapter';
import { PersonaConfig, PersonaMessage, PersonaResponse, ProtocolStatus, Session, ThreadData, AgentState } from './types';
import { parseMessages, convertMessage, createTurnGrouper } from './messages';
import { initializeProtocols } from './runtime/protocols';
import { createPacketListener } from './runtime/listeners';
import { createOnNewHandler, createOnEditHandler, createOnReloadHandler, createOnCancelHandler } from './runtime/handlers';
import { PersonaRuntimeContext } from './runtime/context';
import { createThreadListAdapter } from './runtime/threads';
import { PersonaSessionStorage } from './storage';

export {
  usePersonaRuntime,
  usePersonaRuntimeProtocol,
  usePersonaRuntimeEndpoint,
  usePersonaRuntimeWebRTCProtocol,
  usePersonaRuntimeMessages,
  usePersonaRenameThread,
  useAgentState,
  useSteerAgent,
} from './runtime/context';

export type { PersonaMessage, PersonaResponse, ThreadData };

type StoreMessage = PersonaMessage | ThreadMessageLike;
type StoreMessages = Pick<ExternalStoreAdapter<StoreMessage>, 'messages' | 'convertMessage' | 'setMessages'>;

// Module-level: the external store caches conversions per converter identity.
const convertPersonaMessage = (message: StoreMessage) => convertMessage(message as PersonaMessage);
const passThroughMessage = (message: StoreMessage) => message as ThreadMessageLike;

function PersonaRuntimeProviderInner({
  dev = false,
  baseUrl,
  protocols: _protocols,
  logger,
  children,
  session: initialSession = 'new',
  transformMessages,
  threads = false,
  sessionStorage: _sessionStorage,
  onThreadCreate,
  onThreadSwitch,
  onThreadRename,
  onThreadArchive,
  onThreadUnarchive,
  onThreadDelete,
  enableAttachments = false,
  enableSteer: _enableSteer = false,
  onSteerAccepted,
  onSteerProcessed,
  groupTurns = false,
  ...config
}: Readonly<PersonaConfig>) {
  const [agentState, setAgentState] = useState<AgentState>('idle');
  const [messages, setMessages] = useState<PersonaMessage[]>([]);
  const [protocolsStatus, setProtocolsStatus] = useState<Map<string, ProtocolStatus>>(new Map());
  const didMount = useRef(false);
  const isSwitchingThread = useRef(false);

  // Initialize message storage
  const sessionStorage = useMemo(
    () =>
      _sessionStorage ||
      new PersonaSessionStorage({
        apiKey: config.apiKey!,
        endpoint: baseUrl || 'https://persona.applica.guru/api',
      }),
    [_sessionStorage, baseUrl, config.apiKey],
  );

  // Thread management state
  const [threadList, setThreadList] = useState<ThreadData[]>([]);
  const [currentThreadId, setCurrentThreadId] = useState<string>('');

  // Use refs to avoid re-renders and effect loops
  const currentThreadRef = useRef<Session>(typeof initialSession === 'string' ? initialSession : 'new');
  const hasInitialized = useRef(false);

  // Load threads on mount (only once)
  useEffect(() => {
    if (!threads || hasInitialized.current) {
      return;
    }
    hasInitialized.current = true;

    sessionStorage
      .list(config.agentId, config.userId)
      .then((threads) => {
        if (threads?.length > 0) {
          setThreadList(threads);
          setCurrentThreadId(threads[0].threadId);
        }
      })
      .catch((error) => logger?.error('Failed to load threads:', error));
  }, [threads, sessionStorage, config.agentId, config.userId, logger]);

  // Initialize protocols
  const protocols = useMemo(
    () =>
      initializeProtocols(_protocols, {
        dev: dev || false,
        baseUrl,
        apiKey: config.apiKey,
        agentId: config.agentId,
        userId: config.userId,
        tools: config.tools,
        logger,
      }),
    [_protocols, dev, baseUrl, config.apiKey, config.agentId, config.userId, config.tools, logger],
  );

  // Setup protocol listeners (only once on mount)
  useEffect(() => {
    if (didMount.current) {
      return;
    }
    didMount.current = true;
    logger?.debug(
      'Setting up protocols:',
      protocols.map((p) => p.getName()),
    );
    protocols.forEach((protocol) => {
      protocol.setSession(currentThreadRef.current);
      protocol.clearListeners();

      // Status change listener
      protocol.addStatusChangeListener((status: ProtocolStatus) => {
        logger?.debug(`${protocol.getName()} status: ${status}`);
        setProtocolsStatus((prev) => {
          const next = new Map(prev);
          next.set(protocol.getName(), status);
          return next;
        });

        // Only send initial packets on first connection (not during thread switches)
        if (status === 'connected' && !isSwitchingThread.current) {
          if (config.context) {
            protocol
              .sendPacket({
                type: 'command',
                payload: { command: 'set_initial_context', arguments: config.context },
              })
              .catch(() => {});
          }

          if (config.tools && Array.isArray(config.tools)) {
            protocol
              .sendPacket({
                type: 'command',
                payload: { command: 'set_local_tools', arguments: { tools: config.tools.map((t) => t.schema) } },
              })
              .catch(() => {});
          }
        }
      });

      // Packet listener
      protocol.addPacketListener(
        createPacketListener(
          protocol,
          setMessages,
          setAgentState,
          protocols,
          transformMessages,
          onSteerAccepted,
          onSteerProcessed,
        ),
      );

      // Auto-connect if configured
      if (protocol.autostart && protocol.status === 'disconnected') {
        logger?.debug(`Auto-connecting protocol: ${protocol.getName()}`);
        protocol.connect(currentThreadRef.current).catch(() => {});
        setAgentState('idle');
      }
    });
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, []); // Run only once

  // Load messages for an existing session on mount (without threads)
  useEffect(() => {
    if (threads || !initialSession || initialSession === 'new') {
      return;
    }

    sessionStorage
      .sync(initialSession)
      .then((loadedMessages) => {
        if (loadedMessages.length > 0) {
          const processed = parseMessages(loadedMessages);
          setMessages(transformMessages ? transformMessages(processed) : processed);
          logger?.debug(`Loaded ${loadedMessages.length} messages for session ${initialSession}`);
        }
      })
      .catch((error) => {
        logger?.error('Failed to load session messages:', error);
      });
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, []); // Run only once on mount

  // Handle thread switching
  useEffect(() => {
    if (!threads || !currentThreadId) {
      setMessages([]);
      return;
    }

    if (isSwitchingThread.current || currentThreadId === currentThreadRef.current) {
      return;
    }

    // Mark as switching to prevent concurrent switches
    isSwitchingThread.current = true;
    logger?.info('Switching to thread:', currentThreadId);

    const switchThread = async () => {
      try {
        // Disconnect active protocols in parallel
        const connectedProtocols = protocols.filter((p) => p.status === 'connected');
        if (connectedProtocols.length > 0) {
          await Promise.allSettled(connectedProtocols.map((p) => p.disconnect()));
        }

        // Update thread reference
        currentThreadRef.current = currentThreadId;

        // Set new session on all protocols
        await Promise.allSettled(protocols.map((p) => p.setSession(currentThreadId)));

        // Load messages from storage
        const loadedMessages = await sessionStorage.sync(currentThreadId);
        const processed = parseMessages(loadedMessages);
        setMessages(transformMessages ? transformMessages(processed) : processed);
        logger?.debug(`Loaded ${loadedMessages.length} messages for thread ${currentThreadId}`);

        // Reconnect autostart protocols
        const autostartProtocols = protocols.filter((p) => p.autostart);
        if (autostartProtocols.length > 0) {
          await Promise.allSettled(autostartProtocols.map((p) => p.connect(currentThreadId)));
        }
      } catch (error) {
        logger?.error('Failed to switch thread:', error);
        setMessages([]);
      } finally {
        isSwitchingThread.current = false;
      }
    };

    switchThread();
  }, [currentThreadId, threads, protocols, sessionStorage, transformMessages, logger]);

  // Message handlers
  const onNew = useCallback(
    createOnNewHandler(protocols, setMessages, setAgentState, transformMessages, logger),
    [protocols, transformMessages, logger],
  );

  const onEdit = useCallback(
    createOnEditHandler(protocols, messages, setMessages, setAgentState, transformMessages, logger),
    [protocols, messages, transformMessages, logger],
  );

  const onReload = useCallback(
    createOnReloadHandler(protocols, messages, setMessages, setAgentState, transformMessages, logger),
    [protocols, messages, transformMessages, logger],
  );

  const onCancel = useCallback(createOnCancelHandler(protocols, setAgentState, logger), [protocols, logger]);

  const handleSetMessages = useCallback(
    (newMessages: readonly PersonaMessage[]) => {
      const processed = parseMessages([...newMessages]);
      setMessages(transformMessages ? transformMessages(processed) : processed);
    },
    [transformMessages],
  );

  const getMessages = useCallback(() => messages, [messages]);

  // Thread list adapter
  const threadListAdapter = useMemo(() => {
    if (!threads) return undefined;

    return createThreadListAdapter(
      currentThreadId,
      threadList,
      setCurrentThreadId,
      setThreadList,
      sessionStorage,
      onThreadCreate,
      onThreadSwitch,
      onThreadArchive,
      onThreadUnarchive,
      onThreadDelete,
      onThreadRename,
    );
  }, [
    threads,
    currentThreadId,
    threadList,
    sessionStorage,
    onThreadCreate,
    onThreadSwitch,
    onThreadArchive,
    onThreadUnarchive,
    onThreadDelete,
    onThreadRename,
    logger,
  ]);

  // Create runtime
  const isRunning = agentState === 'running' || agentState === 'cancelling';
  // groupTurns folds each agent turn into one assistant message. Branch
  // switching (setMessages) works on the raw messages, so it is off in that mode.
  const groupTurnMessages = useMemo(() => createTurnGrouper(), []);
  const turnMessages = useMemo(
    () => (groupTurns ? groupTurnMessages(messages) : undefined),
    [groupTurns, groupTurnMessages, messages],
  );
  const storeMessages: StoreMessages = turnMessages
    ? { messages: turnMessages, convertMessage: passThroughMessage }
    : {
        messages,
        convertMessage: convertPersonaMessage,
        setMessages: handleSetMessages as (messages: readonly StoreMessage[]) => void,
      };

  const runtime = useExternalStoreRuntime<StoreMessage>({
    isRunning,
    ...storeMessages,
    onNew,
    onEdit,
    onCancel,
    onReload,
    adapters: {
      ...(enableAttachments
        ? {
            attachments: new CompositeAttachmentAdapter([
              new SimpleImageAttachmentAdapter(),
              new SimpleTextAttachmentAdapter(),
              new SimpleFileAttachmentAdapter(),
            ]),
          }
        : {}),
      ...(threadListAdapter ? { threadList: threadListAdapter } : {}),
    },
  });

  const renameThread = useCallback(
    (threadId: string, name: string) => {
      setThreadList((prev) => prev.map((t) => (t.threadId === threadId ? { ...t, title: name } : t)));
      onThreadRename?.(threadId, name);
      sessionStorage.rename(threadId, name);
    },
    [sessionStorage, onThreadRename],
  );

  const steerAgent = useCallback(
    async (text: string) => {
      const connectedProtocols = protocols.filter(
        (p) => p.getName() !== 'transaction' && p.status === 'connected',
      );
      if (connectedProtocols.length === 0) {
        logger?.warn('No connected protocol available for steer');
        return;
      }
      logger?.debug('Sending steer command:', text);
      await Promise.allSettled(
        connectedProtocols.map((protocol) =>
          protocol.sendPacket({
            type: 'command',
            payload: { command: 'steer', arguments: { text } },
          }),
        ),
      );
    },
    [protocols, logger],
  );

  return (
    <PersonaRuntimeContext.Provider value={{ protocols, protocolsStatus, agentState, getMessages, renameThread, steerAgent }}>
      <AssistantRuntimeProvider runtime={runtime}>{children}</AssistantRuntimeProvider>
    </PersonaRuntimeContext.Provider>
  );
}

export function PersonaRuntimeProvider({ children, ...config }: PropsWithChildren<PersonaConfig>) {
  return <PersonaRuntimeProviderInner {...config}>{children}</PersonaRuntimeProviderInner>;
}