main.js

const { app, BrowserWindow, Tray, Menu, nativeImage, screen, ipcMain, dialog, clipboard } = require('electron');
const path = require('path');
const fs = require('fs-extra');
const http = require('http');
const https = require('https');
const OpenAI = require('openai');

let mainWindow = null;
let tray = null;
let openai = null;

const COLOR_LOCK_DEFAULTS = {
  primary: { enabled: false, value: '#0A84FF' },
  secondary: { enabled: false, value: '#30D158' },
  background: { enabled: false, value: '#111111' },
  surface: { enabled: false, value: '#1C1C1E' },
  accent: { enabled: false, value: '#FFD60A' },
  success: { enabled: false, value: '#30D158' },
  error: { enabled: false, value: '#FF453A' },
  textPrimary: { enabled: false, value: '#F5F5F7' },
  textSecondary: { enabled: false, value: '#8E8E93' },
  border: { enabled: false, value: '#3A3A3C' }
};

// ─── System prompts ──────────────────────────────────────────────────────────

const DESIGN_MD_SYSTEM = `You are a senior design-systems analyst. Analyze the screenshot and produce a compact but high-signal DESIGN.md-style document that another AI agent can use to recreate the visual system faithfully.

Core behavior:
1. Extract reusable system rules, not screenshot trivia.
2. Prefer specific, implementation-useful observations over broad design adjectives.
3. Infer likely real-world values when needed: standard spacing scales, realistic radii, plausible font stacks, likely border/shadow treatments.
4. Do not invent hidden behavior, content, or product structure that is not visually supported.
5. When evidence is partial, choose the most likely design-system interpretation and state it cleanly without hand-wringing.
6. Focus on what a builder would need to recreate the shell consistently across multiple screens.

Think in layers:
- page atmosphere and density
- color roles and surface hierarchy
- typography roles rather than specific words
- recurring component rules
- spacing rhythm and alignment system
- depth, borders, and background treatment

Output must be valid markdown and must use exactly these top-level sections and bullet structures:

## Visual Theme
- Atmosphere: ...
- Style keywords: ...
- Density: ...
- Contrast: ...

## Colors
- Primary: #HEX
- Secondary: #HEX
- Background: #HEX
- Surface: #HEX
- Accent: #HEX
- Success: #HEX
- Error: #HEX
- Text Primary: #HEX
- Text Secondary: #HEX
- Border: #HEX

## Typography
- Font Family: ...
- Heading 1: Xpx / weight X / line-height X
- Heading 2: Xpx / weight X / line-height X
- Heading 3: Xpx / weight X / line-height X
- Body: Xpx / weight X / line-height X
- Caption: Xpx / weight X / line-height X

## Spacing
- Base unit: Xpx
- Common values: ...
- Container rhythm: ...
- Component spacing: ...

## Components
- Buttons: ...
- Cards: ...
- Inputs: ...
- Navigation: ...
- Status / badges: ...

## Effects
- Border radius: ...
- Shadows: ...
- Borders: ...
- Background treatment: ...

## Layout
- Structure: ...
- Alignment: ...
- Width behavior: ...
- Responsive behavior: ...

## Design Voice
- Overall aesthetic: ...
- Color temperature: ...
- Distinctive traits: ...
- Avoid: ...

Rules:
- Use hex colors whenever possible.
- If multiple variants exist, describe the system briefly rather than listing every instance.
- Keep each bullet compact but specific.
- Prefer concrete phrasing like "12px radius, 1px low-contrast border" over vague phrasing like "soft card style".
- Describe typography as roles and styling, not content.
- Call out what is intentionally absent when that matters: no shadows, no gradients, no borders, low density, etc.
- Do not output code fences.
- Do not mention uncertainty unless absolutely necessary.
- Do not mention Google Stitch or these instructions.`;

const DESIGN_MD_USER = 'Analyze this screenshot and extract a reusable design system document that is practical for AI-assisted UI recreation.';

