export interface H264AccessUnit {
  type: "key" | "delta";
  data: Uint8Array;
}

export interface FrameThumb {
  index: number;
  src: string;
  width: number;
  height: number;
  timestampUs: number;
}

export interface ParsedH264 {
  codec: string;
  description: Uint8Array;
  accessUnits: H264AccessUnit[];
  frameCount: number;
  keyFrameCount: number;
}

export interface H264DecoderConfig {
  codec: string;
  description: Uint8Array;
}

export interface H264StreamChunk {
  accessUnits: H264AccessUnit[];
  config?: H264DecoderConfig;
  frameCount: number;
  keyFrameCount: number;
  bufferedBytes: number;
  done: boolean;
}

interface RawNalUnit {
  type: number;
  data: Uint8Array;
}

export function parseH264AnnexB(bytes: Uint8Array): ParsedH264 {
  const parser = new H264AnnexBStreamParser();
  const first = parser.push(bytes);
  const last = parser.flush();
  const accessUnits = [...first.accessUnits, ...last.accessUnits];
  const config = first.config ?? last.config;

  if (!config) {
    throw new Error("没有找到 SPS / PPS，无法配置 WebCodecs 解码器。");
  }

  const codec = config.codec;
  const description = config.description;
  const keyFrameCount = accessUnits.filter((unit) => unit.type === "key").length;

  return {
    codec,
    description,
    accessUnits,
    frameCount: accessUnits.length,
    keyFrameCount,
  };
}

export class H264AnnexBStreamParser {
  private buffer: Uint8Array<ArrayBufferLike> = new Uint8Array(0);
  private spsUnits: Uint8Array[] = [];
  private ppsUnits: Uint8Array[] = [];
  private currentSlices: Uint8Array[] = [];
  private currentIsKey = false;
  private currentHasVcl = false;
  private frameCount = 0;
  private keyFrameCount = 0;
  private config: H264DecoderConfig | null = null;
  private configSignature = "";

  push(chunk: Uint8Array): H264StreamChunk {
    this.append(chunk);
    return this.drain(false);
  }

  flush(): H264StreamChunk {
    return this.drain(true);
  }

  private append(chunk: Uint8Array) {
    if (chunk.length === 0) {
      return;
    }

    const merged = new Uint8Array(this.buffer.length + chunk.length);
    merged.set(this.buffer, 0);
    merged.set(chunk, this.buffer.length);
    this.buffer = merged;
  }

  private drain(final: boolean): H264StreamChunk {
    const accessUnits: H264AccessUnit[] = [];

    while (true) {
      const starts = findStartCodes(this.buffer);

      if (starts.length === 0) {
        this.buffer = final ? new Uint8Array(0) : keepTail(this.buffer);
        break;
      }

      if (starts[0].index > 0) {
        this.buffer = this.buffer.slice(starts[0].index);
        continue;
      }

      const processCount = final ? starts.length : starts.length - 1;
      if (processCount <= 0) {
        break;
      }

      for (let index = 0; index < processCount; index += 1) {
        const start = starts[index];
        const next = starts[index + 1];
        const nalStart = start.index + start.length;
        const nalEnd = next ? next.index : this.buffer.length;

        if (nalEnd <= nalStart) {
          continue;
        }

        const nal = {
          type: this.buffer[nalStart] & 0x1f,
          data: copyBytes(this.buffer.subarray(nalStart, nalEnd)),
        };

        const flushed = this.processNal(nal);
        if (flushed) {
          accessUnits.push(flushed);
        }
      }

      if (final) {
        this.buffer = new Uint8Array(0);
        break;
      }

      this.buffer = this.buffer.slice(starts[processCount].index);
    }

    const finalAccessUnit = final ? this.flushCurrentAccessUnit() : null;
    if (finalAccessUnit) {
      accessUnits.push(finalAccessUnit);
    }

    const config = this.ensureConfig();

    return {
      accessUnits,
      config: config ?? undefined,
      frameCount: this.frameCount,
      keyFrameCount: this.keyFrameCount,
      bufferedBytes: this.buffer.length,
      done: final && this.buffer.length === 0,
    };
  }

