import fs from "node:fs/promises";
import crypto from "node:crypto";
import { fileURLToPath } from "node:url";


const EXPECTED_HASH_RE = /expected:\s*([0-9a-f]{64})/i;


type JsonPrimitive = string | number | boolean | null;
type JsonValue = JsonPrimitive | JsonValue[] | { [key: string]: JsonValue };

interface SendResult {
  success?: boolean;
  error?: string;
}

interface LogEntry {
  action?: string;
  [key: string]: unknown;
}

export const SIGNABLE_FIELDS = [
  "agent_id",
  "target_agent_id",
  "timestamp",
  "nonce",
  "input",
  "output",
  "alert_threshold",
] as const;


export function normalizeA2spaPayload(payload: { [key: string]: JsonValue | undefined }): { [key: string]: JsonValue } {
  if (!payload || typeof payload !== "object" || Array.isArray(payload)) {
    throw new TypeError("payload must be an object");
  }

  const normalized: { [key: string]: JsonValue } = {};
  for (const field of SIGNABLE_FIELDS) {
    if (Object.prototype.hasOwnProperty.call(payload, field)) {
      const value = payload[field];
      if (value === undefined) {
        throw new TypeError(`A2SPA signable field cannot be undefined: ${field}`);
      }
      normalized[field] = value;
    }
  }

  if (!Object.prototype.hasOwnProperty.call(normalized, "alert_threshold")) {
    normalized.alert_threshold = 10;
  }

  return normalized;
}


function pythonCompatibleJsonString(value: string): string {
  return JSON.stringify(value).replace(/[\u007f-\uffff]/g, (character) => {
    return "\\u" + character.charCodeAt(0).toString(16).padStart(4, "0");
  });
}


export function canonicalJsonPythonStyle(value: JsonValue | undefined): string {
  if (value === null) return "null";
  if (Array.isArray(value)) {
    return `[${value.map((item) => canonicalJsonPythonStyle(item)).join(", ")}]`;
  }
  if (typeof value === "object") {
    if (!value) return "null";
    const record = value as { [key: string]: JsonValue };
    const keys = Object.keys(record).sort();
    return `{${keys
      .map((key) => `${pythonCompatibleJsonString(key)}: ${canonicalJsonPythonStyle(record[key])}`)
      .join(", ")}}`;
  }
  if (typeof value === "string") {
    return pythonCompatibleJsonString(value);
  }
  if (typeof value === "number") {
    if (!Number.isFinite(value)) {
      throw new TypeError("A2SPA canonicalization only supports finite JSON numbers");
    }
    return String(value);
  }
  if (typeof value === "boolean") {
    return value ? "true" : "false";
  }
  throw new TypeError(`Unsupported payload value type: ${typeof value}`);
}


export function canonicalizeA2spaValue(value: JsonValue | undefined): string {
  return canonicalJsonPythonStyle(value);
}


export function canonicalizeA2spaPayload(payload: { [key: string]: JsonValue | undefined }): string {
  return canonicalJsonPythonStyle(normalizeA2spaPayload(payload));
}


export function prepareA2spaSignedBytes(payload: { [key: string]: JsonValue | undefined }): Buffer {
  return Buffer.from(canonicalizeA2spaPayload(payload), "utf8");
}


export const prepareSignedBytes = prepareA2spaSignedBytes;


export function computeA2spaPayloadHash(payload: { [key: string]: JsonValue | undefined }): string {
  return crypto.createHash("sha256").update(prepareA2spaSignedBytes(payload)).digest("hex");
}


export const computePayloadHash = computeA2spaPayloadHash;


export function finalizeA2spaPayload(payloadBeforeHash: { [key: string]: JsonValue | undefined }): { [key: string]: JsonValue } {
  const finalized: { [key: string]: JsonValue | undefined } = { ...payloadBeforeHash };
  delete finalized.hash;
  delete finalized.signature;
  finalized.hash = computeA2spaPayloadHash(payloadBeforeHash);
  return finalized as { [key: string]: JsonValue };
}


export const finalizePayload = finalizeA2spaPayload;


