import express from 'express';
import { spawn, execSync } from 'child_process';
import crypto from 'crypto';
import fs from 'fs';
import os from 'os';
import path from 'path';
import { fileURLToPath } from 'url';

const __dirname = path.dirname(fileURLToPath(import.meta.url));

const app = express();
const PORT = process.env.PORT || 7860;
const SECRET_KEY = process.env.SECRET_KEY;

if (!SECRET_KEY) {
  console.error('[ERROR] SECRET_KEY environment variable is not set. Exiting.');
  process.exit(1);
}

app.use(express.json());
app.use(express.static(path.join(__dirname, 'public')));

// ── Session store ─────────────────────────────────────────────────────────────
const sessions = new Set();

// ── ANSI stripper ─────────────────────────────────────────────────────────────
function stripAnsi(str) {
  // eslint-disable-next-line no-control-regex
  return str.replace(/[\u001b\u009b][[()#;?]*(?:[0-9]{1,4}(?:;[0-9]{0,4})*)?[0-9A-ORZcf-nqry=><~]/g, '');
}

// ── Pre-seed shellular config from env vars ───────────────────────────────────
function seedShellularConfig() {
  const hostId    = process.env.SHELLULAR_HOST_ID;
  const keyB64    = process.env.SHELLULAR_KEY;
  const machineId = process.env.SHELLULAR_MACHINE_ID;

  if (!hostId || !keyB64 || !machineId) return;

  const shellularDir = path.join(os.homedir(), '.shellular');
  const configFile   = path.join(shellularDir, 'config.json');
  const keyFile      = path.join(shellularDir, `shellular-${machineId}.e2ee`);

  try {
    fs.mkdirSync(shellularDir, { recursive: true });
    if (!fs.existsSync(configFile)) {
      fs.writeFileSync(configFile, JSON.stringify({ hostId, machineId }), 'utf-8');
      console.log(`[shellular] seeded config: hostId=${hostId}`);
    }
    if (!fs.existsSync(keyFile)) {
      fs.writeFileSync(keyFile, Buffer.from(keyB64, 'base64'), { mode: 0o600 });
      console.log(`[shellular] seeded key: ${keyFile}`);
    }
  } catch (err) {
    console.error('[shellular] failed to seed config:', err.message);
  }
}

seedShellularConfig();

// ── Shellular fetch patch ─────────────────────────────────────────────────────
// WHY THIS EXISTS:
//   The shellular registration API returns `success` as a non-boolean (e.g. the
//   string "true" or "success") but shellular's Zod schema uses a discriminated
//   union that only accepts the strict booleans true | false — so validation
//   crashes with "invalid_union" before the QR is ever generated.
//
// WHY globalThis.fetch PATCHING DOESN'T WORK:
//   Shellular imports undici directly (e.g. `import { fetch } from 'undici'`)
//   rather than using globalThis.fetch, so patching the global has no effect.
//
// THE FIX — patch at the undici global dispatcher level:
//   Every HTTP call undici makes — regardless of which API (fetch, request,
//   stream, pipeline) — goes through the global dispatcher. By replacing it
//   we intercept the raw response bytes before Zod ever sees them and coerce
//   the `success` field to a proper boolean.
// Write the patch file INSIDE shellular's installation directory so that
// `import 'undici'` resolves correctly via Node's ESM node_modules lookup.
// Writing to /tmp/ fails because there is no node_modules there.
function findShellularRoot() {
  try {
    const globalRoot = execSync('npm root -g', { encoding: 'utf-8' }).trim();
    const p = path.join(globalRoot, 'shellular');
    if (fs.existsSync(p)) return p;
  } catch {}
  // Fallback to the path set in the Dockerfile
  return '/home/node/.npm-global/lib/node_modules/shellular';
}
const PATCH_FILE = path.join(findShellularRoot(), 'fetch-patch.mjs');

function writeFetchPatch() {
  const code = `
import { getGlobalDispatcher, setGlobalDispatcher } from 'undici';

const _orig = getGlobalDispatcher();

// ── Handler wrapper — buffers JSON responses and patches the success field ───
class SuccessPatchingHandler {
  constructor(inner) {
    this.inner  = inner;
    this.isJson = false;
    this.chunks = [];
  }
  onConnect(abort)  { return this.inner.onConnect(abort); }
  onError(err)      { return this.inner.onError(err); }
  onUpgrade(...a)   { return this.inner.onUpgrade?.(...a); }
  onBodySent(chunk) { return this.inner.onBodySent?.(chunk); }

  onHeaders(statusCode, headers, resume, statusMessage) {
    // headers is a flat Buffer array: [name, value, name, value …]
    for (let i = 0; i < headers.length; i += 2) {
      if (
        headers[i].toString().toLowerCase() === 'content-type' &&
        headers[i + 1].toString().includes('application/json')
      ) {
        this.isJson = true;
      }
    }
    return this.inner.onHeaders(statusCode, headers, resume, statusMessage);
  }

  onData(chunk) {
    if (this.isJson) {
      this.chunks.push(chunk);
      return true; // signal undici to keep reading
    }
    return this.inner.onData(chunk);
  }

  onComplete(trailers) {
    if (this.isJson && this.chunks.length) {
      const raw = Buffer.concat(this.chunks).toString('utf-8');
      try {
        const json = JSON.parse(raw);
        if ('success' in json && typeof json.success !== 'boolean') {
          const s = json.success;
          const before = s;
          // Coerce: anything truthy except the strings "0" / "false" → true
          json.success = (
            s === true || s === 1 ||
            (typeof s === 'string' && s.length > 0 && s !== '0' && s !== 'false')
          );
          console.log('[patch] coerced success field:', JSON.stringify(before), '→', json.success);
          this.inner.onData(Buffer.from(JSON.stringify(json)));
        } else {
          // Nothing to patch — pass through as-is
          this.inner.onData(Buffer.concat(this.chunks));
        }
      } catch (e) {
        // Not valid JSON — pass through unchanged
        this.inner.onData(Buffer.concat(this.chunks));
      }
    }
    return this.inner.onComplete(trailers);
  }
}

// ── Dispatcher wrapper — routes every request through the patching handler ───
class SuccessPatchingDispatcher {
  dispatch(opts, handler) {
    return _orig.dispatch(opts, new SuccessPatchingHandler(handler));
  }
  close()          { return _orig.close?.(); }
  destroy(err, cb) { return _orig.destroy?.(err, cb); }
}

setGlobalDispatcher(new SuccessPatchingDispatcher());
console.log('[patch] undici dispatcher patched — success field will be coerced to boolean');
`;
  fs.writeFileSync(PATCH_FILE, code, 'utf-8');
  console.log('[patch] undici dispatcher patch written to', PATCH_FILE);
}

writeFetchPatch();

// ── Auth routes ───────────────────────────────────────────────────────────────
app.post('/api/login', (req, res) => {
  const { key } = req.body;
  if (!key || key !== SECRET_KEY) {
    return res.status(401).json({ error: 'Invalid key' });
  }
  const token = crypto.randomUUID();
  sessions.add(token);
  res.json({ token });
});

app.post('/api/logout', (req, res) => {
  const token = req.headers.authorization?.split(' ')[1];
  sessions.delete(token);
  res.json({ ok: true });
});

// ── Auth middleware ────────────────────────────────────────────────────────────
function requireAuth(req, res, next) {
  const token = req.headers.authorization?.split(' ')[1];
  if (!token || !sessions.has(token)) {
    return res.status(401).json({ error: 'Unauthorized' });
  }
  next();
}

// ── Shellular process management ──────────────────────────────────────────────
let shellularProc = null;
let outputBuffer  = '';
const sseClients  = new Set();

function send(res, payload) {
  res.write(`data: ${JSON.stringify(payload)}\n\n`);
}

function broadcast(payload) {
  const frame = `data: ${JSON.stringify(payload)}\n\n`;
  for (const client of sseClients) {
    client.write(frame);
  }
}

let retryTimer = null;

function startShellular() {
  if (shellularProc || retryTimer) return;

  broadcast({ type: 'status', status: 'starting' });

  shellularProc = spawn('shellular', ['--unknown-clients', 'always-allow'], {
    env: {
      ...process.env,
      FORCE_COLOR: '0',
      // Preload the undici dispatcher patch BEFORE shellular initialises.
      // --import runs the ESM module first, so the patched dispatcher is in
      // place before any of shellular's own imports execute.
      NODE_OPTIONS: `--import file://${PATCH_FILE}`,
    },
    stdio: ['ignore', 'pipe', 'pipe'],
  });

  let procOutput = '';

  const handleData = (chunk) => {
    const text = stripAnsi(chunk.toString());
    procOutput   += text;
    outputBuffer += text;
    broadcast({ type: 'output', text });
  };

  shellularProc.stdout.on('data', handleData);
  shellularProc.stderr.on('data', handleData);

  shellularProc.on('error', (err) => {
    // ── DIAGNOSTIC ERROR TYPE: spawn-error ──────────────────────────────────
    // Means the shellular binary could not be found or could not be executed.
    // Most likely cause: npm install -g shellular failed in the Dockerfile,
    // or PATH is missing /home/node/.npm-global/bin.
    // To debug: run `which shellular` and `shellular --version` in the container.
    const text = [
      '',
      '╔══════════════════════════════════════════════════════════╗',
      '║  DIAGNOSTIC: spawn-error                                  ║',
      '╠══════════════════════════════════════════════════════════╣',
      `║  Cannot start the shellular binary: ${err.message}`,
      '║  Likely cause: shellular not installed or PATH wrong.     ║',
      '║  Fix: rebuild the Space (Dockerfile installs shellular).  ║',
      '╚══════════════════════════════════════════════════════════╝',
      '',
    ].join('\n');
    outputBuffer += text;
    broadcast({ type: 'output', text });
    shellularProc = null;
    broadcast({ type: 'status', status: 'error' });
  });

  shellularProc.on('exit', (code, signal) => {
    shellularProc = null;

    // ── Classify the exit reason for clear diagnostics ────────────────────
    const isSchemaError = code === 1 && !signal &&
      procOutput.includes('invalid_union');

    const isPatchBypassed = code === 1 && !signal &&
      procOutput.includes('invalid_union') &&
      !procOutput.includes('[patch] coerced success field');
    // ^ If you see invalid_union BUT "[patch] coerced" never printed, the
    //   undici patch did not load. Check NODE_OPTIONS in the spawn call above.

    const isRateLimit = code === 1 && !signal && !isSchemaError &&
      (procOutput.includes('Too many requests') || procOutput.includes('rate'));

    const isAuthError = code === 1 && !signal &&
      (procOutput.includes('401') || procOutput.includes('Unauthorized') ||
       procOutput.includes('authentication'));

    const isNetworkError = code === 1 && !signal &&
      (procOutput.includes('ECONNREFUSED') || procOutput.includes('ENOTFOUND') ||
       procOutput.includes('ETIMEDOUT') || procOutput.includes('fetch failed'));

    if (isSchemaError) {
      // ── DIAGNOSTIC ERROR TYPE: schema-mismatch ───────────────────────────
      // The shellular registration API returned a response whose shape does
      // not match shellular's Zod schema.
      //
      // If "[patch] coerced success field" DID print above this message,
      // the undici patch ran but the API changed in a different field — look
      // at the raw error above to see which field/path Zod rejected.
      //
      // If "[patch] coerced success field" did NOT print, the patch never
      // ran — NODE_OPTIONS may have been stripped or the patch file missing.
      const patchRan = procOutput.includes('[patch] coerced success field');
      const msg = [
        '',
        '╔══════════════════════════════════════════════════════════╗',
        '║  DIAGNOSTIC: schema-mismatch (invalid_union)              ║',
        '╠══════════════════════════════════════════════════════════╣',
        `║  Undici patch loaded: ${patchRan ? 'YES ✅' : 'NO ❌ — patch never ran!'}`,
        '║',
        patchRan
          ? '║  The API changed a DIFFERENT field than "success".'
          : '║  NODE_OPTIONS may have been stripped. Check spawn env.',
        '║',
        '║  What to do:',
        '║  1. Copy the full output above and share it.',
        '║  2. Look for the "path" field in the [error] JSON —',
        '║     it tells you exactly which field Zod rejected.',
        '║  3. Rebuild the Space to pull shellular@latest.',
        '║  4. Workaround: set SHELLULAR_HOST_ID + SHELLULAR_KEY +',
        '║     SHELLULAR_MACHINE_ID as HF Secrets (skips registration).',
        '╚══════════════════════════════════════════════════════════╝',
        '  Retrying in 5 minutes…',
        '',
      ].join('\n');
      outputBuffer += msg;
      broadcast({ type: 'output', text: msg });
      broadcast({ type: 'status', status: 'retrying' });

      retryTimer = setTimeout(() => {
        retryTimer = null;
        const m = '\n[Retrying after schema error…]\n';
        outputBuffer += m;
        broadcast({ type: 'output', text: m });
        startShellular();
      }, 5 * 60 * 1000);

    } else if (isRateLimit) {
      // ── DIAGNOSTIC ERROR TYPE: rate-limited ──────────────────────────────
      // api.shellular.dev rejected registration because this host or IP has
      // registered too many times in a short window.
      // Fix: save the SHELLULAR_* secrets so registration is skipped entirely.
      const WAIT = 30;
      const msg = [
        '',
        '╔══════════════════════════════════════════════════════════╗',
        '║  DIAGNOSTIC: rate-limited                                 ║',
        '╠══════════════════════════════════════════════════════════╣',
        '║  api.shellular.dev refused registration (too many calls). ║',
        '║  Fix: save SHELLULAR_HOST_ID + SHELLULAR_KEY +            ║',
        '║       SHELLULAR_MACHINE_ID as HF Secrets.                 ║',
        `║  Retrying in ${WAIT}s…                                        ║`,
        '╚══════════════════════════════════════════════════════════╝',
        '',
      ].join('\n');
      outputBuffer += msg;
      broadcast({ type: 'output', text: msg });
      broadcast({ type: 'status', status: 'retrying' });

      retryTimer = setTimeout(() => {
        retryTimer = null;
        const m = '\n[Retrying registration…]\n';
        outputBuffer += m;
        broadcast({ type: 'output', text: m });
        startShellular();
      }, WAIT * 1000);

    } else if (isNetworkError) {
      // ── DIAGNOSTIC ERROR TYPE: network-error ─────────────────────────────
      // shellular could not reach api.shellular.dev at all.
      // Could be a HF Space outbound network restriction or a DNS issue.
      const msg = [
        '',
        '╔══════════════════════════════════════════════════════════╗',
        '║  DIAGNOSTIC: network-error                                ║',
        '╠══════════════════════════════════════════════════════════╣',
        '║  shellular could not reach api.shellular.dev.             ║',
        '║  Possible causes:                                         ║',
        '║  - Outbound network blocked in this HF Space tier.        ║',
        '║  - api.shellular.dev is down.                             ║',
        '║  - DNS resolution failing inside the container.           ║',
        '║  Retrying in 30s…                                         ║',
        '╚══════════════════════════════════════════════════════════╝',
        '',
      ].join('\n');
      outputBuffer += msg;
      broadcast({ type: 'output', text: msg });
      broadcast({ type: 'status', status: 'retrying' });

      retryTimer = setTimeout(() => {
        retryTimer = null;
        startShellular();
      }, 30 * 1000);

    } else if (isAuthError) {
      // ── DIAGNOSTIC ERROR TYPE: auth-error ────────────────────────────────
      // The SHELLULAR_* secrets saved in HF are invalid or expired.
      // The host identity on api.shellular.dev may have been reset.
      // Fix: delete the three SHELLULAR_* secrets from HF, restart the Space
      // to re-register fresh, then save the new values from the setup panel.
      const msg = [
        '',
        '╔══════════════════════════════════════════════════════════╗',
        '║  DIAGNOSTIC: auth-error                                   ║',
        '╠══════════════════════════════════════════════════════════╣',
        '║  Saved SHELLULAR_* secrets were rejected by the API.      ║',
        '║  Fix: delete SHELLULAR_HOST_ID + SHELLULAR_KEY +          ║',
        '║       SHELLULAR_MACHINE_ID from HF Secrets, then restart  ║',
        '║       to re-register and get fresh values.                ║',
        '╚══════════════════════════════════════════════════════════╝',
        '',
      ].join('\n');
      outputBuffer += msg;
      broadcast({ type: 'output', text: msg });
      broadcast({ type: 'status', status: 'error' });

    } else {
      // ── DIAGNOSTIC ERROR TYPE: unknown-exit ──────────────────────────────
      // shellular exited for an unrecognised reason.
      // Share the full output above (especially any [error] JSON lines)
      // to diagnose further.
      const text = code !== 0
        ? `\n[shellular exited — code=${code ?? '?'}, signal=${signal ?? 'none'}]\n`
        : '\n[shellular disconnected]\n';
      outputBuffer += text;
      broadcast({ type: 'output', text });
      broadcast({ type: 'status', status: 'stopped' });
    }
  });

  broadcast({ type: 'status', status: 'running' });
}

function stopShellular() {
  if (retryTimer) { clearTimeout(retryTimer); retryTimer = null; }
  if (!shellularProc) return;
  shellularProc.kill('SIGTERM');
  shellularProc = null;
}

// ── SSE stream ─────────────────────────────────────────────────────────────────
app.get('/api/stream', requireAuth, (req, res) => {
  res.setHeader('Content-Type', 'text/event-stream');
  res.setHeader('Cache-Control', 'no-cache');
  res.setHeader('Connection', 'keep-alive');
  res.setHeader('X-Accel-Buffering', 'no');
  res.flushHeaders();

  send(res, { type: 'status', status: shellularProc ? 'running' : 'stopped' });

  if (outputBuffer) {
    send(res, { type: 'output', text: outputBuffer });
  }

  sseClients.add(res);
  req.on('close', () => sseClients.delete(res));
});

// ── Control endpoints ──────────────────────────────────────────────────────────
app.post('/api/shellular/start', requireAuth, (_req, res) => {
  startShellular();
  res.json({ ok: true, running: !!shellularProc });
});

app.post('/api/shellular/stop', requireAuth, (_req, res) => {
  stopShellular();
  outputBuffer = '';
  broadcast({ type: 'output', text: '' });
  res.json({ ok: true });
});

app.post('/api/shellular/restart', requireAuth, (_req, res) => {
  stopShellular();
  outputBuffer = '';
  broadcast({ type: 'clear' });
  setTimeout(startShellular, 600);
  res.json({ ok: true });
});

app.get('/api/status', requireAuth, (_req, res) => {
  res.json({ running: !!shellularProc });
});

app.get('/api/setup-status', requireAuth, (_req, res) => {
  const seeded = !!(
    process.env.SHELLULAR_HOST_ID &&
    process.env.SHELLULAR_KEY &&
    process.env.SHELLULAR_MACHINE_ID
  );
  res.json({ seeded });
});

app.get('/api/shellular/credentials', requireAuth, (_req, res) => {
  try {
    const shellularDir = path.join(os.homedir(), '.shellular');
    const configRaw    = fs.readFileSync(path.join(shellularDir, 'config.json'), 'utf-8');
    const { hostId, machineId } = JSON.parse(configRaw);
    const keyFile = path.join(shellularDir, `shellular-${machineId}.e2ee`);
    const keyB64  = fs.readFileSync(keyFile).toString('base64');
    res.json({ hostId, machineId, keyB64 });
  } catch {
    res.status(404).json({ error: 'Not registered yet.' });
  }
});

app.get('/api/shellular/qr-data', requireAuth, (_req, res) => {
  try {
    const shellularDir = path.join(os.homedir(), '.shellular');
    const configRaw    = fs.readFileSync(path.join(shellularDir, 'config.json'), 'utf-8');
    const { hostId, machineId } = JSON.parse(configRaw);
    const keyFile = path.join(shellularDir, `shellular-${machineId}.e2ee`);
    const keyB64  = fs.readFileSync(keyFile).toString('base64');
    res.json({ qrData: `${hostId}:${keyB64}` });
  } catch {
    res.status(404).json({ error: 'Config not seeded yet.' });
  }
});

// ── Start ──────────────────────────────────────────────────────────────────────
app.listen(PORT, '0.0.0.0', () => {
  console.log(`Shellular Web UI → http://0.0.0.0:${PORT}`);
});