index.ts

import { Hono } from 'hono';
import { serveStatic } from 'hono/bun';
import postgres from 'postgres';
import { pipeline, type FeatureExtractionPipeline } from '@xenova/transformers';
import { nanoid } from 'nanoid';
import { readdir, stat, mkdir, writeFile } from 'node:fs/promises';
import { join, extname } from 'node:path';

// =============================================================================
// CONFIG
// =============================================================================

const pkg = await Bun.file('package.json').json();
const VERSION = pkg.version;
const APP_NAME = process.env.APP_NAME || 'ShitPostr';
const PORT = parseInt(process.env.PORT || '3000');
const DATABASE_URL = process.env.DATABASE_URL || 'postgres://postgres:postgres@localhost:5432/shitpostr';
const OLLAMA_URL = process.env.OLLAMA_URL || 'http://localhost:11434';
const OLLAMA_MODEL = process.env.OLLAMA_MODEL || 'qwen2.5vl:3b';
const UPLOAD_DIR = process.env.UPLOAD_DIR || '/data/memes/uploads';
// Allow serving both the writable meme volume and the repo-backed template set by default.
const STATIC_DIRS = (process.env.STATIC_DIRS || '/data/memes,/data/templates').split(',').map(d => d.trim());
const IMAGE_EXTENSIONS = new Set(['.jpg', '.jpeg', '.png', '.gif', '.webp', '.avif']);
const SKIP_DB_INIT = process.env.SKIP_DB_INIT === '1';
const TEST_MODE = process.env.TEST_MODE === '1';

// =============================================================================
// PATH NORMALIZATION
// =============================================================================

function canonicalizeDataPath(p: string): string {
  // Normalize host-style "data/..." and docker-style "/data/..." to a single form
  // so scans/uploads from different environments don't create duplicates.
  if (p.startsWith('/data/')) return p;
  if (p.startsWith('data/')) return `/${p}`;
  return p;
}

function dataPathAlternates(p: string): string[] {
  const canonical = canonicalizeDataPath(p);
  if (canonical.startsWith('/data/')) return [canonical, canonical.slice(1)]; // "/data/.." + "data/.."
  return [canonical];
}

async function fileExistsForDataPath(p: string): Promise<boolean> {
  // Support both docker-style (/data/...) and host-style (data/...) paths.
  for (const candidate of dataPathAlternates(p)) {
    if (await Bun.file(candidate).exists()) return true;
  }
  return false;
}

// =============================================================================
// DATABASE
// =============================================================================

const sql = postgres(DATABASE_URL, { max: 10 });

async function initDb() {
  const schema = await Bun.file('schema.sql').text();
  await sql.unsafe(schema);
  console.log('Database initialized');
}

// =============================================================================
// EMBEDDINGS
// =============================================================================

let embedder: FeatureExtractionPipeline | null = null;

async function getEmbedder() {
  if (!embedder) {
    console.log('Loading embedding model...');
    embedder = await pipeline('feature-extraction', 'Xenova/all-MiniLM-L6-v2');
    console.log('Embedding model loaded');
  }
  return embedder;
}

async function embed(text: string): Promise<number[]> {
  const model = await getEmbedder();
  const result = await model(text, { pooling: 'mean', normalize: true });
  return Array.from(result.data as Float32Array);
}

// =============================================================================
// OLLAMA
// =============================================================================

async function describeImage(filePath: string): Promise<string> {
  // Try local path first (data/...), then Docker path (/data/...)
  let file = Bun.file(filePath);
  if (!await file.exists()) {
    const altPath = filePath.startsWith('/') ? filePath.slice(1) : `/${filePath}`;
    file = Bun.file(altPath);
    if (!await file.exists()) {
      throw new Error(`Image not found: ${filePath}`);
    }
  }
  const buffer = await file.arrayBuffer();
  const base64 = Buffer.from(buffer).toString('base64');

  const res = await fetch(`${OLLAMA_URL}/api/generate`, {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify({
      model: OLLAMA_MODEL,
      prompt: 'Describe this meme image briefly and concisely for search purposes. Focus on the visual content, any text visible, and the apparent humor or message. Keep it under 100 words.',
      images: [base64],
      stream: false,
    }),
  });

  if (!res.ok) throw new Error(`Ollama error: ${res.status}`);
  const data = await res.json();
  return data.response;
}

