import type { PreKeyBundle, OneTimePreKey, CryptoProvider } from '@shade/core';
import { NetworkError } from '@shade/core';
import {
  type KTProof,
  type KTProofWire,
  type SignedTreeHead,
  ktProofFromWire,
  verifyBundleAbsence,
  verifyBundleInclusion,
  verifyBundleTombstone,
} from '@shade/key-transparency';
import {
  KTSplitViewError,
  KTVerificationError,
} from '@shade/key-transparency';

/**
 * HTTP transport client for the Shade Prekey Server.
 *
 * Signing: write operations (register, replenish, delete) are signed with
 * the caller's Ed25519 identity signing key. The caller must provide a
 * signing function via the constructor.
 *
 * Usage:
 * ```ts
 * const transport = new ShadeFetchTransport({
 *   baseUrl: 'https://shade.example.com',
 *   crypto,
 *   signingKey: identity.signingPrivateKey,
 * });
 * await transport.register('alice', identity, signedPreKey, oneTimePreKeys);
 * const bundle = await transport.fetchBundle('bob'); // anonymous
 * ```
 */
/** Result of a KT-aware bundle fetch. */
export interface FetchBundleResult {
  bundle: PreKeyBundle;
  /**
   * Verified Signed Tree Head when KT was active and the proof verified.
   * Callers should feed this into a `LightWitness` for split-view tracking.
   */
  ktSth?: SignedTreeHead;
}

/**
 * Optional KT verifier callback. When provided to `ShadeFetchTransport`,
 * the transport verifies every bundle proof before handing the bundle to
 * the SDK. Missing-or-invalid proofs throw `KTVerificationError`.
 *
 * `mode`:
 *   - `'observe'`     verify proof when present, do not fail when missing.
 *   - `'observe-strict'` require a proof; throw if absent.
 */
export interface KTVerifierOptions {
  mode: 'observe' | 'observe-strict';
  logPublicKey: Uint8Array;
  /** Inject `now` for tests; defaults to `Date.now()`. */
  now?: () => number;
  /** Override default 24h freshness. */
  maxStaleMs?: number;
  /** Optional hook to track every observed STH (e.g. into a LightWitness). */
  onObserveSth?: (sth: SignedTreeHead) => void | Promise<void>;
}

export class ShadeFetchTransport {
  private readonly baseUrl: string;
  private readonly crypto: CryptoProvider;
  private readonly signingPrivateKey?: Uint8Array;
  private readonly kt?: KTVerifierOptions;

  constructor(options: {
    baseUrl: string;
    crypto: CryptoProvider;
    signingPrivateKey?: Uint8Array;
    keyTransparency?: KTVerifierOptions;
  }) {
    this.baseUrl = options.baseUrl;
    this.crypto = options.crypto;
    if (options.signingPrivateKey !== undefined) {
      this.signingPrivateKey = options.signingPrivateKey;
    }
    if (options.keyTransparency !== undefined) {
      this.kt = options.keyTransparency;
    }
  }

  private headers(): Record<string, string> {
    return { 'Content-Type': 'application/json' };
  }

  /** Sign a payload using the configured signing key */
  private async sign<T extends Record<string, unknown>>(payload: T): Promise<T & { signedAt: number; signature: string }> {
    if (!this.signingPrivateKey) {
      throw new Error('ShadeFetchTransport: signingPrivateKey is required for write operations');
    }
    const signedAt = Date.now();
    const withTs = { ...payload, signedAt, signature: '' };
    // Canonicalize: sort keys, omit signature
    const { signature, ...rest } = withTs;
    const sorted = Object.keys(rest).sort().reduce<Record<string, unknown>>((acc, k) => {
      acc[k] = (rest as Record<string, unknown>)[k];
      return acc;
    }, {});
    const bytes = new TextEncoder().encode(JSON.stringify(sorted));
    const sig = await this.crypto.sign(this.signingPrivateKey, bytes);
    return { ...payload, signedAt, signature: Buffer.from(sig).toString('base64') };
  }