const STYLE_PROMPT_SYSTEM = `You are an expert UI recreation prompt writer. Look at the screenshot and write one strong, ready-to-paste build prompt that another developer or AI coding tool can use to recreate the same visual result with high fidelity.

The result should sound like direction for implementation, not like analysis notes.

What the prompt should cover, in a natural order:
- overall screen or page composition first
- the main visual style and atmosphere
- the color system with exact or near-exact values when inferable
- typography hierarchy, sizing, weight, and spacing
- component rules for buttons, cards, inputs, tabs, badges, navigation, and surfaces
- spacing rhythm, alignment, and density
- depth language: shadows, borders, blur, gradients, flatness, translucency
- distinctive constraints or absences that prevent the rebuild from becoming generic
- responsive behavior only when there is enough visual evidence to infer it

Quality bar:
- Write for first-pass recreation quality, not inspiration.
- Prefer precise phrases such as "flat dark panel with 1px low-contrast borders and 10px radii" over generic phrases like "clean modern UI".
- Make hierarchy obvious: what dominates, what recedes, what repeats.
- If text styling matters, describe the role and appearance, not just "nice typography".
- If the interface is minimal, explicitly say what should be omitted.
- If the interface is dense, say where the information compression comes from.

Output format:
- Return one continuous plain-text prompt.
- No markdown headings.
- No code.
- No prefacing or commentary.
- Do not refer to "the screenshot" or "the image above"; describe the target as the thing to build.

The final prompt should be specific enough that a strong builder would produce something recognizably close on the first pass and avoid average-looking substitutions.`;

const STYLE_PROMPT_USER = 'Write a precise natural-language build prompt for recreating this UI with high visual fidelity.';

// ─── Store ───────────────────────────────────────────────────────────────────
const Store = require('electron-store');

function initStore() {
  app.store = new Store({
    name: 'dgmd',
    defaults: {
      provider: 'openai',          // openai | anthropic | gemini | ollama
      apiKey: '',                   // OpenAI key
      anthropicKey: '',             // Anthropic key
      geminiKey: '',                // Gemini key
      ollamaUrl: 'http://localhost:11434',
      ollamaModel: '',
      shellOnly: false,
      colorLocks: COLOR_LOCK_DEFAULTS,
      outputDir: path.join(app.getPath('home'), 'design-md-files')
    }
  });
  fs.ensureDirSync(app.store.get('outputDir'));
}

// ─── Generic vision call ─────────────────────────────────────────────────────

async function visionCall(imageData, systemPrompt, userPrompt) {
  const shellOnly = !!app.store.get('shellOnly');
  const colorLocks = app.store.get('colorLocks') || COLOR_LOCK_DEFAULTS;
  const lockedEntries = [
    ['Primary', colorLocks.primary],
    ['Secondary', colorLocks.secondary],
    ['Background', colorLocks.background],
    ['Surface', colorLocks.surface],
    ['Accent', colorLocks.accent],
    ['Success', colorLocks.success],
    ['Error', colorLocks.error],
    ['Text Primary', colorLocks.textPrimary],
    ['Text Secondary', colorLocks.textSecondary],
    ['Border', colorLocks.border]
  ].filter(([, config]) => config?.enabled && config?.value);
  const shellInstruction = shellOnly
    ? '\n\nSHELL ONLY MODE: Ignore visible text content from the image. Do not reproduce, quote, paraphrase, summarize, or infer the actual words, labels, numbers, metrics, headings, button copy, paragraph copy, table values, or OCR-style content shown in the screenshot. Describe only the visual shell: layout, composition, hierarchy, spacing, component structure, generic text roles, text density, typography treatment, color, surfaces, borders, radii, and effects. It is acceptable to say things like "large headline area", "compact label row", or "muted body text blocks", but never include the actual text content.'
    : '';
  const colorLockInstruction = systemPrompt === DESIGN_MD_SYSTEM && lockedEntries.length
    ? `\n\nCOLOR LOCK MODE: In the ## Colors section, force these roles to these exact values and do not infer alternatives for them:\n${lockedEntries.map(([label, config]) => `- ${label}: ${config.value.toUpperCase()}`).join('\n')}\nKeep other color roles inferred normally unless they are also locked.`
    : '';
  const finalSystemPrompt = systemPrompt + shellInstruction + colorLockInstruction;
  const p = app.store.get('provider');
  if (p === 'anthropic') return callAnthropic(imageData, finalSystemPrompt, userPrompt);
  if (p === 'gemini')    return callGemini(imageData, finalSystemPrompt, userPrompt);
  if (p === 'ollama')    return callOllama(imageData, finalSystemPrompt, userPrompt);
  return callOpenAI(imageData, finalSystemPrompt, userPrompt);
}

// ─── OpenAI ──────────────────────────────────────────────────────────────────
function initOpenAI() {
  const k = app.store?.get('apiKey') || process.env.OPENAI_API_KEY;
  if (k) openai = new OpenAI({ apiKey: k });
}

