hexifyer/pipeline/import-from-image.ts

/**
 * One-time hex map terrain import from a source image using Claude vision.
 *
 * Usage:
 *   npx tsx pipeline/import-from-image.ts <image> <map-id> [options]
 *
 * Options:
 *   --model haiku|sonnet   Vision model (default: sonnet)
 *   --dry-run              Classify without writing to DB
 *   --save-every <n>       Persist DB every N hexes (default: 50)
 *
 * Processing order: column-by-column (q ascending, r ascending within q),
 * so NW and W neighbours are always already classified when a hex is reached.
 */

import Anthropic from '@anthropic-ai/sdk';
import sharp from 'sharp';
import { resolve, dirname } from 'path';
import { fileURLToPath } from 'url';
import { initDb, getDb, saveDb } from '../server/db.js';
import { axialToPixel, hexVertices } from '../core/coords.js';
import { gridBoundsForImage } from '../core/hex-grid.js';
import { TERRAIN_TYPES } from '../core/terrain.js';
import { HexEdge, EDGE_DIRECTIONS, ALL_EDGES, type AxialCoord, type PixelCoord } from '../core/types.js';

// ─── Model config ────────────────────────────────────────────────────────────

const MODELS = {
  haiku:  'claude-haiku-4-5-20251001',
  sonnet: 'claude-sonnet-4-6',
} as const;

type ModelKey = keyof typeof MODELS;

// ─── Edge mask helpers ───────────────────────────────────────────────────────

const EDGE_BIT: Record<HexEdge, number> = {
  [HexEdge.NE]: 1,
  [HexEdge.E]:  2,
  [HexEdge.SE]: 4,
  [HexEdge.SW]: 8,
  [HexEdge.W]:  16,
  [HexEdge.NW]: 32,
};

const EDGE_NAME: Record<HexEdge, string> = {
  [HexEdge.NE]: 'NE',
  [HexEdge.E]:  'E',
  [HexEdge.SE]: 'SE',
  [HexEdge.SW]: 'SW',
  [HexEdge.W]:  'W',
  [HexEdge.NW]: 'NW',
};

// ─── Submap extraction ───────────────────────────────────────────────────────

const CROP_RADIUS_FACTOR = 2.8; // multiplied by hexSize → radius of crop around center

interface Crop {
  left: number;
  top: number;
  width: number;
  height: number;
}

function computeCrop(
  cx: number,
  cy: number,
  hexSize: number,
  imageWidth: number,
  imageHeight: number,
): Crop {
  const r = Math.ceil(hexSize * CROP_RADIUS_FACTOR);
  const left   = Math.max(0, Math.round(cx - r));
  const top    = Math.max(0, Math.round(cy - r));
  const right  = Math.min(imageWidth,  Math.round(cx + r));
  const bottom = Math.min(imageHeight, Math.round(cy + r));
  return { left, top, width: right - left, height: bottom - top };
}

/** Build SVG overlay: center hex in red, neighbours in grey. */
function buildHexOverlaySvg(
  center: AxialCoord,
  hexSize: number,
  origin: PixelCoord,
  crop: Crop,
): string {
  const hexCoords = [center, ...ALL_EDGES.map(e => {
    const d = EDGE_DIRECTIONS[e];
    return { q: center.q + d.q, r: center.r + d.r };
  })];

  const polys = hexCoords.map((coord, i) => {
    const px = axialToPixel(coord, hexSize, origin);
    const verts = hexVertices(px.x - crop.left, px.y - crop.top, hexSize);
    const pts = verts.map(v => `${v.x.toFixed(1)},${v.y.toFixed(1)}`).join(' ');
    const isCenter = i === 0;
    return `<polygon points="${pts}" fill="${isCenter ? 'rgba(255,68,68,0.08)' : 'none'}" `
      + `stroke="${isCenter ? '#ff4444' : '#cccccc'}" `
      + `stroke-width="${isCenter ? 3 : 1.5}" stroke-opacity="0.9"/>`;
  });

  return `<svg width="${crop.width}" height="${crop.height}" xmlns="http://www.w3.org/2000/svg">\n`
    + polys.join('\n') + '\n</svg>';
}

async function extractSubmap(
  imagePath: string,
  center: AxialCoord,
  hexSize: number,
  origin: PixelCoord,
  imageWidth: number,
  imageHeight: number,
): Promise<string> {
  const px = axialToPixel(center, hexSize, origin);
  const crop = computeCrop(px.x, px.y, hexSize, imageWidth, imageHeight);
  const svg = buildHexOverlaySvg(center, hexSize, origin, crop);

  const buf = await sharp(imagePath)
    .extract(crop)
    .composite([{ input: Buffer.from(svg), top: 0, left: 0 }])
    .png()
    .toBuffer();

  return buf.toString('base64');
}

// ─── Neighbour context ───────────────────────────────────────────────────────