  /** Register identity and upload prekey bundle + one-time prekeys */
  async register(
    address: string,
    identity: { signingKey: Uint8Array; dhKey: Uint8Array },
    signedPreKey: { keyId: number; publicKey: Uint8Array; signature: Uint8Array },
    oneTimePreKeys?: OneTimePreKey[],
  ): Promise<void> {
    const body: any = {
      address,
      identitySigningKey: toB64(identity.signingKey),
      identityDHKey: toB64(identity.dhKey),
      signedPreKey: {
        keyId: signedPreKey.keyId,
        publicKey: toB64(signedPreKey.publicKey),
        signature: toB64(signedPreKey.signature),
      },
    };

    if (oneTimePreKeys?.length) {
      body.oneTimePreKeys = oneTimePreKeys.map((k) => ({
        keyId: k.keyId,
        publicKey: toB64(k.keyPair.publicKey),
      }));
    }

    const signed = await this.sign(body);

    const res = await fetch(`${this.baseUrl}/v1/keys/register`, {
      method: 'POST',
      headers: this.headers(),
      body: JSON.stringify(signed),
    });
    if (!res.ok) throw new NetworkError(`Register failed: ${res.status}`, res.status);
  }

  /**
   * Fetch a prekey bundle for a peer (anonymous, consumes one one-time prekey).
   *
   * When the transport was created with `keyTransparency`, this also verifies
   * the inclusion proof and (if configured) feeds the STH into the supplied
   * `onObserveSth` hook for split-view tracking.
   */
  async fetchBundle(address: string): Promise<PreKeyBundle> {
    const result = await this.fetchBundleVerified(address);
    return result.bundle;
  }

  /**
   * Same as `fetchBundle` but returns the verified STH alongside the bundle
   * so callers can wire it into a `LightWitness`. Use this when you want
   * direct access to the proof side-channel.
   */
  async fetchBundleVerified(address: string): Promise<FetchBundleResult> {
    const res = await fetch(`${this.baseUrl}/v1/keys/bundle/${encodeURIComponent(address)}`, {
      headers: this.headers(),
    });
    if (!res.ok) {
      // KT-aware 404: read the body for an absence/tombstone proof so the
      // negative answer is also pinned to a tree state.
      if (res.status === 404 && this.kt) {
        const body = (await res.json().catch(() => null)) as
          | { ktProof?: KTProofWire }
          | null;
        if (body?.ktProof) {
          const proof = ktProofFromWire(body.ktProof);
          await this.verifyAbsenceOrTombstone(address, proof);
        } else if (this.kt.mode === 'observe-strict') {
          throw new KTVerificationError(
            `KT-strict: 404 for ${address} but no ktProof in response`,
          );
        }
      }
      throw new NetworkError(`Fetch bundle failed: ${res.status}`, res.status);
    }

    const data = (await res.json()) as {
      registrationId?: number;
      identitySigningKey: string;
      identityDHKey: string;
      signedPreKey: { keyId: number; publicKey: string; signature: string };
      oneTimePreKey?: { keyId: number; publicKey: string };
      ktProof?: KTProofWire;
    };
    const bundle: PreKeyBundle = {
      registrationId: data.registrationId ?? 0,
      identitySigningKey: fromB64(data.identitySigningKey),
      identityDHKey: fromB64(data.identityDHKey),
      signedPreKey: {
        keyId: data.signedPreKey.keyId,
        publicKey: fromB64(data.signedPreKey.publicKey),
        signature: fromB64(data.signedPreKey.signature),
      },
    };
    if (data.oneTimePreKey) {
      bundle.oneTimePreKey = {
        keyId: data.oneTimePreKey.keyId,
        publicKey: fromB64(data.oneTimePreKey.publicKey),
      };
    }

    let ktSth: SignedTreeHead | undefined;
    if (this.kt) {
      if (data.ktProof) {
        const proof = ktProofFromWire(data.ktProof);
        ktSth = await this.verifyInclusion(address, bundle, proof);
        if (this.kt.onObserveSth) await this.kt.onObserveSth(ktSth);
      } else if (this.kt.mode === 'observe-strict') {
        throw new KTVerificationError(
          `KT-strict: bundle for ${address} missing ktProof`,
        );
      }
    }

    return ktSth ? { bundle, ktSth } : { bundle };
  }