export function buildA2spaRequestBody(payload: { [key: string]: JsonValue }, signature: string): {
  payload: { [key: string]: JsonValue };
  signature: string;
} {
  return { payload, signature };
}


export const buildRequestBody = buildA2spaRequestBody;


export function extractExpectedHash(errorText: string | null | undefined): string | null {
  if (!errorText) return null;
  const match = String(errorText).match(EXPECTED_HASH_RE);
  return match ? match[1] : null;
}


export function buildDebugSnapshot(
  payloadBeforeHash: { [key: string]: JsonValue | undefined },
  signature: string | null = null,
  finalPayload: { [key: string]: JsonValue } | null = null,
  serverErrorText: string | null = null,
): { [key: string]: unknown } {
  const canonicalString = canonicalizeA2spaPayload(payloadBeforeHash);
  const signedBytes = prepareA2spaSignedBytes(payloadBeforeHash);
  const requestPayload = finalPayload ?? finalizeA2spaPayload(payloadBeforeHash);
  const snapshot: { [key: string]: unknown } = {
    payload_before_hash: payloadBeforeHash,
    canonical_string: canonicalString,
    canonical_bytes_hex: signedBytes.toString("hex"),
    client_computed_hash: computeA2spaPayloadHash(payloadBeforeHash),
    final_request_body: signature ? buildA2spaRequestBody(requestPayload, signature) : { payload: requestPayload },
    server_expected_hash: extractExpectedHash(serverErrorText),
  };
  if (serverErrorText !== null) {
    snapshot.server_error = serverErrorText;
  }
  return snapshot;
}


export function signA2spaPayload(payload: { [key: string]: JsonValue | undefined }, privateKeyPem: string): string {
  const signer = crypto.createSign("RSA-SHA256");
  signer.update(prepareA2spaSignedBytes(payload));
  signer.end();
  return signer.sign(privateKeyPem, "hex");
}


export const signPayload = signA2spaPayload;


export class A2SPAClient {
  private apiBase: string;
  private apiKey: string;
  private agentId: string;
  private privateKeyPath: string;
  private privateKeyPem?: string;

  constructor(options: { apiBase: string; apiKey: string; agentId: string; privateKeyPath: string }) {
    this.apiBase = options.apiBase.replace(/\/$/, "");
    this.apiKey = options.apiKey;
    this.agentId = options.agentId;
    this.privateKeyPath = options.privateKeyPath;
  }

  static fromEnv(): A2SPAClient {
    return new A2SPAClient({
      apiBase: process.env.A2SPA_API_BASE || "https://aimodularity.com/A2SPA",
      apiKey: process.env.A2SPA_API_KEY || "",
      agentId: process.env.A2SPA_AGENT_ID || "",
      privateKeyPath: process.env.A2SPA_PRIVATE_KEY_PATH || "",
    });
  }

  private async loadPrivateKeyPem(): Promise<string> {
    if (!this.privateKeyPem) {
      this.privateKeyPem = await fs.readFile(this.privateKeyPath, "utf8");
    }
    return this.privateKeyPem;
  }

  buildPayloadFields(
    targetAgentId: string,
    input: JsonValue,
    output: JsonValue,
    alertThreshold = 10,
    extra: { [key: string]: JsonValue } = {},
  ): { [key: string]: JsonValue } {
    const payload: { [key: string]: JsonValue } = {
      agent_id: this.agentId,
      target_agent_id: targetAgentId,
      timestamp: new Date().toISOString(),
      nonce: crypto.randomUUID(),
      input,
      output,
      alert_threshold: alertThreshold,
      ...extra,
    };
    return payload;
  }

  buildPayload(
    targetAgentId: string,
    input: JsonValue,
    output: JsonValue,
    alertThreshold = 10,
    extra: { [key: string]: JsonValue } = {},
  ): { [key: string]: JsonValue } {
    return finalizeA2spaPayload(
      this.buildPayloadFields(targetAgentId, input, output, alertThreshold, extra),
    );
  }