interface ClassifiedHex {
  base: string;
  features: Array<{ terrainId: string; edgeMask: number }>;
}

function buildNeighbourContext(
  center: AxialCoord,
  classified: Map<string, ClassifiedHex>,
): string {
  const lines: string[] = [];
  for (const edge of ALL_EDGES) {
    const d = EDGE_DIRECTIONS[edge];
    const key = `${center.q + d.q},${center.r + d.r}`;
    const name = EDGE_NAME[edge];
    const result = classified.get(key);
    if (result) {
      let desc = result.base;
      if (result.features.length > 0) {
        const feats = result.features.map(f => {
          const exitEdges = ALL_EDGES
            .filter(e => f.edgeMask & EDGE_BIT[e])
            .map(e => EDGE_NAME[e])
            .join('+');
          return `${f.terrainId}(edges:${exitEdges})`;
        }).join(', ');
        desc += ` + ${feats}`;
      }
      lines.push(`  ${name}: ${desc}`);
    } else {
      lines.push(`  ${name}: unknown`);
    }
  }
  return lines.join('\n');
}

// ─── Claude tool definition ──────────────────────────────────────────────────

const AREA_IDS   = TERRAIN_TYPES.filter(t => t.category === 'area').map(t => t.id);
const LINEAR_IDS = TERRAIN_TYPES.filter(t => t.category === 'linear').map(t => t.id);

const classifyTool: Anthropic.Tool = {
  name: 'classify_hex',
  description: 'Classify the terrain of the center hex (red outline) in the map image.',
  input_schema: {
    type: 'object' as const,
    properties: {
      base: {
        type: 'string',
        enum: AREA_IDS,
        description: 'Primary area terrain filling the center hex.',
      },
      features: {
        type: 'array',
        items: {
          type: 'object' as const,
          properties: {
            terrainId: { type: 'string', enum: LINEAR_IDS },
            edgeMask: {
              type: 'integer',
              minimum: 0,
              maximum: 63,
              description: '6-bit mask of edges the feature crosses. '
                + 'Bits: NE=1, E=2, SE=4, SW=8, W=16, NW=32.',
            },
          },
          required: ['terrainId', 'edgeMask'],
        },
        description: 'Linear features (rivers, roads, coastlines) passing through the hex. '
          + 'Omit if none visible.',
      },
      reasoning: {
        type: 'string',
        description: 'One sentence explaining the classification.',
      },
    },
    required: ['base', 'features', 'reasoning'],
  },
};

const SYSTEM_PROMPT = `You are a fantasy cartography expert classifying hexagonal map regions.
You will receive a cropped section of a hand-drawn or digitally painted fantasy map.
The CENTER hex is outlined in RED. Surrounding hexes are outlined in grey for context.

Terrain types (area):
${TERRAIN_TYPES.filter(t => t.category === 'area')
  .map(t => `  ${t.id}: ${t.name}`).join('\n')}

Linear features (can overlay area terrain):
${TERRAIN_TYPES.filter(t => t.category === 'linear')
  .map(t => `  ${t.id}: ${t.name} — use edgeMask to indicate which hex edges it crosses`).join('\n')}

Edge mask bits: NE=1, E=2, SE=4, SW=8, W=16, NW=32.
A river entering from the W edge and exiting via the E edge → edgeMask = 18 (W|E = 16+2).

Focus on the CENTER hex. Use neighbour context only to infer continuity of rivers/roads.`;

// ─── API call ────────────────────────────────────────────────────────────────

async function classifyHex(
  client: Anthropic,
  model: string,
  submapBase64: string,
  neighborContext: string,
): Promise<ClassifiedHex> {
  const userText = `Classify the center hex (red outline).\n\nNeighbour terrain:\n${neighborContext}`;

  const response = await client.messages.create({
    model,
    max_tokens: 512,
    system: SYSTEM_PROMPT,
    tools: [classifyTool],
    tool_choice: { type: 'any' },
    messages: [{
      role: 'user',
      content: [
        {
          type: 'image',
          source: { type: 'base64', media_type: 'image/png', data: submapBase64 },
        },
        { type: 'text', text: userText },
      ],
    }],
  });

  const toolUse = response.content.find(b => b.type === 'tool_use');
  if (!toolUse || toolUse.type !== 'tool_use') {
    throw new Error('No tool_use block in response');
  }

  const input = toolUse.input as { base: string; features: any[]; reasoning: string };
  return { base: input.base, features: input.features ?? [] };
}

// ─── DB helpers ─────────────────────────────────────────────────────────────

interface MapConfig {
  id: number;
  image_width: number;
  image_height: number;
  hex_size: number;
  origin_x: number;
  origin_y: number;
}

