/**
 * axon.js — framework-agnostic Axon protocol client
 *
 * Wire message types:
 *   Client → Relay: AUTH=1, SUBSCRIBE=2, UNSUBSCRIBE=3, PUBLISH=4
 *   Relay → Client: CHALLENGE=10, EVENT=11, EOSE=12, OK=13, ERROR=14
 */

import { ed25519, x25519, edwardsToMontgomeryPriv, edwardsToMontgomeryPub } from '@noble/curves/ed25519';
import { sha256 } from '@noble/hashes/sha256';
import { hkdf } from '@noble/hashes/hkdf';
import { chacha20poly1305 } from '@noble/ciphers/chacha';
import { encode, decode } from '@msgpack/msgpack';

// ---------------------------------------------------------------------------
// Event kind constants
// ---------------------------------------------------------------------------

export const KIND = {
  PROFILE:      0,
  MESSAGE:      1000,
  DM:           2000,
  PROGRESS:     3000,
  JOB_REQ:      5000,
  JOB_RESULT:   6000,
  JOB_FEEDBACK: 7000,
};

// ---------------------------------------------------------------------------
// Hex utilities
// ---------------------------------------------------------------------------

export function toHex(bytes) {
  return Array.from(bytes)
    .map(b => b.toString(16).padStart(2, '0'))
    .join('');
}

export function fromHex(hex) {
  if (hex.length % 2 !== 0) throw new Error('axon: odd hex length');
  const out = new Uint8Array(hex.length / 2);
  for (let i = 0; i < out.length; i++) {
    out[i] = parseInt(hex.slice(i * 2, i * 2 + 2), 16);
  }
  return out;
}

// ---------------------------------------------------------------------------
// Key generation
// ---------------------------------------------------------------------------

export function generateSeed() {
  const seed = new Uint8Array(32);
  crypto.getRandomValues(seed);
  return seed;
}

/**
 * keypairFromSeed(seed) → { seed: Uint8Array(32), pubkey: Uint8Array(32) }
 */
export function keypairFromSeed(seed) {
  const pubkey = ed25519.getPublicKey(seed);
  return { seed: seed instanceof Uint8Array ? seed : new Uint8Array(seed), pubkey };
}

// ---------------------------------------------------------------------------
// Canonical encoding (internal helpers)
// ---------------------------------------------------------------------------

/**
 * canonicalTagsBytes encodes tags in the Axon canonical format (big-endian).
 *
 * Tags are sorted by name, then by first value for ties (ascending).
 *
 * Format:
 *   uint16(num_tags)
 *   for each tag:
 *     uint16(len(name)) || utf8(name)
 *     uint16(num_values)
 *     for each value:
 *       uint32(len(value)) || utf8(value)
 */
function canonicalTagsBytes(tags) {
  const enc = new TextEncoder();

  // Sort a copy: by name ascending, then first value ascending for ties.
  const sorted = tags.slice().sort((a, b) => {
    if (a.name !== b.name) return a.name < b.name ? -1 : 1;
    const va = (a.values && a.values.length > 0) ? a.values[0] : '';
    const vb = (b.values && b.values.length > 0) ? b.values[0] : '';
    return va < vb ? -1 : va > vb ? 1 : 0;
  });

  // First pass: compute total byte length.
  let size = 2; // uint16 num_tags
  for (const tag of sorted) {
    const nameBytes = enc.encode(tag.name);
    size += 2 + nameBytes.length; // uint16(len) + name
    size += 2;                    // uint16(num_values)
    for (const v of (tag.values || [])) {
      const vb = enc.encode(v);
      size += 4 + vb.length;     // uint32(len) + value
    }
  }

  const buf = new Uint8Array(size);
  const dv = new DataView(buf.buffer);
  let off = 0;

  dv.setUint16(off, sorted.length, false); off += 2;

  for (const tag of sorted) {
    const nameBytes = enc.encode(tag.name);
    dv.setUint16(off, nameBytes.length, false); off += 2;
    buf.set(nameBytes, off); off += nameBytes.length;

    const values = tag.values || [];
    dv.setUint16(off, values.length, false); off += 2;

    for (const v of values) {
      const vb = enc.encode(v);
      dv.setUint32(off, vb.length, false); off += 4;
      buf.set(vb, off); off += vb.length;
    }
  }

  return buf;
}