async function callOpenAI(imageData, systemPrompt, userPrompt) {
  if (!openai) initOpenAI();
  if (!openai) throw new Error('Set your OpenAI API key first.');
  const r = await openai.chat.completions.create({
    model: 'gpt-4o',
    messages: [
      { role: 'system', content: systemPrompt },
      { role: 'user', content: [
        { type: 'text', text: userPrompt },
        { type: 'image_url', image_url: { url: imageData } }
      ]}
    ],
    max_tokens: 3000
  });
  return r.choices[0].message.content;
}

// ─── Anthropic (Claude) ─────────────────────────────────────────────────────
async function callAnthropic(imageDataUri, systemPrompt, userPrompt) {
  const key = app.store.get('anthropicKey');
  if (!key) throw new Error('Set your Anthropic API key first.');

  // Extract media type + base64
  const match = imageDataUri.match(/^data:(image\/[^;]+);base64,(.+)$/);
  if (!match) throw new Error('Invalid image data');
  const [, mediaType, base64] = match;

  const body = {
    model: 'claude-sonnet-4-20250514',
    max_tokens: 3000,
    system: systemPrompt,
    messages: [{
      role: 'user',
      content: [
        { type: 'image', source: { type: 'base64', media_type: mediaType, data: base64 } },
        { type: 'text', text: userPrompt }
      ]
    }]
  };

  const data = await jsonFetch('https://api.anthropic.com/v1/messages', {
    method: 'POST',
    headers: {
      'x-api-key': key,
      'anthropic-version': '2023-06-01',
      'content-type': 'application/json'
    },
    body,
    timeout: 120000
  });

  if (data.error) throw new Error(data.error.message || JSON.stringify(data.error));
  return (data.content || []).filter(b => b.type === 'text').map(b => b.text).join('\n');
}

// ─── Google Gemini ───────────────────────────────────────────────────────────
async function callGemini(imageDataUri, systemPrompt, userPrompt) {
  const key = app.store.get('geminiKey');
  if (!key) throw new Error('Set your Gemini API key first.');

  const match = imageDataUri.match(/^data:(image\/[^;]+);base64,(.+)$/);
  if (!match) throw new Error('Invalid image data');
  const [, mimeType, base64] = match;

  const url = `https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent?key=${key}`;
  const body = {
    system_instruction: { parts: [{ text: systemPrompt }] },
    contents: [{
      parts: [
        { inline_data: { mime_type: mimeType, data: base64 } },
        { text: userPrompt }
      ]
    }]
  };

  const data = await jsonFetch(url, {
    method: 'POST',
    headers: { 'content-type': 'application/json' },
    body,
    timeout: 120000
  });

  if (data.error) throw new Error(data.error.message || JSON.stringify(data.error));
  const parts = data.candidates?.[0]?.content?.parts || [];
  return parts.map(p => p.text || '').join('\n');
}

// ─── Ollama ──────────────────────────────────────────────────────────────────
async function fetchOllamaModels(baseUrl) {
  const data = await jsonFetch(`${baseUrl}/api/tags`);
  return (data.models || []).map(m => ({ name: m.name, size: m.size, modified: m.modified_at }));
}

async function callOllama(imageDataUri, systemPrompt, userPrompt) {
  const baseUrl = app.store.get('ollamaUrl') || 'http://localhost:11434';
  const model = app.store.get('ollamaModel');
  if (!model) throw new Error('Select an Ollama model first.');
  const base64 = imageDataUri.replace(/^data:image\/[^;]+;base64,/, '');
  const data = await jsonFetch(`${baseUrl}/api/chat`, {
    method: 'POST', timeout: 300000,
    body: { model, stream: false, messages: [
      { role: 'system', content: systemPrompt },
      { role: 'user', content: userPrompt, images: [base64] }
    ]}
  });
  return data?.message?.content || '';
}

// ─── Shared HTTP helper ──────────────────────────────────────────────────────
function jsonFetch(urlStr, opts = {}) {
  return new Promise((resolve, reject) => {
    const url = new URL(urlStr);
    const mod = url.protocol === 'https:' ? https : http;
    const req = mod.request(url, {
      method: opts.method || 'GET',
      headers: { 'Content-Type': 'application/json', ...(opts.headers || {}) },
      timeout: opts.timeout || 120000
    }, res => {
      let body = '';
      res.on('data', c => { body += c; });
      res.on('end', () => {
        if (res.statusCode >= 200 && res.statusCode < 300) {
          try { resolve(JSON.parse(body)); } catch { resolve(body); }
        } else {
          reject(new Error(`HTTP ${res.statusCode}: ${body.slice(0, 300)}`));
        }
      });
    });
    req.on('error', reject);
    req.on('timeout', () => { req.destroy(); reject(new Error('Request timed out')); });
    if (opts.body) req.write(JSON.stringify(opts.body));
    req.end();
  });
}

