import { createHash } from "node:crypto";
import {
  ArtifactDescriptor,
  CompleteAttemptRequest,
  EnrollResponse,
  type HeartbeatRequest,
  HeartbeatResponse,
  LeaseRenewResponse,
  LeaseResponse,
  type NodeMetadata,
  type RunEvent,
  StageReport,
  WORKER_PROTOCOL_VERSION,
} from "@autoscholar/schemas";
import type { z } from "zod";

export class ControlPlaneError extends Error {
  constructor(
    message: string,
    readonly status: number,
    readonly code?: string,
  ) {
    super(message);
  }
}

export interface WorkerClientOptions {
  baseUrl: string;
  token?: string;
  fetchImpl?: typeof fetch;
  timeoutMs?: number;
  userAgent?: string;
  retryAttempts?: number;
}

/** Worker-side HTTP client for the versioned MacLustr worker protocol (outbound-only). */
export class WorkerClient {
  private token: string | undefined;
  private readonly base: string;
  private readonly fetchImpl: typeof fetch;
  private readonly timeoutMs: number;
  private readonly userAgent: string;
  private readonly retryAttempts: number;

  constructor(opts: WorkerClientOptions) {
    this.base = opts.baseUrl.replace(/\/+$/, "");
    this.token = opts.token;
    this.fetchImpl = opts.fetchImpl ?? fetch;
    this.timeoutMs = opts.timeoutMs ?? 30_000;
    this.userAgent = opts.userAgent ?? `autoscholar-worker/protocol-${WORKER_PROTOCOL_VERSION}`;
    this.retryAttempts = opts.retryAttempts ?? 8;
  }

  setToken(token: string): void {
    this.token = token;
  }

  private async call<T>(
    method: string,
    path: string,
    body?: unknown,
    schema?: z.ZodType<T>,
    raw?: { buffer: Buffer; contentType: string },
  ): Promise<T> {
    // transient failures (network, 5xx) are retried with backoff so a short control-plane restart never kills a run
    const maxAttempts = this.retryAttempts;
    let lastErr: unknown;
    for (let attempt = 1; attempt <= maxAttempts; attempt++) {
      try {
        return await this.callOnce(method, path, body, schema, raw);
      } catch (e) {
        lastErr = e;
        const transient = !(e instanceof ControlPlaneError) || e.status >= 500;
        if (!transient || attempt === maxAttempts) throw e;
        await new Promise((r) => setTimeout(r, Math.min(15_000, 1000 * 2 ** (attempt - 1))));
      }
    }
    throw lastErr;
  }

  private async callOnce<T>(
    method: string,
    path: string,
    body?: unknown,
    schema?: z.ZodType<T>,
    raw?: { buffer: Buffer; contentType: string },
  ): Promise<T> {
    const ctrl = new AbortController();
    const t = setTimeout(() => ctrl.abort(), this.timeoutMs);
    try {
      const headers: Record<string, string> = {
        "user-agent": this.userAgent,
        "x-protocol-version": String(WORKER_PROTOCOL_VERSION),
      };
      if (this.token) headers.authorization = `Bearer ${this.token}`;
      let payload: BodyInit | undefined;
      if (raw) {
        headers["content-type"] = raw.contentType;
        payload = new Blob([new Uint8Array(raw.buffer)]);
      } else if (body !== undefined) {
        headers["content-type"] = "application/json";
        payload = JSON.stringify(body);
      }
      const res = await this.fetchImpl(`${this.base}${path}`, {
        method,
        headers,
        body: payload,
        signal: ctrl.signal,
      });
      const text = await res.text();
      let json: unknown = null;
      try {
        json = text ? JSON.parse(text) : null;
      } catch {
        json = null;
      }
      if (!res.ok) {
        const err = (json as { error?: { message?: string; code?: string } } | null)?.error;
        throw new ControlPlaneError(err?.message ?? `HTTP ${res.status} on ${path}`, res.status, err?.code);
      }
      return schema ? schema.parse(json) : (json as T);
    } finally {
      clearTimeout(t);
    }
  }

