import type { ServerWebSocket } from 'bun';
import { randomUUID } from 'crypto';
import type { ClientMessage, ServerMessage, Message, PiEvent } from './types';
import { sessionManager } from './session-manager';
import * as storage from './storage';
import * as piBridge from './pi-bridge';
import { join, resolve } from 'path';
import { homedir } from 'os';
import { getHomeDir, getUserSettingsFile } from '../../../servers/data-path';
import { logger } from './logger';

// Default model when no user preference is set
const DEFAULT_MODEL = 'opencode/big-pickle';

async function getUserDefaultModel(email: string): Promise<string | null> {
  try {
    const settingsPath = getUserSettingsFile(email);
    const file = Bun.file(settingsPath);
    if (await file.exists()) {
      const settings = await file.json();
      return settings?.chat?.defaultModel || null;
    }
  } catch (err) {
    logger.error('Failed to read user settings for default model', { email, error: String(err) });
  }
  return null;
}

type WSData = {
  userId: number;
  email: string;
  username: string;
  role: string;
  provider: string;
};

const IDLE_TIMEOUT_MS = 60 * 60 * 1000; // 1 hour

const resolveSandboxedCwd = (email: string, cwdRoot?: string, cwd?: string) => {
  const root = !cwdRoot || cwdRoot === 'home' ? getHomeDir(email) : getHomeDir(email);
  if (!cwd || cwd === '~') return root;
  if (cwd.startsWith('~/')) return join(root, cwd.slice(2));
  if (cwd.startsWith('/')) return join(root, cwd.slice(1));
  return root;
};

const resolveHostCwd = (cwdRoot?: string, cwd?: string) => {
  if (cwdRoot === 'officer.dev') return resolve(process.cwd(), '..');
  const root = homedir();
  if (!cwd || cwd === '~') return root;
  if (cwd.startsWith('/')) return cwd;
  if (cwd.startsWith('~/')) return join(root, cwd.slice(2));
  return join(root, cwd);
};

export const resolveBaseCwd = (email: string, cwdRoot?: string, cwd?: string) => {
  return resolveHostCwd(cwdRoot, cwd);
};

const wsToSessionMap = new WeakMap<any, string>();

function sendToClient(ws: ServerWebSocket<WSData> | null, msg: ServerMessage): void {
  if (ws?.readyState === 1) {
    ws.send(JSON.stringify(msg));
  }
}

export async function open(ws: ServerWebSocket<WSData>): Promise<void> {
  // logger.info('WebSocket connection opened', { email: ws.data.email });
}

export function message(ws: ServerWebSocket<WSData>, raw: string | Buffer): void {
  const data = typeof raw === 'string' ? raw : raw.toString();
  
  (async () => {
    try {
      const clientMsg = JSON.parse(data) as ClientMessage;
      
      if (clientMsg.type === 'chat') {
        await handleChat(ws, clientMsg);
      } else if (clientMsg.type === 'resume') {
        await handleResume(ws, clientMsg);
      } else if (clientMsg.type === 'stop') {
        await handleStop(ws);
      }
    } catch (err) {
      logger.error('Error handling WebSocket message', { email: ws.data.email, error: String(err) });
      sendToClient(ws, { type: 'error', message: 'Failed to process message' });
    }
  })();
}

export function close(ws: ServerWebSocket<WSData>): void {
  // logger.info('WebSocket connection closed', { email: ws.data.email });
  
  const sessionId = wsToSessionMap.get(ws);
  if (sessionId) {
    sessionManager.detachWs(sessionId);
    sessionManager.setIdleTimeout(sessionId, IDLE_TIMEOUT_MS);
  }
}

