#!/usr/bin/env node
/**
 * REXO Ollama Worker
 * Picks tasks from PLAN.json, sends to Ollama, writes output files, marks done.
 * No Claude — uses local Ollama (qwen2.5:7b by default).
 *
 * Usage: node ollama-worker.mjs [--model qwen2.5:7b] [--task T799]
 */

import { readFileSync, writeFileSync, mkdirSync, renameSync, appendFileSync, existsSync } from 'fs';
import { dirname, join } from 'path';
import { fileURLToPath } from 'url';
import { execSync } from 'child_process';

const PLAN_PATH = '/opt/amos/rexo-build/plan/PLAN.json';
const ARTIFACTS_PATH = '/opt/amos/rexo-build/artifacts';
const LOG_PATH = '/opt/amos/rexo-build/logs/ollama-worker.log';
const OLLAMA_URL = 'http://127.0.0.1:11434';
const DEFAULT_MODEL = 'qwen2.5:7b';

// Parse args
const args = process.argv.slice(2);
const modelArg = args[indexOf('--model', args) + 1] || DEFAULT_MODEL;
const taskArg = args[indexOf('--task', args) + 1] || null;

function indexOf(flag, arr) {
  const i = arr.indexOf(flag);
  return i === -1 ? -999 : i;
}

function log(msg) {
  const line = `[${new Date().toISOString()}] ${msg}\n`;
  process.stdout.write(line);
  try { appendFileSync(LOG_PATH, line); } catch {}
}

function readPlan() {
  return JSON.parse(readFileSync(PLAN_PATH, 'utf8'));
}

function writePlanAtomic(plan) {
  const tmp = PLAN_PATH + '.tmp.' + process.pid;
  writeFileSync(tmp, JSON.stringify(plan, null, 2));
  renameSync(tmp, PLAN_PATH);
}

function claimTask(targetId) {
  const plan = readPlan();
  const task = plan.tasks.find(t =>
    targetId ? t.id === targetId : (t.status === 'pending' && t.priority >= 7)
  );
  if (!task || task.status !== 'pending') return null;

  task.status = 'in_progress';
  task.worker = `ollama-${process.pid}`;
  task.started_at = new Date().toISOString();
  writePlanAtomic(plan);
  return task;
}

function markDone(taskId) {
  const plan = readPlan();
  const task = plan.tasks.find(t => t.id === taskId);
  if (task) { task.status = 'done'; task.done_at = new Date().toISOString(); }
  writePlanAtomic(plan);
}

function markFailed(taskId, reason) {
  const plan = readPlan();
  const task = plan.tasks.find(t => t.id === taskId);
  if (task) { task.status = 'failed'; task.error = reason; }
  writePlanAtomic(plan);
}

async function ollamaGenerate(model, prompt) {
  const resp = await fetch(`${OLLAMA_URL}/api/generate`, {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify({ model, prompt, stream: false, options: { num_ctx: 8192, temperature: 0.2 } })
  });
  if (!resp.ok) throw new Error(`Ollama error: ${resp.status} ${await resp.text()}`);
  const data = await resp.json();
  return data.response;
}

function extractCodeBlocks(text) {
  // Extract ```html ... ``` or ```css ... ``` or plain code blocks
  const blocks = [];
  const re = /```(?:html|css|js|javascript|xml|json|bash)?\n([\s\S]*?)```/g;
  let m;
  while ((m = re.exec(text)) !== null) {
    blocks.push(m[1]);
  }
  // If no blocks found, return full text (assume it's raw code)
  return blocks.length ? blocks : [text];
}

function extractFilesFromOutput(output, task) {
  // Look for FILE: /path/to/file.html markers in output
  const files = [];
  const fileRe = /FILE:\s*(\/[^\n]+)\n([\s\S]*?)(?=FILE:|$)/g;
  let m;
  while ((m = fileRe.exec(output)) !== null) {
    files.push({ path: m[1].trim(), content: m[2].trim() });
  }

  // If no FILE: markers — use file_targets from task if single file
  if (!files.length && task.file_targets?.length === 1) {
    const code = extractCodeBlocks(output);
    files.push({ path: task.file_targets[0], content: code[0] });
  }

  return files;
}

function writeFiles(files) {
  for (const { path, content } of files) {
    mkdirSync(dirname(path), { recursive: true });
    writeFileSync(path, content, 'utf8');
    log(`  Wrote ${path} (${content.length} bytes)`);
  }
}

async function runTask(task, model) {
  log(`→ Task ${task.id}: ${task.title}`);

  // Build prompt
  const prompt = `You are a precise code generator. Output ONLY the requested file content with no explanation.

TASK: ${task.title}

${task.description}

RULES:
- Output file content only — no commentary before or after
- If multiple files: use FILE: /full/path/to/file.html marker before each file's content
- For HTML: complete self-contained files, inline CSS, Inter font from Google Fonts
- Colors: #0A2540 navy, #635BFF blue, white background
- Mobile-first, no emoji, SVG icons only
- Never truncate — write complete files

OUTPUT:`;

  log(`  Calling Ollama (${model})...`);
  const output = await ollamaGenerate(model, prompt);
  log(`  Got ${output.length} chars from Ollama`);

  // Extract and write files
  const files = extractFilesFromOutput(output, task);

  if (!files.length) {
    throw new Error('No files extracted from Ollama output');
  }

  writeFiles(files);

  // Write artifact
  const artifactPath = join(ARTIFACTS_PATH, `${task.id}-DONE.md`);
  const artifact = `# ${task.id} DONE\n\nTask: ${task.title}\nWorker: ollama/${model}\nCompleted: ${new Date().toISOString()}\n\nFiles written:\n${files.map(f => `- ${f.path}`).join('\n')}\n`;
  writeFileSync(artifactPath, artifact);

  return files;
}

// Main
async function main() {
  log(`Ollama Worker starting (model: ${modelArg}, task: ${taskArg || 'auto'})`);

  // Check Ollama is up
  try {
    const r = await fetch(`${OLLAMA_URL}/api/tags`);
    const d = await r.json();
    const models = d.models?.map(m => m.name) || [];
    log(`Ollama models available: ${models.join(', ')}`);
    if (!models.some(m => m.startsWith(modelArg.split(':')[0]))) {
      log(`WARNING: model ${modelArg} not found — pulling...`);
      execSync(`ollama pull ${modelArg}`, { stdio: 'inherit' });
    }
  } catch (e) {
    log(`ERROR: Ollama not reachable at ${OLLAMA_URL}: ${e.message}`);
    process.exit(1);
  }

  // Claim task
  const task = claimTask(taskArg);
  if (!task) {
    log('No pending tasks found. Exiting.');
    process.exit(0);
  }

  try {
    const files = await runTask(task, modelArg);
    markDone(task.id);
    log(`✅ ${task.id} done — ${files.length} file(s) written`);
  } catch (e) {
    log(`❌ ${task.id} failed: ${e.message}`);
    markFailed(task.id, e.message);

    // Write blocked artifact
    const blockedPath = join(ARTIFACTS_PATH, `${task.id}-BLOCKED.md`);
    writeFileSync(blockedPath, `# ${task.id} BLOCKED\n\nReason: ${e.message}\nWorker: ollama/${modelArg}\nTime: ${new Date().toISOString()}\n`);
    process.exit(1);
  }
}

main().catch(e => { log(`FATAL: ${e.message}`); process.exit(1); });