main.js

const { app, BrowserWindow, protocol, shell, session, net } = require('electron');
const path = require('path');

let mainWindow;
let authWindow;

const PROTOCOL = 'unifyai';
const AUTH_CALLBACK_PATH = 'auth-callback';

function createWindow() {
  mainWindow = new BrowserWindow({
    width: 1200,
    height: 800,
    webPreferences: {
      webviewTag: true,
      contextIsolation: true,
      nodeIntegration: false,
      preload: path.join(__dirname, 'preload.js')
    }
  });

  mainWindow.loadFile('dist/index.html');
  mainWindow.webContents.openDevTools();
  
  // Set up IPC handlers
  mainWindow.webContents.ipc.handle('open-auth-window', async (event, service) => {
    return openAuthWindow(service);
  });
  
  mainWindow.webContents.ipc.handle('check-auth-status', async (event, service) => {
    return checkAuthStatus(service);
  });
  
  mainWindow.webContents.ipc.handle('switch-service', async (event, service) => {
    // Handle service switching logic
    console.log('Switching to service:', service);
    return { success: true, service };
  });
  
  mainWindow.webContents.ipc.handle('share-context', async (event, data) => {
    // Handle context sharing logic
    console.log('Sharing context:', data);
    return { success: true };
  });
  
  mainWindow.webContents.ipc.handle('create-project', async (event, project) => {
    // Handle project creation logic
    console.log('Creating project:', project);
    return { success: true, project };
  });
  
  mainWindow.webContents.ipc.handle('update-project', async (event, projectId, data) => {
    // Handle project update logic
    console.log('Updating project:', projectId, data);
    return { success: true };
  });
}

function handleAuthCallback(url) {
  console.log('Auth callback received:', url);
  
  const urlParams = new URL(url);
  const params = new URLSearchParams(urlParams.search);
  
  if (authWindow && !authWindow.isDestroyed()) {
    authWindow.close();
  }
  
  if (mainWindow) {
    mainWindow.webContents.send('auth-success', {
      url: url,
      params: Object.fromEntries(params)
    });
  }
}

app.setAsDefaultProtocolClient(PROTOCOL);

protocol.registerSchemesAsPrivileged([
  { scheme: PROTOCOL, privileges: { secure: true, standard: true } }
]);

app.on('open-url', (event, url) => {
  event.preventDefault();
  console.log('Received URL:', url);
  
  if (url.startsWith(`${PROTOCOL}://${AUTH_CALLBACK_PATH}`)) {
    handleAuthCallback(url);
  }
});

app.whenReady().then(() => {
  protocol.handle(PROTOCOL, (request) => {
    const url = request.url;
    console.log('Protocol handler:', url);
    
    if (url.includes(AUTH_CALLBACK_PATH)) {
      handleAuthCallback(url);
    }
    
    return new Response('Auth callback received. You can close this window.', {
      headers: { 'content-type': 'text/html' }
    });
  });
  
  createWindow();
});

app.on('window-all-closed', () => {
  if (process.platform !== 'darwin') {
    app.quit();
  }
});

app.on('activate', () => {
  if (BrowserWindow.getAllWindows().length === 0) {
    createWindow();
  }
});