// ─── Window / Tray ───────────────────────────────────────────────────────────

function createWindow() {
  const { width: sw } = screen.getPrimaryDisplay().workAreaSize;
  const cur = screen.getCursorScreenPoint();
  let x = cur.x - 190, y = cur.y + 20;
  if (x < 0) x = 20; if (x + 380 > sw) x = sw - 400; if (y < 0) y = 20;

  mainWindow = new BrowserWindow({
    width: 340, height: 480, x, y,
    minWidth: 300, minHeight: 380, maxWidth: 500, maxHeight: 700,
    show: false, frame: false, resizable: true,
    transparent: true, vibrancy: 'sidebar', visualEffectState: 'active',
    backgroundColor: '#00000000',
    webPreferences: { nodeIntegration: false, contextIsolation: true, preload: path.join(__dirname, 'preload.js') },
    alwaysOnTop: true
  });
  mainWindow.loadFile('index.html');
  mainWindow.on('blur', () => { if (!mainWindow?.webContents?.isDevToolsOpened?.()) mainWindow.hide(); });
  mainWindow.on('closed', () => { mainWindow = null; });
}

function createTray() {
  const iconPath = path.join(__dirname, 'assets', 'iconTemplate.png');
  let icon;
  try { icon = nativeImage.createFromPath(iconPath); if (icon.isEmpty()) icon = createDefaultIcon(); }
  catch { icon = createDefaultIcon(); }
  tray = new Tray(icon.resize({ width: 22, height: 22 }));

  // Left-click: toggle window
  tray.on('click', (_e, bounds) => {
    if (!mainWindow) createWindow();
    if (mainWindow.isVisible()) mainWindow.hide();
    else { mainWindow.setPosition(bounds.x - 170, bounds.y + 4); mainWindow.show(); mainWindow.focus(); }
  });

  // right-click: context menu
  const contextMenu = Menu.buildFromTemplate([
    {
      label: 'Settings',
      click: () => {
        if (!mainWindow) createWindow();
        if (!mainWindow.isVisible()) {
          const { x, y } = screen.getCursorScreenPoint();
          mainWindow.setPosition(x - 170, y + 4);
          mainWindow.show();
        }
        mainWindow.focus();
        mainWindow.webContents.send('open-settings');
      }
    },
    { type: 'separator' },
    {
      label: 'Quit DGMD',
      click: () => { app.isQuitting = true; app.quit(); }
    }
  ]);
  tray.on('right-click', () => { tray.popUpContextMenu(contextMenu); });

  tray.setToolTip('DGMD');
}

function createDefaultIcon() {
  // Refined camera mark fallback for the tray template icon.
  const s = 44, buf = Buffer.alloc(s * s * 4);
  const bodyL = 7, bodyR = 37, bodyT = 14, bodyB = 32;
  const humpL = 13, humpR = 22, humpT = 10, humpB = 15;
  const lensCx = 22, lensCy = 23, lensOuter = 7, lensInner = 3.5;
  const flashL = 29, flashR = 32, flashT = 17, flashB = 20;

  function insideRoundedRect(x, y, l, t, r, b, rad) {
    if (x < l || x > r || y < t || y > b) return false;
    if ((x >= l + rad && x <= r - rad) || (y >= t + rad && y <= b - rad)) return true;
    const cx = x < l + rad ? l + rad : r - rad;
    const cy = y < t + rad ? t + rad : b - rad;
    return (x - cx) ** 2 + (y - cy) ** 2 <= rad ** 2;
  }

  for (let y = 0; y < s; y++) for (let x = 0; x < s; x++) {
    const i = (y*s+x)*4;
    const inBody = insideRoundedRect(x, y, bodyL, bodyT, bodyR, bodyB, 4);
    const inHump = insideRoundedRect(x, y, humpL, humpT, humpR, humpB, 2);
    const lensDist = Math.sqrt((x - lensCx) ** 2 + (y - lensCy) ** 2);
    const inLensRing = lensDist <= lensOuter && lensDist >= lensInner;
    const inFlash = x >= flashL && x <= flashR && y >= flashT && y <= flashB;
    if (inBody || inHump || inLensRing || inFlash) {
      buf[i] = 0;
      buf[i + 1] = 0;
      buf[i + 2] = 0;
      buf[i + 3] = 255;
    }
  }
  return nativeImage.createFromBuffer(buf, { width: s, height: s });
}