/**
 * canonicalPayload builds the deterministic byte blob that is SHA256-hashed
 * to produce the event ID.
 *
 * Layout (big-endian):
 *   [0:2]       uint16 = 32          pubkey length
 *   [2:34]      bytes                pubkey (32 bytes)
 *   [34:42]     uint64               created_at
 *   [42:44]     uint16               kind
 *   [44:48]     uint32               content length
 *   [48:48+n]   bytes                content
 *   [48+n:80+n] bytes                SHA256(canonical_tags_bytes), 32 bytes
 */
function canonicalPayload(pubkey, createdAt, kind, content, tags) {
  const tagsHash = sha256(canonicalTagsBytes(tags));
  const n = content.length;
  // Total: 2 + 32 + 8 + 2 + 4 + n + 32 = 80 + n
  const buf = new Uint8Array(80 + n);
  const dv = new DataView(buf.buffer);

  dv.setUint16(0, 32, false);               // pubkey length
  buf.set(pubkey, 2);                       // pubkey [2:34]
  dv.setBigUint64(34, BigInt(createdAt), false); // created_at [34:42]
  dv.setUint16(42, kind, false);            // kind [42:44]
  dv.setUint32(44, n, false);               // content length [44:48]
  buf.set(content, 48);                     // content [48:48+n]
  buf.set(tagsHash, 48 + n);               // tags hash [48+n:80+n]

  return buf;
}

// ---------------------------------------------------------------------------
// Event signing
// ---------------------------------------------------------------------------

/**
 * signEvent(kp, kind, content, tags=[]) → Event object
 *
 * content may be a string (will be UTF-8 encoded) or Uint8Array.
 * tags is an array of { name: string, values: string[] }.
 *
 * Returns:
 *   { id, pubkey, created_at, kind, content, sig, tags }
 * where binary fields are Uint8Array.
 */
export function signEvent(kp, kind, content, tags = []) {
  const enc = new TextEncoder();
  const contentBytes = typeof content === 'string' ? enc.encode(content) : content;
  const createdAt = Math.floor(Date.now() / 1000);

  const payload = canonicalPayload(kp.pubkey, createdAt, kind, contentBytes, tags);
  const eventID = sha256(payload);
  const sig = ed25519.sign(eventID, kp.seed);

  return {
    id:         eventID,
    pubkey:     kp.pubkey,
    created_at: createdAt,
    kind,
    content:    contentBytes,
    sig,
    tags,
  };
}

// ---------------------------------------------------------------------------
// Challenge signing
// ---------------------------------------------------------------------------

/**
 * signChallenge(kp, nonce, relayURL) → Uint8Array(64)
 *
 * sig = ed25519.sign(SHA256(nonce || utf8(relay_url)), seed)
 */
export function signChallenge(kp, nonce, relayURL) {
  const enc = new TextEncoder();
  const urlBytes = enc.encode(relayURL);
  const msg = new Uint8Array(nonce.length + urlBytes.length);
  msg.set(nonce, 0);
  msg.set(urlBytes, nonce.length);
  const digest = sha256(msg);
  return ed25519.sign(digest, kp.seed);
}

// ---------------------------------------------------------------------------
// Test vector runner
// ---------------------------------------------------------------------------

/**
 * runVectors(vectors) → { pass: bool, results: [{label, pass, got, expected}] }
 *
 * Checks all 6 fields from the vectors JSON:
 *   canonical_tags_hex, canonical_tags_hash, pubkey_hex,
 *   canonical_payload, event_id, sig_hex
 */
export function runVectors(vectors) {
  const results = [];

  function check(label, got, expected) {
    const pass = got === expected;
    results.push({ label, pass, got, expected });
  }

  // Reconstruct inputs from vectors.
  const seed = fromHex(vectors.seed_hex);
  const kp = keypairFromSeed(seed);
  const content = fromHex(vectors.content_hex);
  const tags = vectors.tags_input; // [{ name, values }]

  // 1. canonical_tags_hex
  const tagsBytes = canonicalTagsBytes(tags);
  check('canonical_tags_hex', toHex(tagsBytes), vectors.canonical_tags_hex);

  // 2. canonical_tags_hash
  const tagsHash = sha256(tagsBytes);
  check('canonical_tags_hash', toHex(tagsHash), vectors.canonical_tags_hash);

  // 3. pubkey_hex
  check('pubkey_hex', toHex(kp.pubkey), vectors.pubkey_hex);

  // 4. canonical_payload
  const payload = canonicalPayload(
    kp.pubkey,
    vectors.created_at,
    vectors.kind,
    content,
    tags,
  );
  check('canonical_payload', toHex(payload), vectors.canonical_payload);

  // 5. event_id
  const eventID = sha256(payload);
  check('event_id', toHex(eventID), vectors.event_id);

  // 6. sig_hex
  const sig = ed25519.sign(eventID, seed);
  check('sig_hex', toHex(sig), vectors.sig_hex);

  const pass = results.every(r => r.pass);
  return { pass, results };
}