// Function to open auth window and monitor for successful login
async function openAuthWindow(service = 'chatgpt') {
  authWindow = new BrowserWindow({
    width: 500,
    height: 700,
    webPreferences: {
      nodeIntegration: false,
      contextIsolation: true
    }
  });
  
  const authUrls = {
    chatgpt: 'https://chatgpt.com/auth/login',
    claude: 'https://claude.ai/login'
  };
  
  const domains = {
    chatgpt: '.chatgpt.com',
    claude: '.claude.ai'
  };
  
  const successUrls = {
    chatgpt: ['https://chatgpt.com/', 'https://chatgpt.com/?'],
    claude: ['https://claude.ai/chats', 'https://claude.ai/', 'https://claude.ai/chat']
  };
  
  const authUrl = authUrls[service];
  const domain = domains[service];
  const successPatterns = successUrls[service];
  
  authWindow.loadURL(authUrl);
  authWindow.webContents.openDevTools(); // Open dev tools to see what's happening
  
  // Add a manual way to close and transfer cookies after 30 seconds
  const autoCloseTimer = setTimeout(async () => {
    if (authWindow && !authWindow.isDestroyed()) {
      console.log(`Auto-closing ${service} auth window after 30 seconds`);
      await transferCookiesAndClose();
    }
  }, 30000);
  
  async function transferCookiesAndClose() {
    try {
      // Get cookies from the auth window
      const cookies = await authWindow.webContents.session.cookies.get({ domain });
      console.log(`Found ${cookies.length} cookies for ${service}:`, cookies.map(c => c.name));
      
      if (cookies.length > 0) {
        // Get the webview's partition session
        const webviewSession = session.fromPartition(`persist:${service}`);
        
        // Transfer cookies to the webview's partition
        for (const cookie of cookies) {
          await webviewSession.cookies.set({
            url: authUrls[service].replace('/auth/login', '').replace('/login', ''),
            name: cookie.name,
            value: cookie.value,
            domain: cookie.domain,
            path: cookie.path,
            secure: cookie.secure,
            httpOnly: cookie.httpOnly,
            expirationDate: cookie.expirationDate
          });
        }
        
        // Notify renderer that auth is complete
        if (mainWindow) {
          mainWindow.webContents.send('auth-success', { 
            success: true,
            service: service,
            cookieCount: cookies.length 
          });
        }
      }
      
      clearTimeout(autoCloseTimer);
      authWindow.close();
    } catch (error) {
      console.error(`Error transferring ${service} cookies:`, error);
      clearTimeout(autoCloseTimer);
      authWindow.close();
    }
  }
  
  // Monitor for successful login by checking URL changes
  authWindow.webContents.on('did-navigate', async (event, url) => {
    console.log(`${service} auth window navigated to:`, url);
    
    // For Claude, also check for intermediate redirect URLs
    let isSuccess = false;
    if (service === 'claude') {
      // Claude might redirect through various URLs, so be more flexible
      isSuccess = url.includes('claude.ai') && 
                  !url.includes('/login') && 
                  !url.includes('/auth') &&
                  (url.includes('/chats') || url.includes('/chat') || url === 'https://claude.ai/');
    } else {
      // For other services, use the exact patterns
      isSuccess = successPatterns.some(pattern => 
        url === pattern || url.startsWith(pattern)
      );
    }
    
    if (isSuccess) {
      console.log(`${service} login successful!`);
      // Wait a moment for cookies to be fully set, then transfer
      setTimeout(transferCookiesAndClose, 2000);
    }
  });
  
  // Add manual close button functionality
  authWindow.webContents.on('dom-ready', () => {
    // Inject a "Transfer & Close" button for manual testing
    authWindow.webContents.executeJavaScript(`
      if (!document.getElementById('manual-close-btn')) {
        const btn = document.createElement('button');
        btn.id = 'manual-close-btn';
        btn.innerText = 'Transfer Cookies & Close';
        btn.style.position = 'fixed';
        btn.style.top = '10px';
        btn.style.right = '10px';
        btn.style.zIndex = '9999';
        btn.style.background = '#10b981';
        btn.style.color = 'white';
        btn.style.border = 'none';
        btn.style.padding = '10px';
        btn.style.borderRadius = '4px';
        btn.style.cursor = 'pointer';
        btn.onclick = () => {
          window.postMessage('transfer-cookies', '*');
        };
        document.body.appendChild(btn);
      }
    `);
  });
  
  // Listen for manual transfer message
  authWindow.webContents.on('ipc-message', (event, channel) => {
    if (channel === 'transfer-cookies') {
      transferCookiesAndClose();
    }
  });
  
  authWindow.on('closed', () => {
    authWindow = null;
  });
}

// Check if user is authenticated
async function checkAuthStatus(service = 'chatgpt') {
  const domains = {
    chatgpt: '.chatgpt.com',
    claude: '.claude.ai'
  };
  
  const webviewSession = session.fromPartition(`persist:${service}`);
  const cookies = await webviewSession.cookies.get({ domain: domains[service] });
  
  const hasAuthCookies = cookies.some(cookie => 
    cookie.name.includes('sess') || 
    cookie.name.includes('auth') ||
    cookie.name.includes('__Secure') ||
    cookie.name.includes('token') ||
    cookie.name.includes('sessionid')
  );
  
  return hasAuthCookies;
}