  enroll(enrollmentSecret: string, name: string, node: NodeMetadata, lanIp?: string) {
    return this.call(
      "POST",
      "/api/internal/workers/enroll",
      { enrollmentSecret, name, node, lanIp },
      EnrollResponse,
    );
  }
  heartbeat(body: z.input<typeof HeartbeatRequest>) {
    return this.call("POST", "/api/internal/workers/heartbeat", body, HeartbeatResponse);
  }
  lease(profiles: string[]) {
    return this.call("POST", "/api/internal/workers/lease", { profiles }, LeaseResponse);
  }
  renew(runId: string, fencingToken: number) {
    return this.call(
      "POST",
      `/api/internal/workers/runs/${runId}/renew`,
      { fencingToken },
      LeaseRenewResponse,
    );
  }
  events(runId: string, fencingToken: number, events: RunEvent[]) {
    return this.call<{ ok: true; lastSeq: number }>("POST", `/api/internal/workers/runs/${runId}/events`, {
      fencingToken,
      events,
    });
  }
  stage(runId: string, report: z.input<typeof StageReport>) {
    return this.call<{ ok: true }>(
      "POST",
      `/api/internal/workers/runs/${runId}/stage`,
      StageReport.parse(report),
    );
  }
  async uploadArtifact(
    runId: string,
    fencingToken: number,
    descriptor: Omit<z.input<typeof ArtifactDescriptor>, "sha256" | "bytes">,
    buffer: Buffer,
  ) {
    const sha256 = createHash("sha256").update(buffer).digest("hex");
    const artifact = ArtifactDescriptor.parse({ ...descriptor, sha256, bytes: buffer.length });
    const q = new URLSearchParams({
      fencingToken: String(fencingToken),
      path: artifact.path,
      kind: artifact.kind,
      sha256,
      stage: artifact.stage,
      contentType: artifact.contentType,
      synthetic: artifact.synthetic ? "1" : "0",
    });
    return this.call<{ ok: true; artifactId: string; sha256: string; bytes: number }>(
      "PUT",
      `/api/internal/workers/runs/${runId}/artifacts?${q.toString()}`,
      undefined,
      undefined,
      { buffer, contentType: artifact.contentType },
    );
  }
  listArtifacts(runId: string, fencingToken: number) {
    return this.call<{
      artifacts: Array<{
        path: string;
        version: number;
        sha256: string;
        bytes: number;
        contentType: string;
        kind: string;
      }>;
    }>("GET", `/api/internal/workers/runs/${runId}/artifacts?fencingToken=${fencingToken}`);
  }
  async downloadArtifact(runId: string, fencingToken: number, path: string): Promise<Buffer> {
    const ctrl = new AbortController();
    const t = setTimeout(() => ctrl.abort(), 120_000);
    try {
      const headers: Record<string, string> = {
        "user-agent": this.userAgent,
        "x-protocol-version": String(WORKER_PROTOCOL_VERSION),
      };
      if (this.token) headers.authorization = `Bearer ${this.token}`;
      const res = await this.fetchImpl(
        `${this.base}/api/internal/workers/runs/${runId}/artifacts/${path.split("/").map(encodeURIComponent).join("/")}?fencingToken=${fencingToken}`,
        { headers, signal: ctrl.signal },
      );
      if (!res.ok) throw new ControlPlaneError(`HTTP ${res.status} downloading ${path}`, res.status);
      const buf = Buffer.from(await res.arrayBuffer());
      const sha = res.headers.get("x-artifact-sha256");
      if (sha && createHash("sha256").update(buf).digest("hex") !== sha)
        throw new ControlPlaneError(`digest mismatch downloading ${path}`, 500);
      return buf;
    } finally {
      clearTimeout(t);
    }
  }
  complete(runId: string, body: z.input<typeof CompleteAttemptRequest>) {
    return this.call<{ ok: true }>(
      "POST",
      `/api/internal/workers/runs/${runId}/complete`,
      CompleteAttemptRequest.parse(body),
    );
  }
}

export { WORKER_PROTOCOL_VERSION };