async function checkOllama(): Promise<{ available: boolean; model: string; error?: string }> {
  try {
    const res = await fetch(`${OLLAMA_URL}/api/tags`);
    if (!res.ok) return { available: false, model: OLLAMA_MODEL, error: 'Cannot connect' };
    const data = await res.json();
    const hasModel = data.models?.some((m: { name: string }) => m.name.includes(OLLAMA_MODEL.split(':')[0]));
    return { available: hasModel, model: OLLAMA_MODEL, error: hasModel ? undefined : 'Model not found' };
  } catch (e) {
    return { available: false, model: OLLAMA_MODEL, error: 'Connection failed' };
  }
}

// =============================================================================
// SCANNER
// =============================================================================

let scanState = { status: 'idle' as 'idle' | 'scanning' | 'complete' | 'error', processed: 0, total: 0 };

async function isDirectory(path: string): Promise<boolean> {
  try {
    return (await stat(path)).isDirectory();
  } catch {
    return false;
  }
}

async function resolveDataDir(dir: string): Promise<string> {
  // Support both docker-style (/data/...) and local repo-style (data/...) paths.
  if (await isDirectory(dir)) return dir;
  if (dir.startsWith('/data/')) {
    const local = dir.slice(1); // "/data/..." -> "data/..."
    if (await isDirectory(local)) return local;
  } else if (dir.startsWith('data/')) {
    const docker = `/${dir}`; // "data/..." -> "/data/..."
    if (await isDirectory(docker)) return docker;
  }
  return dir;
}

async function* walkDir(dir: string): AsyncGenerator<string> {
  const entries = await readdir(dir, { withFileTypes: true });
  for (const entry of entries) {
    const fullPath = join(dir, entry.name);
    if (entry.isDirectory() && !entry.name.startsWith('.')) {
      yield* walkDir(fullPath);
    } else if (entry.isFile() && IMAGE_EXTENSIONS.has(extname(entry.name).toLowerCase())) {
      yield fullPath;
    }
  }
}

async function scanDirectory(dir: string): Promise<{ added: number; skipped: number }> {
  scanState = { status: 'scanning', processed: 0, total: 0 };
  let added = 0, skipped = 0;

  try {
    // Count first
    for await (const _ of walkDir(dir)) scanState.total++;
    console.log(`Scanning ${scanState.total} images in ${dir}`);

    // Process
    for await (const filePath of walkDir(dir)) {
      scanState.processed++;
      const canonicalPath = canonicalizeDataPath(filePath);
      const alts = dataPathAlternates(canonicalPath);
      const [existing] = await sql`SELECT 1 FROM memes WHERE file_path = ANY(${alts}::text[])`;
      if (existing) { skipped++; continue; }

      const stats = await stat(filePath);
      await sql`INSERT INTO memes (file_path, meta) VALUES (${canonicalPath}, ${sql.json({
        filesize: stats.size,
        format: extname(filePath).slice(1).toLowerCase(),
      })})`;
      added++;
    }

    scanState.status = 'complete';
    console.log(`Scan complete: ${added} added, ${skipped} skipped`);
    return { added, skipped };
  } catch (e) {
    scanState.status = 'error';
    throw e;
  }
}

// =============================================================================
// HONO APP
// =============================================================================

const app = new Hono();

// Serve static files
app.get('/', (c) => c.html(Bun.file('index.html').text()));