  private async verifyInclusion(
    address: string,
    bundle: PreKeyBundle,
    proof: KTProof,
  ): Promise<SignedTreeHead> {
    if (!this.kt) throw new Error('KT verifier not configured');
    const opts = {
      crypto: this.crypto,
      logPublicKey: this.kt.logPublicKey,
      ...(this.kt.maxStaleMs !== undefined ? { maxStaleMs: this.kt.maxStaleMs } : {}),
      ...(this.kt.now !== undefined ? { nowMs: this.kt.now() } : {}),
    };
    if (proof.body.kind === 'inclusion') {
      return verifyBundleInclusion(opts, address, bundle, proof);
    }
    if (proof.body.kind === 'tombstone') {
      // A tombstone proof on a 200 response would mean the server delivered
      // a deleted address — this is malicious.
      throw new KTVerificationError(`server returned tombstoned bundle for ${address}`);
    }
    if (proof.body.kind === 'absence') {
      throw new KTVerificationError(`server returned bundle but absence proof for ${address}`);
    }
    throw new KTSplitViewError(`unknown proof kind`);
  }

  private async verifyAbsenceOrTombstone(address: string, proof: KTProof): Promise<void> {
    if (!this.kt) return;
    const opts = {
      crypto: this.crypto,
      logPublicKey: this.kt.logPublicKey,
      ...(this.kt.maxStaleMs !== undefined ? { maxStaleMs: this.kt.maxStaleMs } : {}),
      ...(this.kt.now !== undefined ? { nowMs: this.kt.now() } : {}),
    };
    let sth: SignedTreeHead;
    if (proof.body.kind === 'absence') {
      sth = await verifyBundleAbsence(opts, address, proof);
    } else if (proof.body.kind === 'tombstone') {
      sth = await verifyBundleTombstone(opts, address, proof);
    } else {
      throw new KTVerificationError(
        `404 with non-absence/non-tombstone proof for ${address}`,
      );
    }
    if (this.kt.onObserveSth) await this.kt.onObserveSth(sth);
  }

  /** Upload additional one-time prekeys (signed) */
  async replenish(
    address: string,
    keys: Array<{ keyId: number; publicKey: Uint8Array }>,
  ): Promise<number> {
    const body = {
      address,
      oneTimePreKeys: keys.map((k) => ({
        keyId: k.keyId,
        publicKey: toB64(k.publicKey),
      })),
    };
    const signed = await this.sign(body);

    const res = await fetch(`${this.baseUrl}/v1/keys/replenish`, {
      method: 'POST',
      headers: this.headers(),
      body: JSON.stringify(signed),
    });
    if (!res.ok) throw new NetworkError(`Replenish failed: ${res.status}`, res.status);
    const data = (await res.json()) as { remaining: number };
    return data.remaining;
  }

  /** Get remaining one-time prekey count (anonymous) */
  async getKeyCount(address: string): Promise<number> {
    const res = await fetch(`${this.baseUrl}/v1/keys/count/${encodeURIComponent(address)}`, {
      headers: this.headers(),
    });
    if (!res.ok) throw new NetworkError(`Count failed: ${res.status}`, res.status);
    const data = (await res.json()) as { count: number };
    return data.count;
  }

  /** Delete all keys for this identity (signed) */
  async unregister(address: string): Promise<void> {
    const signed = await this.sign({ address });
    const res = await fetch(`${this.baseUrl}/v1/keys/${encodeURIComponent(address)}`, {
      method: 'DELETE',
      headers: this.headers(),
      body: JSON.stringify(signed),
    });
    if (!res.ok) throw new NetworkError(`Unregister failed: ${res.status}`, res.status);
  }
}

function toB64(bytes: Uint8Array): string {
  return Buffer.from(bytes).toString('base64');
}

function fromB64(str: string): Uint8Array {
  return new Uint8Array(Buffer.from(str, 'base64'));
}
