import type { ExtensionAPI } from '@earendil-works/pi-coding-agent';
import { Type } from 'typebox';
import { createHash, randomUUID } from 'node:crypto';
import { readFileSync, appendFileSync, statSync } from 'node:fs';
import net from 'node:net';
import { fileURLToPath } from 'node:url';
import { resolve } from 'node:path';

const SELF = resolve(fileURLToPath(import.meta.url).split('?')[0]!);
const LOG = process.env['GATE0_PI_LOG'];
const CONTEXT_BLOCK =
  process.env['GATE0_PI_CONTEXT_BLOCK'] ??
  [
    'GATE0_PI_ATOMIC_BEGIN',
    'segment-01=alpha-7e31',
    'segment-02=middle-9c42',
    'segment-03=omega-5b83',
    'GATE0_PI_ATOMIC_END',
  ].join('\n');

let sequence = 0;

function sha(value: string | Buffer): string {
  return createHash('sha256').update(value).digest('hex');
}

function procStarttime(): number {
  const text = readFileSync(`/proc/${process.pid}/stat`, 'utf8');
  const close = text.lastIndexOf(')');
  const fields = text.slice(close + 2).trim().split(/\s+/);
  return Number(fields[19]);
}

function log(event: string, details: Record<string, unknown> = {}): void {
  if (!LOG) return;
  sequence += 1;
  appendFileSync(
    LOG,
    `${JSON.stringify({ seq: sequence, event, pid: process.pid, starttime_ticks: procStarttime(), ...details })}\n`,
  );
}

function argvExtensions(): string[] {
  const result: string[] = [];
  for (let i = 0; i < process.argv.length; i += 1) {
    if (process.argv[i] === '--extension' || process.argv[i] === '-e') {
      const candidate = process.argv[i + 1];
      if (candidate) result.push(resolve(candidate));
    }
  }
  return result;
}

interface SourceValidation {
  ok: boolean;
  reason: string;
  fragment?: string;
}

function validateSources(): SourceValidation {
  const manifestPath = process.env['GATE0_SOURCE_MANIFEST'];
  if (!manifestPath) return { ok: true, reason: 'no-manifest-probe-disabled' };
  try {
    const manifest = JSON.parse(readFileSync(manifestPath, 'utf8')) as {
      maxBytes: number;
      fragments: Array<{ path: string; sha256: string }>;
    };
    for (const fragment of manifest.fragments) {
      let fileStat;
      try {
        fileStat = statSync(fragment.path);
      } catch {
        return { ok: false, reason: 'missing', fragment: fragment.path };
      }
      if (!fileStat.isFile()) {
        return { ok: false, reason: 'not-regular-file', fragment: fragment.path };
      }
      if (fileStat.size > manifest.maxBytes) {
        return { ok: false, reason: 'oversize', fragment: fragment.path };
      }
      const bytes = readFileSync(fragment.path);
      if (sha(bytes) !== fragment.sha256) {
        return { ok: false, reason: 'hash-mismatch', fragment: fragment.path };
      }
    }
    return { ok: true, reason: 'all-fragments-valid' };
  } catch (error) {
    return { ok: false, reason: `manifest-error:${error instanceof Error ? error.name : 'unknown'}` };
  }
}

function brokerRequest(payload: Record<string, unknown>): Promise<Record<string, unknown>> {
  const socketPath = process.env['GATE0_GENERATION_SOCKET'];
  if (!socketPath) return Promise.resolve({ skipped: true });
  return new Promise((resolvePromise, reject) => {
    const socket = net.createConnection(socketPath);
    let buffer = '';
    socket.setEncoding('utf8');
    socket.on('connect', () => socket.write(`${JSON.stringify(payload)}\n`));
    socket.on('data', (chunk) => {
      buffer += chunk;
      const newline = buffer.indexOf('\n');
      if (newline < 0) return;
      socket.end();
      resolvePromise(JSON.parse(buffer.slice(0, newline)) as Record<string, unknown>);
    });
    socket.on('error', reject);
  });
}

