import { WebSocket } from 'ws';
import { spawn } from 'child_process';
import path from 'path';
import os from 'os';
import {
  getClient,
  createSession,
  sendPromptAsync,
  subscribeEvents,
  unsubscribeEvents,
  convertToChatMessage,
  extractTextFromParts,
  respondPermission,
} from '../services/aiService.js';

const shellSessions = new Map();

const chatSessions = new Map();

const activeStreams = new Map();

export function handleShellWebSocket(ws, request) {
  const url = new URL(request.url || '', 'http://localhost');
  const sessionId = url.searchParams.get('sessionId') || `shell-${Date.now()}`;
  const cwdParam = url.searchParams.get('cwd');

  console.log('[Shell] New shell session:', sessionId);

  const defaultCwd = cwdParam || process.env.CODEFREE_DIRECTORY || path.join(os.homedir(), 'Desktop');
  const shell = spawn(process.platform === 'win32' ? 'cmd.exe' : '/bin/bash', [], {
    cwd: defaultCwd,
    env: { ...process.env, TERM: 'xterm-256color' },
  });

  shellSessions.set(sessionId, { ws, shell });

  shell.stdout.on('data', (data) => {
    if (ws.readyState === WebSocket.OPEN) {
      ws.send(JSON.stringify({ type: 'output', data: data.toString('base64'), sessionId }));
    }
  });

  shell.stderr.on('data', (data) => {
    if (ws.readyState === WebSocket.OPEN) {
      ws.send(JSON.stringify({ type: 'output', data: data.toString('base64'), sessionId }));
    }
  });

  shell.on('close', (code) => {
    if (ws.readyState === WebSocket.OPEN) {
      ws.send(JSON.stringify({ type: 'exit', code, sessionId }));
    }
    shellSessions.delete(sessionId);
  });

  ws.on('message', (data) => {
    try {
      const message = JSON.parse(data.toString());
      if (message.type === 'input' && shell.stdin.writable) {
        shell.stdin.write(Buffer.from(message.data, 'base64'));
      } else if (message.type === 'resize') {
        // shell resize not supported in simple spawn mode
      }
    } catch (e) {
      console.error('[Shell] Error:', e.message);
    }
  });

  ws.on('close', () => {
    shell.kill();
    shellSessions.delete(sessionId);
  });
}

/**
 * 处理聊天 WebSocket 连接，通过 SDK 与 CodeFree-O 服务通信
 */
