import { randomBytes, randomUUID } from 'node:crypto';
import { createServer, type IncomingMessage, type Server, type ServerResponse } from 'node:http';
import { mkdir, rm } from 'node:fs/promises';
import path from 'node:path';
import { atomicWriteJson } from '../../coding-projects/atomic-json';
import { materializeMakelorePiExtension } from './extensions/makelore-runtime';
import {
  PiProjectWriteLeaseCoordinator,
  type PiProjectWriteLease,
} from './write-lease';
import type { PiGenerationResourceInput } from './worker-pool';
import {
  parsePiSubagentDispatchRequest,
  type PiSubagentScheduler,
} from './subagent';
import type { PiProductTools } from './product-tools';
import type { CodingPluginToolDefinition } from '../../../shared/coding-plugins';
import type { PiSkillEntry } from './resource-loader';
import type { EffectivePluginSnapshot } from '../../coding-plugins/effective-resolver';
import { DevicePackageError } from '../../coding-packages/device-package-manager';
import { AgentBrowserFault } from '../../agent-browser/fault';

const MAX_REQUEST_BYTES = 64 * 1024;
const PRODUCT_TOOL_NAME_PATTERN = /^[A-Za-z][A-Za-z0-9._:-]{0,63}$/u;
const CORE_PRODUCT_TOOL_NAMES = new Set([
  'agent_browser',
  'task_state',
  'changed_file',
  'runtime_context',
]);

interface WorkerRegistrationRecord {
  managedModelRequest?: import('../../../shared/managed-model-capabilities').ManagedModelRequest;
  token: string;
  conversationId: string;
  generation: number;
  projectId: string;
  projectPath: string | null;
  skillIds: string[];
  catalogRevision?: number;
  allowedToolNames: string[];
  tools: CodingPluginToolDefinition[];
  projectWriteLeaseToolNames: string[];
  role: 'parent' | 'child';
  effectiveSnapshot?: EffectivePluginSnapshot;
  devicePackageGeneration?: number;
  devicePackageIds?: string[];
  contextFile: string;
  runId: string | null;
  leases: Map<string, PiProjectWriteLease>;
  waiters: Map<string, AbortController>;
}

export interface PiExtensionWorkerRegistration {
  extensionPath: string;
  env: NodeJS.ProcessEnv;
  sensitiveValues: string[];
  allowedToolNames: readonly string[];
  dispose(): Promise<void>;
}

export interface RegisterPiExtensionWorkerInput {
  managedModelRequest?: import('../../../shared/managed-model-capabilities').ManagedModelRequest;
  conversationId: string;
  generation: number;
  projectId: string;
  projectPath?: string;
  skillEntries?: readonly PiSkillEntry[];
  catalogRevision?: number;
  tools?: readonly CodingPluginToolDefinition[];
  /** Exact Main-owned resolver output used for this worker generation. */
  effectiveSnapshot?: EffectivePluginSnapshot;
  devicePackageGeneration?: number;
  devicePackageIds?: readonly string[];
  extensionsDir: string;
  role?: 'parent' | 'child';
  runId?: string;
}

interface LeaseBridgeRequest {
  action: 'lease.acquire' | 'lease.release';
  conversationId: string;
  workerGeneration: number;
  runId: string;
  resourceId: string;
  leaseId?: string;
}

interface SubagentBridgeRequest {
  action: 'subagent.dispatch';
  conversationId: string;
  workerGeneration: number;
  runId: string;
  resourceId: string;
  request: unknown;
}

interface ProductToolBridgeRequest {
  action: 'product.invoke' | 'product.invoke.stream';
  conversationId: string;
  workerGeneration: number;
  runId: string;
  resourceId: string;
  toolName: string;
  input: unknown;
}

interface ChangeRefreshBridgeRequest {
  action: 'changes.bash';
  conversationId: string;
  workerGeneration: number;
  runId: string;
  resourceId: string;
}

interface ChangeTouchedBridgeRequest {
  action: 'changes.touched';
  conversationId: string;
  workerGeneration: number;
  runId: string;
  resourceId: string;
  paths: string[];
}