function loadMapConfig(mapId: number): MapConfig {
  const db = getDb();
  const rows = db.exec('SELECT id, image_width, image_height, hex_size, origin_x, origin_y FROM hex_maps WHERE id = ?', [mapId]);
  if (rows.length === 0 || rows[0].values.length === 0) {
    throw new Error(`Map ${mapId} not found`);
  }
  const [id, image_width, image_height, hex_size, origin_x, origin_y] = rows[0].values[0] as number[];
  return { id, image_width, image_height, hex_size, origin_x, origin_y };
}

function writeHex(mapId: number, coord: AxialCoord, result: ClassifiedHex): void {
  const db = getDb();
  db.run(
    `INSERT INTO hexes (map_id, q, r, base_terrain, updated_at)
     VALUES (?, ?, ?, ?, datetime('now'))
     ON CONFLICT(map_id, q, r)
     DO UPDATE SET base_terrain = excluded.base_terrain, updated_at = datetime('now')`,
    [mapId, coord.q, coord.r, result.base],
  );
  const idRows = db.exec('SELECT id FROM hexes WHERE map_id = ? AND q = ? AND r = ?', [mapId, coord.q, coord.r]);
  const hexId = idRows[0].values[0][0] as number;
  db.run('DELETE FROM hex_features WHERE hex_id = ?', [hexId]);
  for (const f of result.features) {
    if (f.edgeMask === 0) continue;
    db.run('INSERT INTO hex_features (hex_id, terrain_id, edge_mask) VALUES (?, ?, ?)', [hexId, f.terrainId, f.edgeMask]);
  }
}

// ─── Processing order ────────────────────────────────────────────────────────

/** Sort axial coords: q ascending, then r ascending within each q column. */
function sortCoords(coords: AxialCoord[]): AxialCoord[] {
  return [...coords].sort((a, b) => a.q !== b.q ? a.q - b.q : a.r - b.r);
}

// ─── Main ────────────────────────────────────────────────────────────────────

async function main() {
  const args = process.argv.slice(2);

  const imagePath = args[0];
  const mapId     = parseInt(args[1], 10);

  if (!imagePath || isNaN(mapId)) {
    console.error('Usage: npx tsx pipeline/import-from-image.ts <image> <map-id> [--model haiku|sonnet] [--dry-run] [--save-every n]');
    process.exit(1);
  }

  const modelKey  = (args.includes('--model') ? args[args.indexOf('--model') + 1] : 'sonnet') as ModelKey;
  const dryRun    = args.includes('--dry-run');
  const saveEvery = args.includes('--save-every') ? parseInt(args[args.indexOf('--save-every') + 1], 10) : 50;

  if (!(modelKey in MODELS)) {
    console.error(`Unknown model "${modelKey}". Use: haiku or sonnet`);
    process.exit(1);
  }

  const model = MODELS[modelKey];
  console.log(`Model:   ${modelKey} (${model})`);
  console.log(`Dry run: ${dryRun}`);

  await initDb();
  const mapConfig = loadMapConfig(mapId);
  const { image_width, image_height, hex_size, origin_x, origin_y } = mapConfig;
  const origin: PixelCoord = { x: origin_x, y: origin_y };

  console.log(`Map ${mapId}: ${image_width}×${image_height}px, hexSize=${hex_size}, origin=(${origin_x},${origin_y})`);

  const { coords } = gridBoundsForImage(image_width, image_height, hex_size, origin);
  const sorted = sortCoords(coords);
  console.log(`Total hexes: ${sorted.length}`);

  const client = new Anthropic();
  const classified = new Map<string, ClassifiedHex>();
  let done = 0;
  let errors = 0;

  for (const coord of sorted) {
    const key = `${coord.q},${coord.r}`;
    done++;

    const neighborContext = buildNeighbourContext(coord, classified);

    let result: ClassifiedHex;
    try {
      const submap = await extractSubmap(imagePath, coord, hex_size, origin, image_width, image_height);
      result = await classifyHex(client, model, submap, neighborContext);
    } catch (err) {
      console.error(`  [${done}/${sorted.length}] (${coord.q},${coord.r}) ERROR: ${err}`);
      errors++;
      result = { base: 'plains', features: [] };
    }

    classified.set(key, result);

    const featureStr = result.features.length > 0
      ? ` + ${result.features.map(f => f.terrainId).join(', ')}`
      : '';
    console.log(`  [${done}/${sorted.length}] (${coord.q},${coord.r}) → ${result.base}${featureStr}`);

    if (!dryRun) {
      writeHex(mapId, coord, result);
      if (done % saveEvery === 0) {
        saveDb();
        console.log(`  [saved at ${done}]`);
      }
    }
  }

  if (!dryRun) {
    saveDb();
    console.log(`\nDone. ${done} hexes written, ${errors} errors.`);
  } else {
    console.log(`\nDry run complete. ${done} hexes classified (not written), ${errors} errors.`);
  }
}

main().catch(err => {
  console.error('Fatal:', err);
  process.exit(1);
});