function markerPaths(value: unknown, path = '$'): string[] {
  const matches: string[] = [];
  if (typeof value === 'string') {
    if (value.includes(CONTEXT_BLOCK)) matches.push(path);
    return matches;
  }
  if (Array.isArray(value)) {
    value.forEach((item, index) => matches.push(...markerPaths(item, `${path}[${index}]`)));
    return matches;
  }
  if (value && typeof value === 'object') {
    for (const [key, item] of Object.entries(value as Record<string, unknown>)) {
      matches.push(...markerPaths(item, `${path}.${key}`));
    }
  }
  return matches;
}

function assistantToolIds(message: unknown): string[] {
  if (!message || typeof message !== 'object') return [];
  const candidate = message as { role?: string; content?: unknown };
  if (candidate.role !== 'assistant' || !Array.isArray(candidate.content)) return [];
  return candidate.content
    .filter(
      (block): block is { type: 'toolCall'; id: string } =>
        Boolean(
          block &&
            typeof block === 'object' &&
            (block as { type?: string }).type === 'toolCall' &&
            typeof (block as { id?: unknown }).id === 'string',
        ),
    )
    .map((block) => block.id);
}

function assistantText(message: unknown): string {
  if (!message || typeof message !== 'object') return '';
  const candidate = message as { role?: string; content?: unknown };
  if (candidate.role !== 'assistant' || !Array.isArray(candidate.content)) return '';
  return candidate.content
    .filter(
      (block): block is { type: 'text'; text: string } =>
        Boolean(
          block &&
            typeof block === 'object' &&
            (block as { type?: string }).type === 'text' &&
            typeof (block as { text?: unknown }).text === 'string',
        ),
    )
    .map((block) => block.text)
    .join('');
}