type BridgeRequest = LeaseBridgeRequest
  | SubagentBridgeRequest
  | ProductToolBridgeRequest
  | ChangeRefreshBridgeRequest
  | ChangeTouchedBridgeRequest;

export interface PiExtensionSubagentBridge {
  scheduler: PiSubagentScheduler;
  trackGenerationResource?(input: PiGenerationResourceInput): () => void;
}

function recordValue(value: unknown): value is Record<string, unknown> {
  return value !== null && typeof value === 'object' && !Array.isArray(value);
}

function bridgeRequest(value: unknown): value is BridgeRequest {
  if (!recordValue(value)) return false;
  const common = typeof value.conversationId === 'string'
    && Number.isSafeInteger(value.workerGeneration)
    && typeof value.runId === 'string'
    && typeof value.resourceId === 'string';
  if (!common) return false;
  if (value.action === 'subagent.dispatch') return 'request' in value;
  if (value.action === 'product.invoke' || value.action === 'product.invoke.stream') {
    return typeof value.toolName === 'string'
      && PRODUCT_TOOL_NAME_PATTERN.test(value.toolName)
      && 'input' in value;
  }
  if (value.action === 'changes.bash') return true;
  if (value.action === 'changes.touched') {
    return Array.isArray(value.paths)
      && value.paths.length > 0
      && value.paths.length <= 200
      && value.paths.every((filePath) => typeof filePath === 'string');
  }
  return (value.action === 'lease.acquire' || value.action === 'lease.release')
    && (value.leaseId === undefined || typeof value.leaseId === 'string');
}

export class PiManagedExtensionHost {
  private readonly leases: PiProjectWriteLeaseCoordinator;
  private readonly registrations = new Map<string, WorkerRegistrationRecord>();
  private readonly runBindings = new Map<string, string>();
  private readonly requestFlights = new Set<Promise<void>>();
  private readonly extensionMaterializations = new Map<string, Promise<string>>();
  private subagentBridge: PiExtensionSubagentBridge | undefined;
  private productTools: PiProductTools | undefined;
  private server: Server | null = null;
  private bridgeUrl: string | null = null;
  private startFlight: Promise<string> | null = null;
  private closing = false;

  constructor(leases = new PiProjectWriteLeaseCoordinator()) {
    this.leases = leases;
  }

  configureSubagents(bridge: PiExtensionSubagentBridge): void {
    this.subagentBridge = bridge;
  }

  configureProductTools(productTools: PiProductTools): void {
    this.productTools = productTools;
  }

  getDiagnostics(): {
    registrations: { parent: number; child: number };
    writeLeases: { active: number; waiting: number };
    bridgeRequests: number;
  } {
    const records = [...this.registrations.values()];
    return {
      registrations: {
        parent: records.filter(({ role }) => role === 'parent').length,
        child: records.filter(({ role }) => role === 'child').length,
      },
      writeLeases: {
        active: records.reduce((total, { leases }) => total + leases.size, 0),
        waiting: records.reduce((total, { waiters }) => total + waiters.size, 0),
      },
      bridgeRequests: this.requestFlights.size,
    };
  }