// ---------------------------------------------------------------------------
// DM encryption (X25519 ECDH + ChaCha20-Poly1305)
// ---------------------------------------------------------------------------

/**
 * dhSharedSecret(senderSeed, recipientEd25519Pub) → Uint8Array(32)
 *
 * Converts both keys to X25519 form, then runs the DH scalar multiplication.
 * Matches Go's DHSharedSecret exactly:
 *   - private: SHA-512(seed)[0:32] clamped per RFC 8032
 *   - public:  Edwards y → Montgomery u via (1+y)/(1-y) mod p
 */
function dhSharedSecret(localSeed, remoteEd25519Pub) {
  const xPriv = edwardsToMontgomeryPriv(localSeed);
  const xPub  = edwardsToMontgomeryPub(remoteEd25519Pub);
  return x25519.getSharedSecret(xPriv, xPub);
}

/**
 * dmKey derives a 32-byte symmetric key from the raw DH shared secret
 * via HKDF-SHA256, matching Go's dmKey().
 */
function dmKey(shared) {
  return hkdf(sha256, shared, undefined, 'axon-dm-v1', 32);
}

/**
 * dmAD builds associated data by concatenating sender + recipient Ed25519
 * public keys, matching Go's dmAD().
 */
function dmAD(senderPub, recipientPub) {
  const ad = new Uint8Array(senderPub.length + recipientPub.length);
  ad.set(senderPub);
  ad.set(recipientPub, senderPub.length);
  return ad;
}

/**
 * encryptDM(senderSeed, recipientPubkey, plaintext) → Uint8Array
 *
 * Returns nonce(12) || ciphertext+tag, matching Go's EncryptDM.
 * plaintext may be string or Uint8Array.
 */
export function encryptDM(senderSeed, recipientPubkey, plaintext) {
  const enc = new TextEncoder();
  const pt = typeof plaintext === 'string' ? enc.encode(plaintext) : plaintext;
  const shared = dhSharedSecret(senderSeed, recipientPubkey);
  const key = dmKey(shared);
  const nonce = crypto.getRandomValues(new Uint8Array(12));
  const senderPub = ed25519.getPublicKey(senderSeed);
  const ad = dmAD(senderPub, recipientPubkey);
  const ct = chacha20poly1305(key, nonce, ad).encrypt(pt);
  const out = new Uint8Array(12 + ct.length);
  out.set(nonce);
  out.set(ct, 12);
  return out;
}

/**
 * decryptDM(recipientSeed, senderPubkey, blob) → Uint8Array
 *
 * blob is nonce(12) || ciphertext+tag. Throws on authentication failure.
 */
export function decryptDM(recipientSeed, senderPubkey, blob) {
  if (blob.length < 12) throw new Error('axon: dm blob too short');
  const shared = dhSharedSecret(recipientSeed, senderPubkey);
  const key = dmKey(shared);
  const nonce = blob.slice(0, 12);
  const ct    = blob.slice(12);
  const recipientPub = ed25519.getPublicKey(recipientSeed);
  const ad = dmAD(senderPubkey, recipientPub);
  return chacha20poly1305(key, nonce, ad).decrypt(ct);
}

// ---------------------------------------------------------------------------
// AxonClient
// ---------------------------------------------------------------------------

const MSG = {
  // Client → Relay
  AUTH:        1,
  SUBSCRIBE:   2,
  UNSUBSCRIBE: 3,
  PUBLISH:     4,
  // Relay → Client
  CHALLENGE:   10,
  EVENT:       11,
  EOSE:        12,
  OK:          13,
  ERROR:       14,
};

export class AxonClient {
  #ws = null;
  #kp = null;
  #relayURL = '';
  #authed = false;
  #listeners = {};

  constructor(relayURL, kp) {
    this.#relayURL = relayURL;
    this.#kp = kp;
  }

  // -------------------------------------------------------------------------
  // Event emitter
  // -------------------------------------------------------------------------