// Serve images from any /data/* directory via /images/*
// e.g., /images/memes/foo.jpg -> /data/memes/foo.jpg
app.get('/images/*', async (c) => {
  const subPath = c.req.path.slice('/images'.length); // e.g., /memes/uploads/xxx.jpg

  // Security: prevent directory traversal
  if (subPath.includes('..')) return c.text('Forbidden', 403);

  // Try local path first (data/...), then Docker path (/data/...)
  const localPath = `data${subPath}`;
  const dockerPath = `/data${subPath}`;

  // Validate path is under an allowed directory
  const isAllowed = STATIC_DIRS.some(dir =>
    dockerPath.startsWith(dir) || localPath.startsWith(dir.replace(/^\//, ''))
  );
  if (!isAllowed) return c.text('Forbidden', 403);

  let file = Bun.file(localPath);
  if (!await file.exists()) {
    file = Bun.file(dockerPath);
    if (!await file.exists()) {
      return c.text('Not found', 404);
    }
  }

  const ext = subPath.split('.').pop()?.toLowerCase();
  const types: Record<string, string> = { jpg: 'image/jpeg', jpeg: 'image/jpeg', png: 'image/png', gif: 'image/gif', webp: 'image/webp', avif: 'image/avif' };
  c.header('Content-Type', types[ext || ''] || 'application/octet-stream');
  c.header('Content-Length', String(file.size));
  c.header('Cache-Control', 'public, max-age=31536000, immutable');
  return c.body(file.stream());
});

// Health check
app.get('/health', (c) => c.json({ ok: true }));
app.get('/api/version', (c) => c.json({ version: VERSION, appName: APP_NAME }));

// Config (for frontend) - DO NOT expose token
app.get('/api/config', (c) => c.json({
  ziplineEnabled: !!(process.env.ZIPLINE_URL && process.env.ZIPLINE_TOKEN)
}));

// Proxy upload to Zipline (secure - token stays server-side)
app.post('/api/zipline-upload', async (c) => {
  const ziplineUrl = process.env.ZIPLINE_URL;
  const ziplineToken = process.env.ZIPLINE_TOKEN;

  if (!ziplineUrl || !ziplineToken) {
    return c.json({ error: 'Zipline not configured' }, 400);
  }

  const formData = await c.req.formData();
  const file = formData.get('file') as File;
  if (!file) return c.json({ error: 'No file provided' }, 400);

  try {
    const buffer = await file.arrayBuffer();
    const uploadFormData = new FormData();
    uploadFormData.append('file', new Blob([buffer]), file.name);

    const response = await fetch(ziplineUrl, {
      method: 'POST',
      headers: { 'Authorization': ziplineToken },
      body: uploadFormData
    });

    if (!response.ok) throw new Error(`HTTP ${response.status}`);
    const data = await response.json();
    const url = data.files?.[0]?.url || data.url;
    return c.json({ url });
  } catch (e) {
    return c.json({ error: String(e) }, 500);
  }
});

// =============================================================================
// API: Memes
// =============================================================================

app.get('/api/memes', async (c) => {
  if (TEST_MODE) return c.json({ memes: [], count: 0 });
  const { status, starred, limit = '50', offset = '0' } = c.req.query();
  const conditions: string[] = [];
  const values: unknown[] = [];
  let i = 1;

  if (status) { conditions.push(`status = $${i++}`); values.push(status); }
  if (starred === 'true') { conditions.push(`starred = $${i++}`); values.push(true); }

  const where = conditions.length ? `WHERE ${conditions.join(' AND ')}` : '';
  const query = `SELECT * FROM memes ${where} ORDER BY created_at DESC LIMIT $${i++} OFFSET $${i++}`;
  values.push(parseInt(limit), parseInt(offset));

  const memes = await sql.unsafe(query, values);
  return c.json({ memes, count: memes.length });
});

app.get('/api/memes/:id', async (c) => {
  const [meme] = await sql`SELECT * FROM memes WHERE id = ${c.req.param('id')}`;
  return meme ? c.json(meme) : c.json({ error: 'Not found' }, 404);
});

app.patch('/api/memes/:id', async (c) => {
  const id = c.req.param('id');
  const body = await c.req.json();
  const updates: string[] = [];
  const values: unknown[] = [];
  let i = 1;

  if (body.title !== undefined) { updates.push(`title = $${i++}`); values.push(body.title); }
  if (body.description !== undefined) { updates.push(`description = $${i++}`); values.push(body.description); }
  if (body.tags !== undefined) { updates.push(`tags = $${i++}`); values.push(body.tags); }
  if (body.starred !== undefined) { updates.push(`starred = $${i++}`); values.push(body.starred); }
  if (body.status !== undefined) { updates.push(`status = $${i++}`); values.push(body.status); }

  if (!updates.length) return c.json({ error: 'No updates' }, 400);
  values.push(id);

  const [meme] = await sql.unsafe(
    `UPDATE memes SET ${updates.join(', ')} WHERE id = $${i} RETURNING *`,
    values
  );
  return meme ? c.json(meme) : c.json({ error: 'Not found' }, 404);
});

app.delete('/api/memes/:id', async (c) => {
  const result = await sql`DELETE FROM memes WHERE id = ${c.req.param('id')}`;
  return c.json({ success: result.count > 0 });
});

app.post('/api/memes/:id/star', async (c) => {
  const [meme] = await sql`
    UPDATE memes SET starred = NOT starred WHERE id = ${c.req.param('id')} RETURNING *
  `;
  return meme ? c.json(meme) : c.json({ error: 'Not found' }, 404);
});

app.post('/api/memes/:id/share', async (c) => {
  const id = c.req.param('id');
  const body = await c.req.json();
  const { url, textBoxes } = body;

  const [meme] = await sql`SELECT * FROM memes WHERE id = ${id}`;
  if (!meme) return c.json({ error: 'Not found' }, 404);

  const meta = meme.meta || {};
  const shares = meta.shares || [];
  shares.push({ url, textBoxes: textBoxes || [], created_at: new Date().toISOString() });

  const [updated] = await sql`
    UPDATE memes SET meta = ${sql.json({ ...meta, shares })} WHERE id = ${id} RETURNING *
  `;
  return c.json(updated);
});

app.post('/api/memes/:id/generate', async (c) => {
  const id = c.req.param('id');
  const [meme] = await sql`SELECT * FROM memes WHERE id = ${id}`;
  if (!meme) return c.json({ error: 'Not found' }, 404);

  await sql`UPDATE memes SET status = 'processing' WHERE id = ${id}`;

  try {
    const description = await describeImage(meme.file_path);
    const embedding = await embed(description);
    await sql`UPDATE memes SET description = ${description}, embedding = ${`[${embedding.join(',')}]`}::vector, status = 'complete' WHERE id = ${id}`;
    return c.json({ success: true });
  } catch (e) {
    await sql`UPDATE memes SET status = 'error' WHERE id = ${id}`;
    return c.json({ error: String(e) }, 500);
  }
});

// =============================================================================
// API: Search
// =============================================================================

app.get('/api/search', async (c) => {
  const { q, mode = 'vector', limit = '20' } = c.req.query();
  if (!q) return c.json({ results: [] });

  if (mode === 'vector') {
    const embedding = await embed(q);
    const embStr = `[${embedding.join(',')}]`;
    const results = await sql`
      SELECT *, 1 - (embedding <=> ${embStr}::vector) as score
      FROM memes WHERE embedding IS NOT NULL
      ORDER BY embedding <=> ${embStr}::vector
      LIMIT ${parseInt(limit)}
    `;
    return c.json({ results: results.map(m => ({ meme: m, score: m.score })) });
  } else {
    const results = await sql`
      SELECT *, ts_rank(to_tsvector('english', COALESCE(description, '')), plainto_tsquery('english', ${q})) as score
      FROM memes
      WHERE to_tsvector('english', COALESCE(description, '')) @@ plainto_tsquery('english', ${q})
      ORDER BY score DESC
      LIMIT ${parseInt(limit)}
    `;
    return c.json({ results: results.map(m => ({ meme: m, score: m.score })) });
  }
});

// =============================================================================
// API: Stats & Utilities
// =============================================================================

app.get('/api/random', async (c) => {
  const { status } = c.req.query();

  if (status) {
    const [meme] = await sql`SELECT * FROM memes WHERE embedding IS NOT NULL AND status = ${status} ORDER BY RANDOM() LIMIT 1`;
    return meme ? c.json(meme) : c.json({ error: 'No memes found' }, 404);
  }

  const [meme] = await sql`SELECT * FROM memes WHERE embedding IS NOT NULL ORDER BY RANDOM() LIMIT 1`;
  return meme ? c.json(meme) : c.json({ error: 'No memes found' }, 404);
});

app.get('/api/memes/:id/duplicates', async (c) => {
  const id = c.req.param('id');
  const { threshold = '0.85' } = c.req.query();
  const thresholdNum = parseFloat(threshold);
  if (isNaN(thresholdNum) || thresholdNum < 0 || thresholdNum > 1) {
    return c.json({ error: 'Invalid threshold: must be between 0 and 1' }, 400);
  }

  const [meme] = await sql`SELECT id, file_path, title, description, embedding FROM memes WHERE id = ${id}`;
  if (!meme || !meme.embedding) return c.json({ error: 'Meme not found or no embedding' }, 404);

  const results = await sql`
    SELECT
      id,
      file_path,
      title,
      description,
      1 - (embedding <=> ${meme.embedding}::vector) as similarity
    FROM memes
    WHERE id != ${id} AND embedding IS NOT NULL
      AND (1 - (embedding <=> ${meme.embedding}::vector)) >= ${thresholdNum}
    ORDER BY embedding <=> ${meme.embedding}::vector
    LIMIT 20
  `;
  return c.json({ duplicates: results.map(r => ({ meme: { id: r.id, file_path: r.file_path, title: r.title, description: r.description }, similarity: r.similarity })) });
});

app.get('/api/stats', async (c) => {
  if (TEST_MODE) {
    return c.json({ total: 0, pending: 0, processing: 0, complete: 0, error: 0, starred: 0 });
  }
  const [stats] = await sql`
    SELECT
      COUNT(*)::int as total,
      COUNT(*) FILTER (WHERE status = 'pending')::int as pending,
      COUNT(*) FILTER (WHERE status = 'processing')::int as processing,
      COUNT(*) FILTER (WHERE status = 'complete')::int as complete,
      COUNT(*) FILTER (WHERE status = 'error')::int as error,
      COUNT(*) FILTER (WHERE starred)::int as starred
    FROM memes
  `;
  return c.json(stats);
});

app.get('/api/ollama-status', async (c) => c.json(await checkOllama()));

app.post('/api/cleanup', async (c) => {
  const memes = await sql`SELECT id, file_path FROM memes`;
  let deleted = 0;
  for (const meme of memes) {
    if (!await fileExistsForDataPath(meme.file_path)) {
      await sql`DELETE FROM memes WHERE id = ${meme.id}`;
      deleted++;
    }
  }
  return c.json({ checked: memes.length, deleted });
});

app.post('/api/reset-processing', async (c) => {
  const result = await sql`UPDATE memes SET status = 'pending' WHERE status = 'processing'`;
  return c.json({ reset: result.count });
});

app.post('/api/reset-errors', async (c) => {
  const result = await sql`UPDATE memes SET status = 'pending' WHERE status = 'error'`;
  return c.json({ reset: result.count });
});

// =============================================================================
// API: Upload & Scan
// =============================================================================

app.post('/api/upload', async (c) => {
  const resolvedUploadDir = await resolveDataDir(UPLOAD_DIR);
  await mkdir(resolvedUploadDir, { recursive: true }).catch(() => {});
  const formData = await c.req.formData();
  const files = formData.getAll('file') as File[];
  const results: { id?: string; name: string; success: boolean; error?: string }[] = [];

  for (const file of files) {
    try {
      if (!file.type.startsWith('image/')) {
        results.push({ name: file.name, success: false, error: 'Not an image' });
        continue;
      }
      const ext = extname(file.name) || `.${file.type.split('/')[1]}`;
      const filename = `${nanoid(12)}${ext}`;
      const filePath = join(resolvedUploadDir, filename);
      await writeFile(filePath, Buffer.from(await file.arrayBuffer()));
      const canonicalPath = canonicalizeDataPath(filePath);
      const [meme] = await sql`INSERT INTO memes (file_path, title, meta) VALUES (${canonicalPath}, ${file.name.replace(/\.[^/.]+$/, '')}, ${sql.json({ filesize: file.size })}) RETURNING *`;
      results.push({ id: meme.id, name: file.name, success: true });
    } catch (e) {
      results.push({ name: file.name, success: false, error: String(e) });
    }
  }

  return c.json({ uploaded: results.filter(r => r.success).length, failed: results.filter(r => !r.success).length, results });
});

app.post('/api/scan', async (c) => {
  const body = await c.req.json().catch(() => ({}));
  const requestedDir = body.directory || '/data/memes';
  const dir = await resolveDataDir(requestedDir);
  if (!await isDirectory(dir)) {
    return c.json({ error: `Directory not found: ${requestedDir}` }, 400);
  }

  if (scanState.status === 'scanning') {
    return c.json({ error: 'Scan already in progress' }, 409);
  }

  // Run async
  scanDirectory(dir).catch(e => console.error('Scan error:', e));
  return c.json({ success: true, message: `Started scanning ${dir}` });
});

app.get('/api/scan/status', (c) => c.json(scanState));

// =============================================================================
// API: Batch Generate
// =============================================================================

app.post('/api/generate-pending', async (c) => {
  const ollama = await checkOllama();
  if (!ollama.available) return c.json({ error: `Ollama unavailable: ${ollama.error}` }, 400);

  const pending = await sql`SELECT * FROM memes WHERE status = 'pending' LIMIT 1000`;
  let processed = 0, failed = 0;

  for (const meme of pending) {
    try {
      await sql`UPDATE memes SET status = 'processing' WHERE id = ${meme.id}`;
      const description = await describeImage(meme.file_path);
      const embedding = await embed(description);
      await sql`UPDATE memes SET description = ${description}, embedding = ${`[${embedding.join(',')}]`}::vector, status = 'complete' WHERE id = ${meme.id}`;
      processed++;
      console.log(`[${processed}/${pending.length}] Generated: ${meme.id}`);
    } catch (e) {
      await sql`UPDATE memes SET status = 'error' WHERE id = ${meme.id}`;
      failed++;
    }
  }

  return c.json({ processed, failed, total: pending.length });
});

// =============================================================================
// START
// =============================================================================

console.log(`Initializing ${APP_NAME}...`);
if (!SKIP_DB_INIT) {
  await initDb();
} else {
  console.warn('Skipping database init (SKIP_DB_INIT=1).');
}
getEmbedder().catch(() => {}); // Warm up

export default { port: PORT, fetch: app.fetch };
console.log(`${APP_NAME} running on http://localhost:${PORT}`);