  async registerWorker(input: RegisterPiExtensionWorkerInput): Promise<PiExtensionWorkerRegistration> {
    if (!Number.isSafeInteger(input.generation) || input.generation <= 0) {
      throw new Error('Worker generation must be a positive safe integer');
    }
    const bridgeUrl = await this.start();
    await mkdir(input.extensionsDir, { recursive: true });
    let materialization = this.extensionMaterializations.get(input.extensionsDir);
    if (!materialization) {
      materialization = materializeMakelorePiExtension(input.extensionsDir).catch((error) => {
        this.extensionMaterializations.delete(input.extensionsDir);
        throw error;
      });
      this.extensionMaterializations.set(input.extensionsDir, materialization);
    }
    const extensionPath = await materialization;
    const role = input.role ?? 'parent';
    if (role === 'child' && !input.runId?.trim()) {
      throw new Error('Child extension registration requires a parent run id');
    }
    if (this.productTools && !input.projectPath?.trim()) {
      throw new Error('Product tools require a worker project path');
    }
    const skillEntries = [...new Map(
      (input.skillEntries ?? [])
        .filter((entry) => entry.id.trim() && entry.entryPath.trim())
        .map((entry) => [entry.id.trim(), {
          id: entry.id.trim(),
          entryPath: entry.entryPath.trim().replaceAll('\\', '/'),
          ...(entry.packageRoot?.trim() ? { packageRoot: path.resolve(entry.packageRoot) } : {}),
        }]),
    ).values()];
    const tools = role === 'child'
      ? []
      : [...new Map(
        (input.tools ?? []).map((tool) => [tool.name, structuredClone(tool)]),
      ).values()];
    const allowedToolNames = tools.map(({ name }) => name);
    const projectWriteLeaseToolNames = tools
      .filter(({ projectWriteLease }) => projectWriteLease)
      .map(({ name }) => name);
    const token = randomBytes(32).toString('base64url');
    const contextFile = path.join(input.extensionsDir, `worker-${randomUUID()}.json`);
    const record: WorkerRegistrationRecord = {
      managedModelRequest: input.managedModelRequest,
      token,
      conversationId: input.conversationId,
      generation: input.generation,
      projectId: input.projectId,
      projectPath: input.projectPath?.trim() || null,
      skillIds: skillEntries.map(({ id }) => id),
      ...(input.catalogRevision === undefined ? {} : { catalogRevision: input.catalogRevision }),
      allowedToolNames,
      tools,
      projectWriteLeaseToolNames,
      role,
      ...(input.effectiveSnapshot ? { effectiveSnapshot: input.effectiveSnapshot } : {}),
      ...(input.devicePackageGeneration === undefined
        ? {}
        : { devicePackageGeneration: input.devicePackageGeneration }),
      ...(input.devicePackageIds ? { devicePackageIds: [...input.devicePackageIds] } : {}),
      contextFile,
      runId: role === 'child'
        ? input.runId as string
        : this.runBindings.get(input.conversationId) ?? null,
      leases: new Map(),
      waiters: new Map(),
    };
    this.registrations.set(token, record);
    await this.writeContext(record);
    let disposed = false;
    return {
      extensionPath,
      env: {
        MAKELORE_PI_BRIDGE_URL: bridgeUrl,
        MAKELORE_PI_WORKER_TOKEN: token,
        MAKELORE_PI_CONTEXT_FILE: contextFile,
        MAKELORE_PI_WORKER_ROLE: role,
        ...(record.projectPath ? { MAKELORE_PI_PROJECT_PATH: record.projectPath } : {}),
      },
      sensitiveValues: [token],
      allowedToolNames: [...allowedToolNames],
      dispose: async () => {
        if (disposed) return;
        disposed = true;
        this.disposeRecord(record);
        await rm(contextFile, { force: true });
      },
    };
  }

  async bindRun(conversationId: string, generation: number, runId: string, managedModelRequest?: import('../../../shared/managed-model-capabilities').ManagedModelRequest): Promise<void> {
    const record = this.findWorker(conversationId, generation);
    if (!record) throw new Error('Pi extension worker registration is unavailable');
    if (this.productTools && record.projectPath) {
      await this.productTools.beginRun({ conversationId, runId, projectPath: record.projectPath });
    }
    this.releaseWorkerResources(record);
    this.runBindings.set(conversationId, runId);
    record.runId = runId;
    if (managedModelRequest) record.managedModelRequest = structuredClone(managedModelRequest);
    await this.writeContext(record);
  }

  async clearRun(conversationId: string, generation: number, runId?: string): Promise<void> {
    const record = this.findWorker(conversationId, generation);
    if (!runId || this.runBindings.get(conversationId) === runId) {
      this.runBindings.delete(conversationId);
    }
    if (!record || (runId && record.runId !== runId)) return;
    if (this.productTools && record.runId) {
      await this.productTools.settleRun(conversationId, record.runId).catch(() => null);
    }
    this.releaseWorkerResources(record);
    record.runId = null;
    await this.writeContext(record);
  }