  on(event, fn) {
    if (!this.#listeners[event]) this.#listeners[event] = [];
    this.#listeners[event].push(fn);
    return this;
  }

  #emit(name, ...args) {
    for (const fn of (this.#listeners[name] || [])) {
      try { fn(...args); } catch { /* ignore listener errors */ }
    }
  }

  // -------------------------------------------------------------------------
  // Wire encoding
  // -------------------------------------------------------------------------

  #send(type, payload) {
    if (!this.#ws || this.#ws.readyState !== WebSocket.OPEN) return;
    this.#ws.send(encode([type, payload]));
  }

  // -------------------------------------------------------------------------
  // Connection lifecycle
  // -------------------------------------------------------------------------

  /**
   * connect() → Promise
   *
   * Opens the WebSocket, performs the CHALLENGE → AUTH → OK handshake,
   * then resolves. Rejects on error or if auth fails.
   */
  connect() {
    return new Promise((resolve, reject) => {
      let settled = false;
      const settle = (ok, value) => {
        if (settled) return;
        settled = true;
        if (ok) resolve(value);
        else reject(value);
      };

      const ws = new WebSocket(this.#relayURL);
      ws.binaryType = 'arraybuffer';
      this.#ws = ws;

      ws.onopen = () => {
        // Wait for CHALLENGE from relay; do not resolve yet.
      };

      ws.onclose = (e) => {
        this.#authed = false;
        settle(false, new Error(`axon: ws closed (${e.code})`));
        this.#emit('disconnected', e.code, e.reason);
      };

      ws.onerror = (e) => {
        settle(false, new Error('axon: ws error'));
        this.#emit('error', 0, 'websocket error');
      };

      ws.onmessage = (e) => {
        let msg;
        try {
          msg = decode(new Uint8Array(e.data));
        } catch (err) {
          this.#emit('error', 0, 'decode error: ' + err.message);
          return;
        }

        if (!Array.isArray(msg) || msg.length < 2) {
          this.#emit('error', 0, 'malformed message');
          return;
        }

        const [type, payload] = msg;

        if (!this.#authed) {
          // Pre-auth: only handle CHALLENGE and OK.
          if (type === MSG.CHALLENGE) {
            const nonce = payload.nonce;
            const sig = signChallenge(this.#kp, nonce, this.#relayURL);
            this.#send(MSG.AUTH, { pubkey: this.#kp.pubkey, sig });
          } else if (type === MSG.OK) {
            this.#authed = true;
            settle(true, undefined);
            this.#emit('connected');
          } else if (type === MSG.ERROR) {
            settle(false, new Error(`axon: auth error ${payload.code}: ${payload.message}`));
            this.#emit('error', payload.code, payload.message);
          }
          return;
        }

        // Post-auth: handle all message types.
        switch (type) {
          case MSG.EVENT: {
            const ev = payload.event;
            this.#emit('event', payload.sub_id, ev);
            break;
          }
          case MSG.EOSE:
            this.#emit('eose', payload.sub_id);
            break;
          case MSG.OK:
            this.#emit('ok', payload.message);
            break;
          case MSG.ERROR:
            this.#emit('error', payload.code, payload.message);
            break;
          default:
            // Unknown message type — ignore.
            break;
        }
      };
    });
  }

  disconnect() {
    if (this.#ws) {
      this.#ws.close(1000);
      this.#ws = null;
    }
    this.#authed = false;
  }

  // -------------------------------------------------------------------------
  // Protocol operations
  // -------------------------------------------------------------------------

  /**
   * publish(kind, content, tags=[])
   *
   * Signs and publishes an event. content may be string or Uint8Array.
   */
  publish(kind, content, tags = []) {
    const event = signEvent(this.#kp, kind, content, tags);
    this.#send(MSG.PUBLISH, { event });
  }

  /**
   * subscribe(subID, filter={})
   *
   * Sends a SUBSCRIBE message. All filter fields are sent explicitly.
   */
  subscribe(subID, filter = {}) {
    this.#send(MSG.SUBSCRIBE, {
      sub_id: subID,
      filter: {
        ids:     filter.ids     ?? [],
        authors: filter.authors ?? [],
        kinds:   filter.kinds   ?? [],
        since:   filter.since   ?? 0,
        until:   filter.until   ?? 0,
        limit:   filter.limit   ?? 0,
        tags:    filter.tags    ?? [],
      },
    });
  }

  /**
   * unsubscribe(subID)
   */
  unsubscribe(subID) {
    this.#send(MSG.UNSUBSCRIBE, { sub_id: subID });
  }
}