export function handleChatWebSocket(ws, request, connectedClients, broadcastToClients) {
  const url = new URL(request.url || '', 'http://localhost');
  const sessionId = url.searchParams.get('sessionId') || `chat-${Date.now()}`;

  console.log('[Chat] New chat session:', sessionId);

  let currentSessionId = null;
  let streamingMessageId = null;
  let streamAbortController = null;

  chatSessions.set(sessionId, { ws, sdkSessionId: null });

  initSessionAndStream(ws, sessionId);

  ws.on('message', (data) => {
    try {
      const message = JSON.parse(data.toString());

      if (message.type === 'chat_message') {
        handleChatMessage(ws, sessionId, message.content);
      } else if (message.type === 'abort') {
        handleAbort(sessionId);
      }
    } catch (e) {
      console.error('[Chat] Error:', e.message);
    }
  });

  ws.on('close', () => {
    console.log('[Chat] Client disconnected:', sessionId);
    cleanupStream(sessionId);
    chatSessions.delete(sessionId);
  });

  /**
   * 初始化 SDK 会话并启动 SSE 事件流订阅
   */
  async function initSessionAndStream(ws, localSessionId) {

      const session = await createSession('New Chat');
      const sdkSessionId = session.id;

      const chatSession = chatSessions.get(localSessionId);
      if (chatSession) {
        chatSession.sdkSessionId = sdkSessionId;
      }
      currentSessionId = sdkSessionId;

      if (ws.readyState === WebSocket.OPEN) {
        ws.send(JSON.stringify({
          type: 'session_created',
          sessionId: localSessionId,
          session: {
            id: localSessionId,
            name: session.title || 'New Chat',
            messages: [],
          },
        }));
      }

      startEventStream(localSessionId, sdkSessionId, ws);
  
  }

  /**
   * 启动 SSE 事件流订阅，将 SDK 事件转发为 WebSocket 消息
   */
  async function startEventStream(localSessionId, sdkSessionId, ws) {
    streamAbortController = new AbortController();
    activeStreams.set(localSessionId, streamAbortController);

    try {
      await subscribeEvents(
        async (event) => {
          if (ws.readyState !== WebSocket.OPEN) return;

          switch (event.type) {
            case 'message.updated': {
              const msgInfo = event.properties.info;
              if (msgInfo.sessionID !== sdkSessionId) return;

              if (msgInfo.role === 'user') {
                ws.send(JSON.stringify({
                  type: 'message_received',
                  sessionId: localSessionId,
                  message: convertToChatMessage(msgInfo, []),
                }));
              }
              break;
            }

            case 'message.part.updated': {
              const part = event.properties.part;
              if (part.sessionID !== sdkSessionId) return;

              if (part.type === 'text') {
                const delta = event.properties.delta || '';
                const fullText = part.text || '';

                if (!streamingMessageId) {
                  streamingMessageId = part.messageID;
                  ws.send(JSON.stringify({
                    type: 'assistant_message_chunk',
                    sessionId: localSessionId,
                    messageId: part.messageID,
                    delta,
                    content: fullText,
                  }));
                } else if (streamingMessageId === part.messageID) {
                  ws.send(JSON.stringify({
                    type: 'assistant_message_chunk',
                    sessionId: localSessionId,
                    messageId: part.messageID,
                    delta,
                    content: fullText,
                  }));
                }
              } else if (part.type === 'tool') {
                const toolState = part.state;
                if (toolState.status === 'running') {
                  ws.send(JSON.stringify({
                    type: 'tool_status',
                    sessionId: localSessionId,
                    tool: part.tool,
                    status: 'running',
                    title: toolState.title || part.tool,
                  }));
                }
              }
              break;
            }

            case 'session.idle': {
              if (event.properties.sessionID !== sdkSessionId) return;

              if (streamingMessageId) {
                try {
                  const sdk = getClient();
                  const messagesResult = await sdk.session.messages({
                    path: { id: sdkSessionId },
                  });
                  if (messagesResult.data) {
                    const lastAssistant = [...messagesResult.data]
                      .reverse()
                      .find((m) => m.info.role === 'assistant' && m.info.id === streamingMessageId);
                    if (lastAssistant) {
                      const chatMsg = convertToChatMessage(lastAssistant.info, lastAssistant.parts);
                      ws.send(JSON.stringify({
                        type: 'assistant_message',
                        sessionId: localSessionId,
                        message: chatMsg,
                      }));
                    }
                  }
                } catch (e) {
                  console.error('[Chat] Failed to fetch final message:', e.message);
                }
                streamingMessageId = null;
              }
              break;
            }

            case 'permission.updated': {
              const perm = event.properties;
              if (perm.sessionID !== sdkSessionId) return;

              try {
                await respondPermission(perm.sessionID, perm.id, 'always');
              } catch (e) {
                console.error('[Chat] Failed to auto-approve permission:', e.message);
              }
              break;
            }

            case 'session.error': {
              if (event.properties.sessionID && event.properties.sessionID !== sdkSessionId) return;
              const errorMsg = event.properties.error || 'AI 服务发生错误';
              sendErrorMessage(ws, localSessionId, streamingMessageId, errorMsg);
              streamingMessageId = null;
              break;
            }

            default:
              break;
          }
        },
        { signal: streamAbortController.signal }
      );
    } catch (error) {
      if (error.name !== 'AbortError') {
        console.error('[Chat] Event stream error:', error.message);
      }
    }
  }

  /**
   * 处理用户聊天消息，通过 SDK promptAsync 发送
   */
  async function handleChatMessage(ws, localSessionId, content) {
    const chatSession = chatSessions.get(localSessionId);
    if (!chatSession || !chatSession.sdkSessionId) {
      sendErrorMessage(ws, localSessionId, null, '会话尚未初始化，请稍后重试');
      return;
    }

    const userMessage = {
      id: `msg-${Date.now()}`,
      role: 'user',
      content,
      timestamp: new Date().toISOString(),
    };

    if (ws.readyState === WebSocket.OPEN) {
      ws.send(JSON.stringify({
        type: 'message_received',
        sessionId: localSessionId,
        message: userMessage,
      }));
    }

    try {
      await sendPromptAsync(chatSession.sdkSessionId, content);
    } catch (error) {
      console.error('[Chat] Failed to send prompt:', error);
      sendErrorMessage(ws, localSessionId, null, error.message);
    }
  }

  /**
   * 中止当前会话的 AI 响应
   */
  async function handleAbort(localSessionId) {
    const chatSession = chatSessions.get(localSessionId);
    if (!chatSession || !chatSession.sdkSessionId) return;

    try {
      const { abortSession } = await import('../services/aiService.js');
      await abortSession(chatSession.sdkSessionId);
    } catch (error) {
      console.error('[Chat] Failed to abort session:', error.message);
    }
  }

  /**
   * 清理事件流
   */
  function cleanupStream(localSessionId) {
    const controller = activeStreams.get(localSessionId);
    if (controller) {
      controller.abort();
      activeStreams.delete(localSessionId);
    }
    streamAbortController = null;
  }
}

/**
 * 发送错误消息到前端
 */
function sendErrorMessage(ws, sessionId, messageId, errorMessage) {
  const errorContent = `抱歉，AI 服务暂时不可用，请稍后重试。\n\n\`错误信息: ${errorMessage}\``;

  const assistantMessage = {
    id: messageId || `msg-${Date.now()}-error`,
    role: 'assistant',
    content: errorContent,
    timestamp: new Date().toISOString(),
  };

  if (ws.readyState === WebSocket.OPEN) {
    ws.send(JSON.stringify({
      type: 'assistant_message',
      sessionId,
      message: assistantMessage,
    }));
  }
}
