import { WebSocket } from "ws";
import type { BufferedAudioFrame } from "../shared/realtime-voice-protocol";
import type {
  RealtimeSttRecognizer,
  SendVoiceMessage,
  SttEnvironment,
} from "./realtime-stt-provider";

const CONNECT_TIMEOUT_MS = 8000;
const MAX_PENDING_FRAMES = 600;

function transcriptWords(text: string): string[] {
  return text
    .replace(/[^\p{L}\p{N}]+/gu, " ")
    .trim()
    .split(/\s+/)
    .filter(Boolean);
}

function endpointingMs(environment: SttEnvironment): number {
  const value = Number(environment.SHIHAB_ENDPOINTING_MS ?? 800);
  return Number.isFinite(value) && value >= 300 && value <= 3000
    ? Math.round(value)
    : 800;
}

export function shihabWebSocketUrl(
  environment: SttEnvironment = process.env,
): string {
  const url = new URL("wss://api.shihab.co.uk/v1/listen/stream");
  url.search = new URLSearchParams({
    model: environment.SHIHAB_MODEL?.trim() || "shihab-asr-v1",
    language: "ar-PS",
    encoding: "linear16",
    sample_rate: "16000",
    channels: "1",
    interim_results: "true",
    endpointing: String(endpointingMs(environment)),
    store: "false",
  }).toString();
  return url.toString();
}

export class ShihabRealtimeRecognizer implements RealtimeSttRecognizer {
  private socket: WebSocket | null = null;
  private connectPromise: Promise<void> | null = null;
  private active = true;
  private ready = false;
  private lastAcknowledgedSequence = -1;
  private pendingFrames = new Map<number, BufferedAudioFrame>();
  private finalEventHistory = new Map<number, Record<string, unknown>>();
  private latestInterimEvent: Record<string, unknown> | null = null;
  private utteranceIndex = 0;
  private revision = 0;

  constructor(
    readonly sessionId: string,
    private readonly emit: SendVoiceMessage,
    private readonly environment: SttEnvironment = process.env,
  ) {}

  start(): Promise<void> {
    this.active = true;
    return this.ensureConnected();
  }

  accept(frame: BufferedAudioFrame): void {
    if (frame.audioSequence <= this.lastAcknowledgedSequence) {
      this.emitAck();
      return;
    }
    if (!this.pendingFrames.has(frame.audioSequence)) {
      if (this.pendingFrames.size >= MAX_PENDING_FRAMES) {
        this.emitError("Shihab replay buffer is full");
        return;
      }
      this.pendingFrames.set(frame.audioSequence, frame);
    }
    this.sendFrame(frame);
    if (!this.socket || this.socket.readyState !== WebSocket.OPEN) {
      void this.ensureConnected().catch((error) => this.emitError(error));
    }
  }

  stop(): void {
    this.active = false;
    if (this.socket?.readyState === WebSocket.OPEN) {
      this.socket.send(JSON.stringify({ type: "CloseStream" }));
    }
  }

  close(): void {
    this.active = false;
    this.socket?.close();
    this.socket = null;
    this.connectPromise = null;
    this.ready = false;
  }

  acknowledgedSequence(): number {
    return this.lastAcknowledgedSequence;
  }

  replayTranscriptEvents(): void {
    for (const event of [...this.finalEventHistory.entries()]
      .sort(([left], [right]) => left - right)
      .map(([, value]) => value)) {
      this.emit(event);
    }
    if (this.latestInterimEvent) this.emit(this.latestInterimEvent);
  }

  private ensureConnected(): Promise<void> {
    if (this.socket?.readyState === WebSocket.OPEN && this.ready) {
      return Promise.resolve();
    }
    if (this.connectPromise) return this.connectPromise;

    const apiKey = this.environment.SHIHAB_API_KEY?.trim();
    if (!apiKey) return Promise.reject(new Error("SHIHAB_API_KEY is missing"));

    const promise = new Promise<void>((resolve, reject) => {
      const socket = new WebSocket(shihabWebSocketUrl(this.environment), {
        headers: { Authorization: `Bearer ${apiKey}` },
        perMessageDeflate: false,
      });
      this.socket = socket;
      let settled = false;
      const timeout = setTimeout(() => {
        if (settled) return;
        settled = true;
        socket.close();
        reject(new Error("Timed out connecting to ShihabASR"));
      }, CONNECT_TIMEOUT_MS);

      socket.on("open", () => {
        if (settled) return;
        settled = true;
        clearTimeout(timeout);
        this.ready = true;
        for (const frame of [...this.pendingFrames.values()].sort(
          (left, right) => left.audioSequence - right.audioSequence,
        )) {
          this.sendFrame(frame);
        }
        resolve();
      });

      socket.on("message", (raw, isBinary) => {
        if (isBinary) return;
        let message: Record<string, any>;
        try {
          message = JSON.parse(raw.toString());
        } catch {
          return;
        }
        if (message.type === "Results") this.forwardResult(message);
        else if (message.type === "Error") {
          this.emitError(
            `${message.code ?? "shihab_error"}: ${message.message ?? "ShihabASR error"}`,
          );
        }
      });

      socket.on("error", (error) => {
        if (!settled) {
          settled = true;
          clearTimeout(timeout);
          reject(error);
        } else this.emitError(error);
      });

      socket.on("close", (code) => {
        clearTimeout(timeout);
        if (this.socket === socket) this.socket = null;
        this.ready = false;
        this.connectPromise = null;
        if (this.active && code !== 1000) {
          this.emitError(`ShihabASR stream closed (${code})`);
        }
      });
    }).finally(() => {
      this.connectPromise = null;
    });

    this.connectPromise = promise;
    return promise;
  }

  private sendFrame(frame: BufferedAudioFrame): void {
    if (!this.ready || this.socket?.readyState !== WebSocket.OPEN) return;
    this.socket.send(Buffer.from(frame.audio, "base64"));
    this.lastAcknowledgedSequence = Math.max(
      this.lastAcknowledgedSequence,
      frame.audioSequence,
    );
    this.pendingFrames.delete(frame.audioSequence);
    this.emitAck();
  }

  private emitAck(): void {
    this.emit({
      type: "audio_ack",
      sessionId: this.sessionId,
      audioSequence: this.lastAcknowledgedSequence,
    });
  }

  private forwardResult(message: Record<string, any>): void {
    const text = String(message.transcript ?? "").trim();
    if (!text) return;
    const isFinal = message.is_final === true;
    if (this.environment.STT_LOG_TRANSCRIPTS?.trim().toLowerCase() === "true") {
      console.log(
        `[voice-ws][shihab][${this.sessionId}] ${isFinal ? "final" : "interim"}: ${JSON.stringify(text)}`,
      );
    }
    const event: Record<string, unknown> = {
      type: isFinal ? "final" : "interim",
      sessionId: this.sessionId,
      utteranceId: `${this.sessionId}:${this.utteranceIndex}`,
      utteranceIndex: this.utteranceIndex,
      revision: ++this.revision,
      audioSequence: this.lastAcknowledgedSequence,
      text,
      words: isFinal ? transcriptWords(text) : undefined,
    };
    if (isFinal) {
      this.finalEventHistory.set(this.utteranceIndex, event);
      this.latestInterimEvent = null;
      this.utteranceIndex++;
      this.revision = 0;
    } else this.latestInterimEvent = event;
    this.emit(event);
  }

  private emitError(error: unknown): void {
    this.emit({
      type: "error",
      sessionId: this.sessionId,
      message: error instanceof Error ? error.message : String(error),
    });
  }
}