// ─── IPC Handlers ────────────────────────────────────────────────────────────

ipcMain.handle('get-settings', () => ({
  provider:     app.store.get('provider'),
  hasApiKey:    !!app.store.get('apiKey'),
  hasAnthropicKey: !!app.store.get('anthropicKey'),
  hasGeminiKey: !!app.store.get('geminiKey'),
  ollamaUrl:    app.store.get('ollamaUrl'),
  ollamaModel:  app.store.get('ollamaModel'),
  shellOnly:    !!app.store.get('shellOnly'),
  colorLocks:   app.store.get('colorLocks') || COLOR_LOCK_DEFAULTS,
  outputDir:    app.store.get('outputDir')
}));

ipcMain.handle('set-provider',      (_e, v) => { app.store.set('provider', v);      return { success: true }; });
ipcMain.handle('set-api-key',       (_e, v) => { app.store.set('apiKey', v); initOpenAI(); return { success: true }; });
ipcMain.handle('set-anthropic-key', (_e, v) => { app.store.set('anthropicKey', v);  return { success: true }; });
ipcMain.handle('set-gemini-key',    (_e, v) => { app.store.set('geminiKey', v);     return { success: true }; });
ipcMain.handle('set-ollama-url',    (_e, v) => { app.store.set('ollamaUrl', v);     return { success: true }; });
ipcMain.handle('set-ollama-model',  (_e, v) => { app.store.set('ollamaModel', v);   return { success: true }; });
ipcMain.handle('set-shell-only',    (_e, v) => { app.store.set('shellOnly', !!v);   return { success: true }; });
ipcMain.handle('set-color-locks',   (_e, v) => { app.store.set('colorLocks', v || COLOR_LOCK_DEFAULTS); return { success: true }; });
ipcMain.handle('reset-credentials', () => {
  app.store.set('apiKey', '');
  app.store.set('anthropicKey', '');
  app.store.set('geminiKey', '');
  openai = null;
  return { success: true };
});

ipcMain.handle('fetch-ollama-models', async (_e, baseUrl) => {
  try { return { success: true, models: await fetchOllamaModels(baseUrl || app.store.get('ollamaUrl')) }; }
  catch (err) { return { error: err.message || 'Cannot reach Ollama' }; }
});

ipcMain.handle('set-output-dir', async () => {
  const r = await dialog.showOpenDialog(mainWindow, { properties: ['openDirectory','createDirectory'], defaultPath: app.store.get('outputDir') });
  if (!r.canceled && r.filePaths[0]) { app.store.set('outputDir', r.filePaths[0]); fs.ensureDirSync(r.filePaths[0]); return { path: r.filePaths[0] }; }
  return null;
});

// Paste image from clipboard
ipcMain.handle('paste-image', () => {
  const img = clipboard.readImage();
  if (img.isEmpty()) return { error: 'No image in clipboard' };
  const png = img.toPNG();
  const b64 = `data:image/png;base64,${png.toString('base64')}`;
  return { success: true, base64Data: b64, size: png.length };
});

// Analyze image — both outputs in parallel
ipcMain.handle('analyze-image', async (_e, { base64Data }) => {
  try {
    const [designMd, stylePrompt] = await Promise.all([
      visionCall(base64Data, DESIGN_MD_SYSTEM, DESIGN_MD_USER),
      visionCall(base64Data, STYLE_PROMPT_SYSTEM, STYLE_PROMPT_USER)
    ]);
    return { success: true, designMd, stylePrompt };
  } catch (error) {
    console.error('Error analyzing image:', error);
    return { error: error.message || 'Failed to analyze image' };
  }
});

// File ops
ipcMain.handle('save-design-md', async (_e, { content, filename }) => {
  const p = path.join(app.store.get('outputDir'), filename || 'design.md');
  await fs.writeFile(p, content); return { success: true, path: p };
});
ipcMain.handle('copy-to-clipboard', (_e, text) => { clipboard.writeText(text); return { success: true }; });
ipcMain.handle('open-output-folder', () => { require('electron').shell.openPath(app.store.get('outputDir')); return { success: true }; });

// ─── App lifecycle ───────────────────────────────────────────────────────────
app.whenReady().then(() => {
  // Hide from dock & Cmd+Tab — menubar-only app
  if (app.dock) app.dock.hide();

  initStore(); initOpenAI(); createWindow(); createTray();
  app.on('window-all-closed', e => e.preventDefault());
});
app.on('before-quit', () => { if (tray) tray.destroy(); });