function createEventHandler(sessionId: string, model: string, cwd: string, storageDir: string) {
  return async (event: PiEvent): Promise<void> => {
    const session = sessionManager.getSession(sessionId);
    if (!session) return;

    const ws = session.ws as ServerWebSocket<WSData> | null;

    switch (event.type) {
      case 'delta': {
        sendToClient(ws, { type: 'assistant:delta', text: event.text });
        session.streamBuffer += event.text;
        break;
      }

      case 'text': {
        // Flush streaming buffer as complete text
        const text = event.text || session.streamBuffer;
        if (text) {
          sendToClient(ws, { type: 'assistant:text', text });
          
          const assistantMsg: Message = {
            id: randomUUID(),
            timestamp: Date.now(),
            role: 'assistant',
            text,
            model,
          };
          session.messages.push(assistantMsg);
          session.meta.messageCount += 1;
          session.streamBuffer = '';
        }
        break;
      }

      case 'tool:start': {
        // Flush any pending streaming text first
        if (session.streamBuffer) {
          sendToClient(ws, { type: 'assistant:text', text: session.streamBuffer });
          
          const assistantMsg: Message = {
            id: randomUUID(),
            timestamp: Date.now(),
            role: 'assistant',
            text: session.streamBuffer,
            model,
          };
          session.messages.push(assistantMsg);
          session.meta.messageCount += 1;
          session.streamBuffer = '';
        }

        sendToClient(ws, {
          type: 'tool:start',
          toolCallId: event.toolCallId,
          toolName: event.toolName,
          toolInput: event.toolInput,
        });

        const toolMsg: Message = {
          id: randomUUID(),
          timestamp: Date.now(),
          role: 'tool',
          toolCallId: event.toolCallId,
          toolName: event.toolName,
          toolInput: event.toolInput,
        };
        session.messages.push(toolMsg);
        session.meta.messageCount += 1;
        break;
      }

      case 'tool:result': {
        sendToClient(ws, {
          type: 'tool:result',
          toolCallId: event.toolCallId,
          output: event.output,
          isError: event.isError,
        });

        // Update existing tool message with output
        for (let i = session.messages.length - 1; i >= 0; i--) {
          const m = session.messages[i]!;
          if (m.role === 'tool' && m.toolCallId === event.toolCallId) {
            m.output = event.output;
            m.isError = event.isError;
            break;
          }
        }
        break;
      }

      case 'result': {
        // Flush any remaining streaming buffer
        if (session.streamBuffer) {
          sendToClient(ws, { type: 'assistant:text', text: session.streamBuffer });
          
          const assistantMsg: Message = {
            id: randomUUID(),
            timestamp: Date.now(),
            role: 'assistant',
            text: session.streamBuffer,
            model,
            cost: event.cost,
          };
          session.messages.push(assistantMsg);
          session.meta.messageCount += 1;
          session.streamBuffer = '';
        }

        sendToClient(ws, { type: 'result', sessionId, cost: event.cost });
        
        session.isGenerating = false;
        session.meta.cost.inputTokens += event.cost.inputTokens;
        session.meta.cost.outputTokens += event.cost.outputTokens;
        session.meta.cost.totalUSD += event.cost.totalUSD;
        session.meta.updatedAt = Date.now();

        // Save session to disk
        try {
          await storage.saveSession(storageDir, sessionId, session.meta, session.messages);
          logger.info('Session saved to disk', { sessionId, messageCount: session.messages.length });
        } catch (err) {
          logger.error('Failed to save session', { sessionId, error: String(err) });
        }
        break;
      }

      case 'error': {
        sendToClient(ws, { type: 'error', message: event.message });
        session.isGenerating = false;
        break;
      }

      case 'stopped': {
        sendToClient(ws, { type: 'stopped' });
        session.isGenerating = false;
        break;
      }
    }
  };
}

async function handleChat(
  ws: ServerWebSocket<WSData>,
  msg: { prompt: string; sessionId?: string; model?: string; cwd?: string; cwdRoot?: string; sandboxed?: boolean; groupSlug?: string; attachmentIds?: string[] }
): Promise<void> {
  const { email, username, userId } = ws.data;
  const sessionId = msg.sessionId || randomUUID();
  
  // Use provided model, or fall back to user default, or use system default
  let model = msg.model;
  let modelSource = 'client-provided';
  let userDefault = null;
  if (!model) {
    userDefault = await getUserDefaultModel(email);
    if (userDefault) {
      model = userDefault;
      modelSource = 'user-settings';
    } else {
      model = DEFAULT_MODEL;
      modelSource = 'system-default';
    }
  }
  
  logger.info('Model selected for chat', { 
    sessionId, 
    model, 
    modelSource,
    clientModel: msg.model || null,
    userDefault,
  });
  
  const homeDir = getHomeDir(email);
  const sandboxed = msg.sandboxed ?? false;
  const cwd = sandboxed
    ? resolveSandboxedCwd(email, msg.cwdRoot, msg.cwd)
    : resolveHostCwd(msg.cwdRoot, msg.cwd);
  const groupSlug = msg.groupSlug || null;
  const session = sessionManager.getOrCreate(sessionId, email, cwd, model, groupSlug);
  session.sandboxed = sandboxed;
  session.userId = userId;
  sessionManager.attachWs(sessionId, ws);
  wsToSessionMap.set(ws as any, sessionId);

  sendToClient(ws, { type: 'session:init', sessionId, model, cwd });
  if (!session.piProcess) {
    try {
      const onEvent = createEventHandler(sessionId, model, cwd, homeDir);
      session.piProcess = await piBridge.spawnPi(cwd, model, email, onEvent, sandboxed ? { userId, username, email, homeDir } : undefined);
      logger.info('Spawned Pi process for session', { sessionId, model, cwd, sandboxed });
    } catch (err) {
      logger.error('Failed to spawn Pi process', { sessionId, model, error: String(err) });
      sendToClient(ws, { type: 'error', message: 'Failed to start Pi process' });
      return;
    }
  }

  // Add user message to session
  const userMsg: Message = {
    id: randomUUID(),
    timestamp: Date.now(),
    role: 'user',
    text: msg.prompt,
  };
  session.messages.push(userMsg);
  session.meta.messageCount += 1;
  session.meta.updatedAt = Date.now();

  if (!session.meta.title) {
    session.meta.title = msg.prompt.slice(0, 100);
  }

  // Send prompt to Pi
  const requestId = randomUUID();
  session.isGenerating = true;
  piBridge.sendPrompt(session.piProcess, msg.prompt, requestId);
}

