import { fromBase64, toBase64 } from "../crypto";

const encoder = new TextEncoder();
const MAGIC = "SSHSIG";
const GIT_NAMESPACE = "git";
const ARMOR = /-----BEGIN SSH SIGNATURE-----([\s\S]*?)-----END SSH SIGNATURE-----/;
const ED25519_KEY_BYTES = 32;
const UNCOMPRESSED_POINT = 4;

class SshReader {
  private position = 0;
  constructor(private readonly bytes: Uint8Array) {}
  raw(length: number): Uint8Array {
    if (this.position + length > this.bytes.byteLength) throw new Error("Truncated SSH data");
    const out = this.bytes.subarray(this.position, this.position + length);
    this.position += length;
    return out;
  }
  uint32(): number {
    const view = new DataView(this.bytes.buffer, this.bytes.byteOffset + this.position, 4);
    this.raw(4);
    return view.getUint32(0);
  }
  string(): Uint8Array {
    return this.raw(this.uint32());
  }
  text(): string {
    return new TextDecoder().decode(this.string());
  }
}

function sshString(value: Uint8Array | string): Uint8Array {
  const bytes = typeof value === "string" ? encoder.encode(value) : value;
  const out = new Uint8Array(4 + bytes.byteLength);
  new DataView(out.buffer).setUint32(0, bytes.byteLength);
  out.set(bytes, 4);
  return out;
}

function concat(parts: Uint8Array[]): Uint8Array<ArrayBuffer> {
  const out = new Uint8Array(parts.reduce((sum, part) => sum + part.byteLength, 0));
  let offset = 0;
  for (const part of parts) {
    out.set(part, offset);
    offset += part.byteLength;
  }
  return out;
}

function base64Url(bytes: Uint8Array): string {
  return toBase64(bytes).replaceAll("+", "-").replaceAll("/", "_").replace(/=+$/, "");
}

function stripMpint(bytes: Uint8Array): Uint8Array {
  let start = 0;
  while (start < bytes.byteLength - 1 && bytes[start] === 0) start++;
  return bytes.subarray(start);
}

function padTo(bytes: Uint8Array, length: number): Uint8Array {
  const stripped = stripMpint(bytes);
  if (stripped.byteLength > length) throw new Error("Integer too large");
  const out = new Uint8Array(length);
  out.set(stripped, length - stripped.byteLength);
  return out;
}

export type SshPublicKey = { type: string; blob: Uint8Array<ArrayBuffer>; fingerprint: string };

/** SHA256 fingerprint in the `SHA256:<base64>` form printed by `ssh-keygen -l`. */
export async function sshFingerprint(blob: Uint8Array<ArrayBuffer>): Promise<string> {
  const digest = new Uint8Array(await crypto.subtle.digest("SHA-256", blob));
  return `SHA256:${toBase64(digest).replace(/=+$/, "")}`;
}

const SUPPORTED_KEY_TYPES = new Set([
  "ssh-ed25519",
  "ssh-rsa",
  "ecdsa-sha2-nistp256",
  "ecdsa-sha2-nistp384",
  "ecdsa-sha2-nistp521",
]);

/** Parses an `authorized_keys`-style line such as `ssh-ed25519 AAAA... comment`. */
export async function parseSshPublicKey(line: string): Promise<SshPublicKey> {
  const [type, data] = line.trim().split(/\s+/);
  if (!type || !data || !SUPPORTED_KEY_TYPES.has(type)) throw new Error("Unsupported SSH key type");
  const blob = fromBase64(data);
  if (new SshReader(blob).text() !== type) throw new Error("SSH key type does not match its data");
  return { type, blob, fingerprint: await sshFingerprint(blob) };
}

type EcCurve = { curve: "P-256" | "P-384" | "P-521"; hash: string; size: number };
const EC_CURVES: Record<string, EcCurve> = {
  "ecdsa-sha2-nistp256": { curve: "P-256", hash: "SHA-256", size: 32 },
  "ecdsa-sha2-nistp384": { curve: "P-384", hash: "SHA-384", size: 48 },
  "ecdsa-sha2-nistp521": { curve: "P-521", hash: "SHA-512", size: 66 },
};