export default function register(pi: ExtensionAPI) {
  const localProviderUrl = process.env['GATE0_LOCAL_PROVIDER_URL'];
  if (localProviderUrl) {
    pi.registerProvider('gate0-local', {
      baseUrl: localProviderUrl,
      apiKey: 'gate0-probe-not-a-secret',
      api: 'openai-completions',
      models: [
        {
          id: 'gate0-model',
          name: 'Gate0 deterministic local model',
          reasoning: false,
          input: ['text'],
          cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
          contextWindow: 32_000,
          maxTokens: 1_024,
        },
      ],
    });
  }

  const extensions = argvExtensions();
  const lastPosition = extensions.length > 0 && extensions.at(-1) === SELF;
  interface RequestCycle {
    nonce: string;
    verified: boolean;
    sourceReason: string;
  }
  let buildingCycle: RequestCycle | undefined;
  const inFlightCycles: RequestCycle[] = [];
  const toolNonce = new Map<string, { nonce: string; verified: boolean; sourceReason: string }>();

  pi.on('session_start', async (event) => {
    const broker = await brokerRequest({ action: 'lifecycle', phase: 'start', reason: event.reason });
    log('session_start', {
      reason: event.reason,
      extensions,
      self: SELF,
      lastPosition,
      gateState: lastPosition ? 'UNVERIFIED_READY' : 'CLOSED_NOT_LAST',
      broker,
    });
  });

  pi.on('session_shutdown', async (event) => {
    const broker = await brokerRequest({ action: 'lifecycle', phase: 'shutdown', reason: event.reason });
    log('session_shutdown', { reason: event.reason, broker });
  });

  pi.on('context', async (event) => {
    const validation = validateSources();
    buildingCycle = {
      nonce: randomUUID(),
      sourceReason: validation.reason,
      verified: lastPosition && validation.ok,
    };
    const inputJson = JSON.stringify(event.messages);
    const injected = {
      role: 'custom' as const,
      customType: 'gate0-context',
      content: CONTEXT_BLOCK,
      display: false,
      timestamp: Date.now(),
    };
    const outputMessages = buildingCycle.verified
      ? [...event.messages, injected]
      : [...event.messages];
    const outputPrefix = outputMessages.slice(0, event.messages.length);
    const sourceBroker = validation.ok
      ? { action: 'none', reason: 'source-valid' }
      : await brokerRequest({ action: 'source-invalid', reason: validation.reason });
    log('context_return', {
      requestNonce: buildingCycle.nonce,
      sourceValidation: validation,
      sourceBroker,
      lastPosition,
      promotion: false,
      injectionDecision: buildingCycle.verified ? 'ONE_ATOMIC_AGENT_MESSAGE' : 'REFUSED',
      inputCount: event.messages.length,
      outputCount: outputMessages.length,
      prefixHashBefore: sha(inputJson),
      prefixHashAfter: sha(JSON.stringify(outputPrefix)),
      prefixPreservedByReturn: sha(inputJson) === sha(JSON.stringify(outputPrefix)),
      blockLength: CONTEXT_BLOCK.length,
      blockSha256: sha(CONTEXT_BLOCK),
    });
    return { messages: outputMessages };
  });

  pi.on('before_provider_request', async (event) => {
    const paths = markerPaths(event.payload);
    const cycle = buildingCycle;
    buildingCycle = undefined;
    if (cycle) inFlightCycles.push(cycle);
    log('before_provider_request', {
      requestNonce: cycle?.nonce,
      inFlightDepth: inFlightCycles.length,
      markerOccurrences: paths.length,
      markerPaths: paths,
      finalPayloadValid: Boolean(cycle?.verified && paths.length === 1),
    });
  });

  pi.on('after_provider_response', async (event) => {
    const cycle = inFlightCycles[0];
    log('after_provider_response', {
      requestNonce: cycle?.nonce,
      status: event.status,
      assistantContentAvailableAtThisHook: false,
      timing: 'headers/status before stream consumption',
    });
  });

  pi.on('message_end', async (event) => {
    const role = (event.message as { role?: string }).role;
    const ids = assistantToolIds(event.message);
    const text = assistantText(event.message);
    const cycle = role === 'assistant' ? inFlightCycles.shift() : undefined;
    if (ids.length > 0 && cycle) {
      for (const id of ids) {
        toolNonce.set(id, {
          nonce: cycle.nonce,
          verified: cycle.verified,
          sourceReason: cycle.sourceReason,
        });
      }
    }
    log('message_end', {
      role,
      assistantContentObserved: role === 'assistant',
      requestNonce: cycle?.nonce,
      inFlightDepthAfter: inFlightCycles.length,
      toolCallIds: ids,
      nonceMappings: ids.map((id) => ({ toolCallId: id, requestNonce: cycle?.nonce })),
      exactContextBlockCopied: text.includes(CONTEXT_BLOCK),
      assistantTextSha256: text ? sha(text) : null,
    });
  });

  pi.on('tool_call', async (event) => {
    const mapping = toolNonce.get(event.toolCallId);
    const allowed = Boolean(lastPosition && mapping?.verified);
    log('tool_call', {
      toolCallId: event.toolCallId,
      toolName: event.toolName,
      mapping: mapping ?? null,
      allowed,
      reason: !lastPosition
        ? 'closed-not-last'
        : !mapping
          ? 'unknown-tool-call-id'
          : !mapping.verified
            ? `unverified-source:${mapping.sourceReason}`
            : 'exact-tool-call-id-mapped-to-verified-request-nonce',
    });
    if (!allowed) return { block: true, reason: 'Gate0 probe refused unverified tool batch' };
  });

  pi.on('agent_settled', async () => {
    log('agent_settled', { retainedNonceMappingsBeforeClear: toolNonce.size });
    toolNonce.clear();
  });

  pi.registerTool({
    name: 'gate0_nonce_probe',
    label: 'Gate0 Nonce Probe',
    description: 'Gate0-only harmless tool used to prove toolCallId to request-nonce correlation.',
    parameters: Type.Object({ label: Type.String() }),
    async execute(toolCallId, params) {
      const broker = await brokerRequest({ action: 'promote-probe' });
      log('tool_execute', { toolCallId, label: params.label, broker });
      return {
        content: [{ type: 'text', text: `gate0_nonce_probe executed for ${params.label}` }],
        details: { harmless: true },
      };
    },
  });

  pi.registerCommand('gate0-reload', {
    description: 'Trigger a real same-PID Pi extension/runtime reload.',
    handler: async (_args, ctx) => {
      log('reload_command_before');
      await ctx.reload();
      return;
    },
  });
}