async function handleResume(
  ws: ServerWebSocket<WSData>,
  msg: { sessionId: string; cwd?: string; cwdRoot?: string }
): Promise<void> {
  const { email } = ws.data;
  const { sessionId } = msg;

  try {
    let session = sessionManager.getSession(sessionId);

    if (!session) {
      const homeDir = getHomeDir(email);

      try {
        const { meta, messages } = await storage.loadSession(homeDir, sessionId);
        
        session = sessionManager.getOrCreate(sessionId, email, meta.cwd, meta.model);
        session.messages = messages;
        session.meta = meta;
        
        logger.info('Loaded session from disk', { sessionId, messageCount: messages.length });
      } catch (err) {
        logger.error('Failed to load session from disk', { sessionId, error: String(err) });
        sendToClient(ws, { type: 'error', message: 'Session not found', errorCode: 'SESSION_NOT_FOUND' });
        return;
      }
    }

    sessionManager.attachWs(sessionId, ws);
    wsToSessionMap.set(ws as any, sessionId);

    sendToClient(ws, { type: 'session:init', sessionId, model: session.model, cwd: session.cwd });

    // Spawn fresh Pi process if needed
    if (!session.piProcess) {
      try {
        const homeDir = getHomeDir(email);
        const sandbox = session.sandboxed && session.userId ? { userId: session.userId, username: ws.data.username, email, homeDir } : undefined;
        const onEvent = createEventHandler(sessionId, session.model, session.cwd, homeDir);
        session.piProcess = await piBridge.spawnPi(session.cwd, session.model, email, onEvent, sandbox);
        logger.info('Spawned fresh Pi process for resumed session', { sessionId, model: session.model, sandboxed: session.sandboxed });
      } catch (err) {
        logger.error('Failed to spawn Pi process for resume', { sessionId, error: String(err) });
        sendToClient(ws, { type: 'error', message: 'Failed to start Pi process' });
        return;
      }
    }

    sendToClient(ws, {
      type: 'sync:messages',
      sessionId,
      messages: session.messages,
      isGenerating: session.isGenerating,
      streamingText: session.streamBuffer,
    });
    
    logger.info('Session resumed successfully', { sessionId, messageCount: session.messages.length });
  } catch (err) {
    logger.error('Unexpected error in handleResume', { sessionId, error: String(err) });
    sendToClient(ws, { type: 'error', message: 'Failed to resume session' });
  }
}

async function handleStop(ws: ServerWebSocket<WSData>): Promise<void> {
  const sessionId = wsToSessionMap.get(ws);
  
  if (sessionId) {
    const session = sessionManager.getSession(sessionId);
    
    if (session?.piProcess) {
      try {
        piBridge.abort(session.piProcess, randomUUID());
        logger.info('Sent abort to Pi process', { sessionId });
        session.isGenerating = false;
      } catch (err) {
        logger.error('Failed to abort Pi process', { sessionId, error: String(err) });
      }
    }
  }
  
  sendToClient(ws, { type: 'stopped' });
}

export const piWebsocket = {
  open,
  message,
  close,
  drain() {},
};