  async close(): Promise<void> {
    const server = this.server;
    this.server = null;
    this.bridgeUrl = null;
    this.startFlight = null;
    this.closing = true;
    const closed = server ? new Promise<void>((resolve, reject) => {
      server.close((error) => error ? reject(error) : resolve());
    }) : Promise.resolve();
    for (const record of [...this.registrations.values()]) this.disposeRecord(record);
    this.runBindings.clear();
    await Promise.allSettled([...this.requestFlights]);
    server?.closeIdleConnections();
    await closed;
    this.closing = false;
  }

  private start(): Promise<string> {
    if (this.bridgeUrl) return Promise.resolve(this.bridgeUrl);
    if (this.startFlight) return this.startFlight;
    this.startFlight = new Promise<string>((resolve, reject) => {
      const server = createServer((request, response) => {
        const flight = this.handle(request, response).finally(() => {
          this.requestFlights.delete(flight);
        });
        this.requestFlights.add(flight);
      });
      server.once('error', reject);
      server.listen(0, '127.0.0.1', () => {
        server.removeListener('error', reject);
        const address = server.address();
        if (!address || typeof address === 'string') {
          reject(new Error('Pi extension bridge did not bind a loopback port'));
          return;
        }
        this.server = server;
        this.closing = false;
        this.bridgeUrl = `http://127.0.0.1:${address.port}/v1/worker`;
        resolve(this.bridgeUrl);
      });
    }).finally(() => {
      if (!this.bridgeUrl) this.startFlight = null;
    });
    return this.startFlight;
  }