  async buildSignedRequest(
    targetAgentId: string,
    input: JsonValue,
    output: JsonValue,
    alertThreshold = 10,
    extra: { [key: string]: JsonValue } = {},
  ): Promise<{
    payloadBeforeHash: { [key: string]: JsonValue };
    payload: { [key: string]: JsonValue };
    signature: string;
    requestBody: { payload: { [key: string]: JsonValue }; signature: string };
  }> {
    const payloadBeforeHash = this.buildPayloadFields(targetAgentId, input, output, alertThreshold, extra);
    const payload = finalizeA2spaPayload(payloadBeforeHash);
    const signature = signA2spaPayload(payloadBeforeHash, await this.loadPrivateKeyPem());
    return {
      payloadBeforeHash,
      payload,
      signature,
      requestBody: buildA2spaRequestBody(payload, signature),
    };
  }

  async postRequestBody(
    requestBody: { payload: { [key: string]: JsonValue }; signature: string },
  ): Promise<Response> {
    return fetch(`${this.apiBase}/api/verify_payload`, {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
        "x-api-key": this.apiKey,
      },
      body: JSON.stringify(requestBody),
    });
  }

  async sendPayload(
    targetAgentId: string,
    input: JsonValue,
    output: JsonValue,
    alertThreshold = 10,
    extra: { [key: string]: JsonValue } = {},
    debug = false,
  ): Promise<SendResult> {
    const signedRequest = await this.buildSignedRequest(targetAgentId, input, output, alertThreshold, extra);
    const response = await this.postRequestBody(signedRequest.requestBody);
    if (!response.ok) {
      const errorText = await response.text();
      if (debug) {
        console.error(
          JSON.stringify(
            buildDebugSnapshot(
              signedRequest.payloadBeforeHash,
              signedRequest.signature,
              signedRequest.payload,
              errorText,
            ),
            null,
            2,
          ),
        );
      }
      throw new Error(errorText);
    }
    return (await response.json()) as SendResult;
  }

  async fetchLogs(agentId = this.agentId): Promise<LogEntry[]> {
    const params = new URLSearchParams({ agent_id: agentId });
    const response = await fetch(`${this.apiBase}/api/logs_for_agent?${params.toString()}`, {
      headers: { "x-api-key": this.apiKey },
    });
    if (!response.ok) {
      throw new Error(await response.text());
    }
    const body = (await response.json()) as { logs?: LogEntry[] };
    return body.logs || [];
  }

  async fetchInbox(agentId = this.agentId, after = "", limit = 50): Promise<LogEntry[]> {
    const params = new URLSearchParams({ agent_id: agentId, limit: String(limit) });
    if (after) params.set("after", after);
    const response = await fetch(`${this.apiBase}/api/inbox_for_agent?${params.toString()}`, {
      headers: { "x-api-key": this.apiKey },
    });
    if (!response.ok) {
      throw new Error(await response.text());
    }
    const body = (await response.json()) as { messages?: LogEntry[] };
    return body.messages || [];
  }

  async runAndSend(
    targetAgentId: string,
    input: JsonValue,
    runner: (input: JsonValue) => Promise<JsonValue> | JsonValue,
    alertThreshold = 10,
  ): Promise<SendResult> {
    const output = await runner(input);
    return this.sendPayload(targetAgentId, input, output, alertThreshold);
  }
}


async function existingAgent(input: JsonValue): Promise<JsonValue> {
  const message = typeof input === "object" && input && !Array.isArray(input) ? (input as { message?: string }).message : undefined;
  return {
    status: "ready",
    summary: `Processed: ${message ?? "no message"}`,
  };
}


async function main(): Promise<void> {
  const client = A2SPAClient.fromEnv();
  const input = { message: "Hello from my existing TypeScript agent" };
  const output = await existingAgent(input);
  const signedRequest = await client.buildSignedRequest(
    process.env.A2SPA_TARGET_AGENT_ID || "",
    input,
    output,
  );
  const firstResponse = await client.postRequestBody(signedRequest.requestBody);
  console.log(`First request: ${firstResponse.status} ${await firstResponse.text()}`);
  const secondResponse = await client.postRequestBody(signedRequest.requestBody);
  console.log(`Second request: ${secondResponse.status} ${await secondResponse.text()}`);
}


if (process.argv[1] && fileURLToPath(import.meta.url) === process.argv[1]) {
  main().catch((error) => {
    console.error(error);
    process.exit(1);
  });
}