async function verifyRaw(
  keyBlob: Uint8Array,
  signatureAlgo: string,
  signature: Uint8Array,
  data: Uint8Array<ArrayBuffer>,
): Promise<boolean> {
  const key = new SshReader(keyBlob);
  const type = key.text();
  const sig = new Uint8Array(signature);
  if (type === "ssh-ed25519") {
    const publicKey = key.string();
    if (publicKey.byteLength !== ED25519_KEY_BYTES) return false;
    const imported = await crypto.subtle.importKey("raw", new Uint8Array(publicKey), { name: "Ed25519" }, false, [
      "verify",
    ]);
    return crypto.subtle.verify({ name: "Ed25519" }, imported, sig, data);
  }
  if (type === "ssh-rsa") {
    const exponent = stripMpint(key.string());
    const modulus = stripMpint(key.string());
    const hash = signatureAlgo === "rsa-sha2-512" ? "SHA-512" : signatureAlgo === "rsa-sha2-256" ? "SHA-256" : null;
    if (!hash) return false;
    const imported = await crypto.subtle.importKey(
      "jwk",
      { kty: "RSA", n: base64Url(modulus), e: base64Url(exponent), ext: true },
      { name: "RSASSA-PKCS1-v1_5", hash },
      false,
      ["verify"],
    );
    return crypto.subtle.verify({ name: "RSASSA-PKCS1-v1_5" }, imported, sig, data);
  }
  const curve = EC_CURVES[type];
  if (!curve) return false;
  key.string();
  const point = key.string();
  if (point[0] !== UNCOMPRESSED_POINT || point.byteLength !== 1 + 2 * curve.size) return false;
  const imported = await crypto.subtle.importKey(
    "jwk",
    {
      kty: "EC",
      crv: curve.curve,
      x: base64Url(point.subarray(1, 1 + curve.size)),
      y: base64Url(point.subarray(1 + curve.size)),
      ext: true,
    },
    { name: "ECDSA", namedCurve: curve.curve },
    false,
    ["verify"],
  );
  const parts = new SshReader(signature);
  const rawSignature = concat([padTo(parts.string(), curve.size), padTo(parts.string(), curve.size)]);
  return crypto.subtle.verify({ name: "ECDSA", hash: curve.hash }, imported, rawSignature, data);
}

export type SshSignature = { publicKey: Uint8Array<ArrayBuffer>; verify: () => Promise<boolean> };

/** Parses an armored SSHSIG over `payload` in the `git` namespace; `verify` checks it against the embedded key. */
export function parseSshSignature(armored: string, payload: Uint8Array): SshSignature | null {
  const match = ARMOR.exec(armored);
  if (!match) return null;
  const bytes = fromBase64(match[1].replace(/\s+/g, ""));
  const reader = new SshReader(bytes);
  if (new TextDecoder().decode(reader.raw(MAGIC.length)) !== MAGIC) return null;
  if (reader.uint32() !== 1) return null;
  const publicKey = new Uint8Array(reader.string());
  const namespace = reader.text();
  const reserved = reader.string();
  const hashAlgorithm = reader.text();
  const signatureBlob = new SshReader(reader.string());
  const signatureAlgo = signatureBlob.text();
  const signature = signatureBlob.string();
  const hashName = hashAlgorithm === "sha512" ? "SHA-512" : hashAlgorithm === "sha256" ? "SHA-256" : null;
  return {
    publicKey,
    verify: async () => {
      if (namespace !== GIT_NAMESPACE || !hashName) return false;
      const digest = new Uint8Array(await crypto.subtle.digest(hashName, new Uint8Array(payload)));
      const signed = concat([
        encoder.encode(MAGIC),
        sshString(namespace),
        sshString(reserved),
        sshString(hashAlgorithm),
        sshString(digest),
      ]);
      return verifyRaw(publicKey, signatureAlgo, signature, signed);
    },
  };
}