  private async handle(request: IncomingMessage, response: ServerResponse): Promise<void> {
    if (request.method !== 'POST' || request.url !== '/v1/worker') {
      this.respond(response, 404, { error: 'Not found' });
      return;
    }
    const authorization = request.headers.authorization;
    const token = authorization?.startsWith('Bearer ') ? authorization.slice(7) : '';
    const record = this.registrations.get(token);
    if (!record) {
      this.respond(response, 401, { error: 'Worker token is invalid' });
      return;
    }
    try {
      const value = await this.readBody(request);
      if (!bridgeRequest(value)) {
        this.respond(response, 400, { error: 'Bridge request is invalid' });
        return;
      }
      if (value.conversationId !== record.conversationId
        || value.workerGeneration !== record.generation
        || value.runId !== record.runId) {
        this.respond(response, 409, { error: 'Worker run identity is stale' });
        return;
      }
      if (!value.resourceId.trim()) {
        this.respond(response, 400, { error: 'Bridge resource id is required' });
        return;
      }
      if (value.action === 'subagent.dispatch') {
        await this.dispatchSubagents(request, response, record, value);
        return;
      }
      if (value.action === 'changes.bash') {
        if (!this.productTools) {
          this.respond(response, 503, { error: 'Conversation change tracker is unavailable' });
          return;
        }
        await this.productTools.markBash(record.conversationId, value.runId);
        this.respond(response, 200, { marked: true });
        return;
      }
      if (value.action === 'changes.touched') {
        if (!this.productTools) {
          this.respond(response, 503, { error: 'Conversation change tracker is unavailable' });
          return;
        }
        await this.productTools.recordTouchedPaths(
          record.conversationId,
          value.runId,
          value.paths,
        );
        this.respond(response, 200, { recorded: true });
        return;
      }
      if (value.action === 'product.invoke' || value.action === 'product.invoke.stream') {
        if (record.role !== 'parent') {
          this.respond(response, 403, { error: 'Child workers cannot invoke parent product tools' });
          return;
        }
        if (!CORE_PRODUCT_TOOL_NAMES.has(value.toolName)
          && !record.allowedToolNames.includes(value.toolName)) {
          this.respond(response, 403, { error: 'Product tool is not enabled for this worker' });
          return;
        }
        if (!this.productTools || !record.projectPath) {
          this.respond(response, 503, { error: 'Product tools are unavailable' });
          return;
        }
        const context = {
          conversationId: record.conversationId,
          workerGeneration: record.generation,
          runId: value.runId,
          resourceId: value.resourceId,
          projectId: record.projectId,
          projectPath: record.projectPath,
          skillIds: record.skillIds,
          ...(record.effectiveSnapshot ? { effectiveSnapshot: record.effectiveSnapshot } : {}),
        };
        if (value.action === 'product.invoke.stream') {
          response.writeHead(200, {
            'content-type': 'application/x-ndjson; charset=utf-8',
            ...(this.closing ? { connection: 'close' } : {}),
          });
          try {
            const productResult = await this.productTools.execute(
              value.toolName,
              context,
              value.input,
              (result) => {
                if (!response.writableEnded && !response.destroyed) {
                  response.write(`${JSON.stringify({ result })}\n`);
                }
              },
            );
            if (!response.writableEnded && !response.destroyed) {
              response.end(`${JSON.stringify({ result: productResult, done: true })}\n`);
            }
          } catch {
            if (!response.writableEnded && !response.destroyed) {
              response.end(`${JSON.stringify({ error: 'Product tool invocation failed' })}\n`);
            }
          }
          return;
        }
        const productResult = await this.productTools.execute(
          value.toolName,
          context,
          value.input,
        );
        this.respond(response, 200, { result: productResult });
        return;
      }
      if (value.action === 'lease.release') {
        const lease = record.leases.get(value.resourceId);
        if (!lease || !value.leaseId || lease.id !== value.leaseId) {
          this.respond(response, 409, { error: 'Project write lease is stale' });
          return;
        }
        record.leases.delete(value.resourceId);
        lease.release();
        this.respond(response, 200, { released: true });
        return;
      }
      if (record.leases.has(value.resourceId) || record.waiters.has(value.resourceId)) {
        this.respond(response, 409, { error: 'Project write lease resource already exists' });
        return;
      }
      const controller = new AbortController();
      record.waiters.set(value.resourceId, controller);
      const cancel = () => {
        if (!response.writableEnded) controller.abort();
      };
      request.once('aborted', cancel);
      response.once('close', cancel);
      try {
        const lease = await this.leases.acquire(
          record.projectId,
          `${record.conversationId}:${record.generation}:${value.resourceId}`,
          controller.signal,
        );
        if (record.runId !== value.runId || this.registrations.get(token) !== record) {
          lease.release();
          this.respond(response, 409, { error: 'Worker run identity is stale' });
          return;
        }
        record.leases.set(value.resourceId, lease);
        this.respond(response, 200, { leaseId: lease.id });
      } finally {
        request.removeListener('aborted', cancel);
        response.removeListener('close', cancel);
        record.waiters.delete(value.resourceId);
      }
    } catch (error) {
      if (!response.writableEnded && error instanceof AgentBrowserFault) {
        const message = error.message.slice(0, 512);
        this.respond(response, 400, {
          code: error.code,
          error: `${error.code}: ${message}`,
          retryable: error.retryable,
          ...(error.generation === undefined ? {} : { generation: error.generation }),
          ...(error.outcome === undefined ? {} : { outcome: error.outcome }),
        });
        return;
      }
      if (!response.writableEnded && error instanceof DevicePackageError) {
        this.respond(response, 400, {
          code: error.code,
          error: error.message.slice(0, 512),
        });
        return;
      }
      if (!response.writableEnded) this.respond(response, 400, { error: 'Bridge request failed' });
    }
  }

