import fs from "node:fs/promises";
import crypto from "node:crypto";


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;
}

function canonicalJson(value: JsonValue): string {
  if (value === null) return "null";
  if (Array.isArray(value)) {
    return `[${value.map((item) => canonicalJson(item)).join(", ")}]`;
  }
  if (typeof value === "object") {
    const record = value as { [key: string]: JsonValue };
    const keys = Object.keys(record).sort();
    return `{${keys
      .map((key) => `${JSON.stringify(key)}: ${canonicalJson(record[key])}`)
      .join(", ")}}`;
  }
  return JSON.stringify(value);
}


function stripSecurityFields(payload: { [key: string]: JsonValue }): { [key: string]: JsonValue } {
  const { hash, signature, ...rest } = payload;
  return rest;
}


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;
  }

  private sha256Hex(value: { [key: string]: JsonValue }): string {
    return crypto.createHash("sha256").update(canonicalJson(stripSecurityFields(value))).digest("hex");
  }

  private async sign(payload: { [key: string]: JsonValue }): Promise<string> {
    const signer = crypto.createSign("RSA-SHA256");
    signer.update(canonicalJson(stripSecurityFields(payload)));
    signer.end();
    return signer.sign(await this.loadPrivateKeyPem(), "hex");
  }

  buildPayload(targetAgentId: string, input: JsonValue, output: JsonValue, alertThreshold = 10): { [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,
    };
    payload.hash = this.sha256Hex(payload);
    return payload;
  }

  async sendPayload(targetAgentId: string, input: JsonValue, output: JsonValue, alertThreshold = 10): Promise<SendResult> {
    const payload = this.buildPayload(targetAgentId, input, output, alertThreshold);
    const signature = await this.sign(payload);
    const response = await fetch(`${this.apiBase}/api/verify_payload`, {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
        "x-api-key": this.apiKey,
      },
      body: JSON.stringify({ payload, signature }),
    });
    if (!response.ok) {
      throw new Error(await response.text());
    }
    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 result = await client.runAndSend(
    process.env.A2SPA_TARGET_AGENT_ID || "",
    { message: "Hello from my existing TypeScript agent" },
    existingAgent,
  );
  console.log(result);
}


void main();