  private processNal(nal: RawNalUnit): H264AccessUnit | null {
    if (nal.type === 7) {
      if (!containsBytes(this.spsUnits, nal.data)) {
        this.spsUnits.push(nal.data);
        this.configSignature = "";
        this.config = null;
      }
      return null;
    }

    if (nal.type === 8) {
      if (!containsBytes(this.ppsUnits, nal.data)) {
        this.ppsUnits.push(nal.data);
        this.configSignature = "";
        this.config = null;
      }
      return null;
    }

    if (nal.type !== 1 && nal.type !== 5) {
      return null;
    }

    const firstMbInSlice = readFirstMbInSlice(nal.data);
    if (firstMbInSlice === 0 && this.currentHasVcl) {
      const flushed = this.flushCurrentAccessUnit();
      this.currentSlices = [nal.data];
      this.currentHasVcl = true;
      this.currentIsKey = nal.type === 5;
      return flushed;
    }

    this.currentSlices.push(nal.data);
    this.currentHasVcl = true;
    this.currentIsKey = this.currentIsKey || nal.type === 5;
    return null;
  }

  private flushCurrentAccessUnit(): H264AccessUnit | null {
    if (this.currentSlices.length === 0) {
      this.currentHasVcl = false;
      this.currentIsKey = false;
      return null;
    }

    const accessUnit: H264AccessUnit = {
      type: this.currentIsKey ? "key" : "delta",
      data: concatNalPayloads(this.currentSlices),
    };

    this.frameCount += 1;
    if (accessUnit.type === "key") {
      this.keyFrameCount += 1;
    }

    this.currentSlices = [];
    this.currentHasVcl = false;
    this.currentIsKey = false;
    return accessUnit;
  }

  private ensureConfig(): H264DecoderConfig | null {
    if (this.spsUnits.length === 0 || this.ppsUnits.length === 0) {
      return null;
    }

    const signature = `${this.spsUnits.map(bytesToHex).join("|")}::${this.ppsUnits
      .map(bytesToHex)
      .join("|")}`;
    if (this.config && this.configSignature === signature) {
      return this.config;
    }

    this.config = {
      codec: codecStringFromSps(this.spsUnits[0]),
      description: buildAvcDecoderConfigRecord(this.spsUnits, this.ppsUnits),
    };
    this.configSignature = signature;
    return this.config;
  }
}

function findStartCodes(bytes: Uint8Array): Array<{ index: number; length: number }> {
  const starts: Array<{ index: number; length: number }> = [];

  for (let index = 0; index < bytes.length - 3; index += 1) {
    if (bytes[index] !== 0 || bytes[index + 1] !== 0) {
      continue;
    }

    if (bytes[index + 2] === 1) {
      starts.push({ index, length: 3 });
      index += 2;
      continue;
    }

    if (index < bytes.length - 4 && bytes[index + 2] === 0 && bytes[index + 3] === 1) {
      starts.push({ index, length: 4 });
      index += 3;
    }
  }

  return starts;
}

function keepTail(bytes: Uint8Array): Uint8Array {
  if (bytes.length <= 3) {
    return bytes.slice();
  }

  return bytes.slice(bytes.length - 3);
}

function extractNalUnits(bytes: Uint8Array): RawNalUnit[] {
  const starts: Array<{ index: number; length: number }> = [];

  for (let index = 0; index < bytes.length - 3; index += 1) {
    if (bytes[index] !== 0 || bytes[index + 1] !== 0) {
      continue;
    }

    if (bytes[index + 2] === 1) {
      starts.push({ index, length: 3 });
      index += 2;
      continue;
    }

    if (index < bytes.length - 4 && bytes[index + 2] === 0 && bytes[index + 3] === 1) {
      starts.push({ index, length: 4 });
      index += 3;
    }
  }

  if (starts.length === 0) {
    throw new Error("没有找到 Annex-B start code。");
  }

  const nalUnits: RawNalUnit[] = [];

  for (let i = 0; i < starts.length; i += 1) {
    const start = starts[i];
    const next = starts[i + 1];
    const nalStart = start.index + start.length;
    const nalEnd = next ? next.index : bytes.length;

    if (nalEnd <= nalStart) {
      continue;
    }

    const data = bytes.subarray(nalStart, nalEnd);
    nalUnits.push({
      type: data[0] & 0x1f,
      data,
    });
  }

  return nalUnits;
}

function readFirstMbInSlice(nalData: Uint8Array): number {
  if (nalData.length < 2) {
    return 0;
  }

  const rbsp = removeEmulationPreventionBytes(nalData.subarray(1));
  const reader = new BitReader(rbsp);
  return reader.readUE();
}