  private async dispatchSubagents(
    request: IncomingMessage,
    response: ServerResponse,
    record: WorkerRegistrationRecord,
    value: SubagentBridgeRequest,
  ): Promise<void> {
    if (record.role !== 'parent') {
      this.respond(response, 403, { error: 'Child workers cannot dispatch subagents' });
      return;
    }
    if (!this.subagentBridge) {
      this.respond(response, 503, { error: 'Subagent scheduler is unavailable' });
      return;
    }
    try {
      parsePiSubagentDispatchRequest(value.request);
    } catch {
      this.respond(response, 400, { error: 'Subagent dispatch request is invalid' });
      return;
    }
    if (record.waiters.has(value.resourceId) || record.leases.has(value.resourceId)) {
      this.respond(response, 409, { error: 'Bridge resource already exists' });
      return;
    }
    const controller = new AbortController();
    record.waiters.set(value.resourceId, controller);
    const cancel = () => controller.abort();
    request.once('aborted', cancel);
    response.once('close', cancel);
    let untrack = () => undefined;
    try {
      untrack = this.subagentBridge.trackGenerationResource?.({
        conversationId: record.conversationId,
        kind: 'child',
        id: value.resourceId,
        cancel,
      }) ?? untrack;
      response.writeHead(200, {
        'content-type': 'application/x-ndjson; charset=utf-8',
        ...(this.closing ? { connection: 'close' } : {}),
      });
      const result = await this.subagentBridge.scheduler.dispatch({
        conversationId: record.conversationId,
        workerGeneration: record.generation,
        runId: value.runId,
        projectId: record.projectId,
        request: value.request,
      }, {
        signal: controller.signal,
        onUpdate: (details) => {
          if (!response.writableEnded && !response.destroyed) {
            response.write(`${JSON.stringify({ details })}\n`);
          }
        },
      });
      if (!response.writableEnded && !response.destroyed) {
        response.end(`${JSON.stringify({ details: result.details, done: true })}\n`);
      }
    } catch {
      if (!response.headersSent) this.respond(response, 400, { error: 'Subagent dispatch failed' });
      else if (!response.writableEnded && !response.destroyed) response.end();
    } finally {
      request.removeListener('aborted', cancel);
      response.removeListener('close', cancel);
      record.waiters.delete(value.resourceId);
      untrack();
    }
  }

  private readBody(request: IncomingMessage): Promise<unknown> {
    return new Promise((resolve, reject) => {
      const chunks: Buffer[] = [];
      let bytes = 0;
      request.on('data', (chunk: Buffer) => {
        bytes += chunk.length;
        if (bytes > MAX_REQUEST_BYTES) {
          reject(new Error('Bridge request is too large'));
          request.destroy();
          return;
        }
        chunks.push(chunk);
      });
      request.once('end', () => {
        try {
          resolve(JSON.parse(Buffer.concat(chunks).toString('utf8')));
        } catch (error) {
          reject(error);
        }
      });
      request.once('error', reject);
    });
  }

  private respond(response: ServerResponse, status: number, body: Record<string, unknown>): void {
    if (response.writableEnded) return;
    response.writeHead(status, {
      'content-type': 'application/json; charset=utf-8',
      ...(this.closing ? { connection: 'close' } : {}),
    });
    response.end(JSON.stringify(body));
  }

  private findWorker(conversationId: string, generation: number): WorkerRegistrationRecord | undefined {
    return [...this.registrations.values()].find((record) => (
      record.role === 'parent'
      && record.conversationId === conversationId && record.generation === generation
    ));
  }

  private async writeContext(record: WorkerRegistrationRecord): Promise<void> {
    await atomicWriteJson(record.contextFile, {
      conversationId: record.conversationId,
      workerGeneration: record.generation,
      role: record.role,
      skillIds: record.skillIds,
      ...(record.catalogRevision === undefined ? {} : { catalogRevision: record.catalogRevision }),
      ...(record.effectiveSnapshot ? { effectivePluginSnapshot: record.effectiveSnapshot } : {}),
      ...(record.devicePackageGeneration === undefined
        ? {}
        : { devicePackageGeneration: record.devicePackageGeneration }),
      ...(record.devicePackageIds ? { devicePackageIds: record.devicePackageIds } : {}),
      allowedToolNames: record.allowedToolNames,
      tools: record.tools,
      projectWriteLeaseToolNames: record.projectWriteLeaseToolNames,
      ...(record.runId ? { runId: record.runId } : {}),
      ...(record.managedModelRequest ? { managedModelRequest: record.managedModelRequest } : {}),
    });
  }

  private disposeRecord(record: WorkerRegistrationRecord): void {
    if (this.registrations.get(record.token) !== record) return;
    this.registrations.delete(record.token);
    this.releaseWorkerResources(record);
  }

  private releaseWorkerResources(record: WorkerRegistrationRecord): void {
    for (const controller of record.waiters.values()) controller.abort();
    record.waiters.clear();
    for (const lease of record.leases.values()) lease.release();
    record.leases.clear();
  }
}