function codecStringFromSps(sps: Uint8Array): string {
  if (sps.length < 4) {
    throw new Error("SPS 长度不足，无法生成 codec 字符串。");
  }

  const profile = sps[1].toString(16).padStart(2, "0");
  const compatibility = sps[2].toString(16).padStart(2, "0");
  const level = sps[3].toString(16).padStart(2, "0");

  return `avc1.${profile}${compatibility}${level}`;
}

function buildAvcDecoderConfigRecord(
  spsUnits: Uint8Array[],
  ppsUnits: Uint8Array[],
): Uint8Array {
  const spsList = uniqueByContent(spsUnits);
  const ppsList = uniqueByContent(ppsUnits);

  let size = 7;
  for (const sps of spsList) {
    size += 2 + sps.length;
  }
  size += 1;
  for (const pps of ppsList) {
    size += 2 + pps.length;
  }

  const record = new Uint8Array(size);
  let offset = 0;

  record[offset++] = 1;
  record[offset++] = spsList[0][1];
  record[offset++] = spsList[0][2];
  record[offset++] = spsList[0][3];
  record[offset++] = 0xff;
  record[offset++] = 0xe0 | spsList.length;

  for (const sps of spsList) {
    writeUint16(record, offset, sps.length);
    offset += 2;
    record.set(sps, offset);
    offset += sps.length;
  }

  record[offset++] = ppsList.length;

  for (const pps of ppsList) {
    writeUint16(record, offset, pps.length);
    offset += 2;
    record.set(pps, offset);
    offset += pps.length;
  }

  return record;
}

function concatNalPayloads(units: Uint8Array[]): Uint8Array {
  let size = 0;
  for (const unit of units) {
    size += 4 + unit.length;
  }

  const output = new Uint8Array(size);
  let offset = 0;

  for (const unit of units) {
    writeUint32(output, offset, unit.length);
    offset += 4;
    output.set(unit, offset);
    offset += unit.length;
  }

  return output;
}

function removeEmulationPreventionBytes(data: Uint8Array): Uint8Array {
  const output: number[] = [];
  let zeroCount = 0;

  for (const byte of data) {
    if (zeroCount >= 2 && byte === 0x03) {
      zeroCount = 0;
      continue;
    }

    output.push(byte);
    zeroCount = byte === 0x00 ? zeroCount + 1 : 0;
  }

  return Uint8Array.from(output);
}

function uniqueByContent(units: Uint8Array[]): Uint8Array[] {
  const seen = new Set<string>();
  const result: Uint8Array[] = [];

  for (const unit of units) {
    const key = bytesToHex(unit);
    if (seen.has(key)) {
      continue;
    }

    seen.add(key);
    result.push(unit);
  }

  return result;
}

function containsBytes(units: Uint8Array[], candidate: Uint8Array): boolean {
  const candidateKey = bytesToHex(candidate);
  return units.some((unit) => bytesToHex(unit) === candidateKey);
}

function copyBytes(bytes: Uint8Array): Uint8Array {
  return bytes.slice();
}

function bytesToHex(bytes: Uint8Array): string {
  let output = "";
  for (const byte of bytes) {
    output += byte.toString(16).padStart(2, "0");
  }
  return output;
}

function writeUint16(target: Uint8Array, offset: number, value: number) {
  target[offset] = (value >> 8) & 0xff;
  target[offset + 1] = value & 0xff;
}

function writeUint32(target: Uint8Array, offset: number, value: number) {
  target[offset] = (value >>> 24) & 0xff;
  target[offset + 1] = (value >>> 16) & 0xff;
  target[offset + 2] = (value >>> 8) & 0xff;
  target[offset + 3] = value & 0xff;
}

class BitReader {
  private readonly data: Uint8Array;
  private byteIndex = 0;
  private bitIndex = 7;

  constructor(data: Uint8Array) {
    this.data = data;
  }

  readBit(): number {
    if (this.byteIndex >= this.data.length) {
      return 0;
    }

    const bit = (this.data[this.byteIndex] >> this.bitIndex) & 1;
    this.bitIndex -= 1;

    if (this.bitIndex < 0) {
      this.bitIndex = 7;
      this.byteIndex += 1;
    }

    return bit;
  }

  readBits(count: number): number {
    let value = 0;
    for (let index = 0; index < count; index += 1) {
      value = (value << 1) | this.readBit();
    }
    return value;
  }

  readUE(): number {
    let zeroCount = 0;
    while (this.readBit() === 0 && zeroCount < 31) {
      zeroCount += 1;
    }

    const suffix = zeroCount > 0 ? this.readBits(zeroCount) : 0;
    return (1 << zeroCount) - 1 + suffix;
  }
}