import type {
  ChatCitationsEvent,
  ChatDeltaEvent,
  ChatDoneEvent,
  ChatErrorEvent,
  ChatEvidence,
  ChatMetaEvent,
  ChatRetrievalEvent,
  ChatStreamEvent,
  ChatTimings,
  ChatUsageEvent,
} from "./types";

export type ChatStreamErrorKind =
  | "aborted"
  | "forbidden"
  | "not_ready"
  | "unavailable"
  | "validation"
  | "network"
  | "invalid_stream"
  | "http";

export class ChatStreamError extends Error {
  readonly kind: ChatStreamErrorKind;
  readonly status: number | null;

  constructor(kind: ChatStreamErrorKind, message: string, status: number | null = null) {
    super(message);
    this.name = "ChatStreamError";
    this.kind = kind;
    this.status = status;
  }
}

const UUID_PATTERN = /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i;
const PROFILE_HASH_PATTERN = /^[0-9a-f]{64}$/;
const LABEL_PATTERN = /^S[1-9]\d*$/;
const EVENT_NAMES = new Set(["meta", "retrieval", "delta", "citations", "usage", "done", "error"]);

type JsonObject = Record<string, unknown>;
type StreamStage =
  | "expect_meta"
  | "expect_retrieval"
  | "expect_delta_or_error"
  | "accepting_delta"
  | "expect_usage"
  | "expect_done"
  | "terminal";

function invalidStream(): never {
  throw new ChatStreamError(
    "invalid_stream",
    "回答流不完整或格式无效。已停止展示后续内容，请重新发起问题。",
  );
}

function isObject(value: unknown): value is JsonObject {
  return typeof value === "object" && value !== null && !Array.isArray(value);
}

function hasExactKeys(value: JsonObject, keys: readonly string[]): boolean {
  const actual = Object.keys(value).sort();
  const expected = [...keys].sort();
  return actual.length === expected.length && actual.every((key, index) => key === expected[index]);
}

function isFiniteNumber(
  value: unknown,
  minimum: number,
  maximum = Number.POSITIVE_INFINITY,
): value is number {
  return (
    typeof value === "number" && Number.isFinite(value) && value >= minimum && value <= maximum
  );
}

function isPositiveInteger(value: unknown): value is number {
  return Number.isInteger(value) && typeof value === "number" && value >= 1;
}

function isNullableNonnegativeInteger(value: unknown): value is number | null {
  return value === null || (Number.isInteger(value) && typeof value === "number" && value >= 0);
}

function isNullableString(value: unknown): value is string | null {
  return value === null || typeof value === "string";
}

function parseProfile(value: unknown) {
  if (
    !isObject(value) ||
    !hasExactKeys(value, ["profile_hash", "model", "dimension", "synthetic"]) ||
    typeof value.profile_hash !== "string" ||
    !PROFILE_HASH_PATTERN.test(value.profile_hash) ||
    typeof value.model !== "string" ||
    value.model.length === 0 ||
    value.dimension !== 1024 ||
    typeof value.synthetic !== "boolean"
  ) {
    invalidStream();
  }
  return {
    profile_hash: value.profile_hash,
    model: value.model,
    dimension: 1024 as const,
    synthetic: value.synthetic,
  };
}

function parseEvidence(value: unknown): ChatEvidence {
  const keys = [
    "label",
    "rank",
    "vector_rank",
    "citation_id",
    "document_id",
    "source_name",
    "snippet",
    "section_path",
    "page_start",
    "page_end",
    "page_label",
    "vector_score",
    "rerank_score",
  ] as const;
  if (!isObject(value) || !hasExactKeys(value, keys)) invalidStream();
  if (
    typeof value.label !== "string" ||
    !LABEL_PATTERN.test(value.label) ||
    !isPositiveInteger(value.rank) ||
    value.label !== `S${value.rank}` ||
    !isPositiveInteger(value.vector_rank) ||
    typeof value.citation_id !== "string" ||
    !UUID_PATTERN.test(value.citation_id) ||
    typeof value.document_id !== "string" ||
    !UUID_PATTERN.test(value.document_id) ||
    typeof value.source_name !== "string" ||
    value.source_name.length < 1 ||
    value.source_name.length > 240 ||
    typeof value.snippet !== "string" ||
    value.snippet.length < 1 ||
    value.snippet.length > 1_200 ||
    !Array.isArray(value.section_path) ||
    value.section_path.some((part) => typeof part !== "string") ||
    (value.page_start !== null && !isPositiveInteger(value.page_start)) ||
    (value.page_end !== null && !isPositiveInteger(value.page_end)) ||
    (value.page_start === null) !== (value.page_end === null) ||
    (typeof value.page_start === "number" &&
      typeof value.page_end === "number" &&
      value.page_end < value.page_start) ||
    typeof value.page_label !== "string" ||
    value.page_label.length === 0 ||
    !isFiniteNumber(value.vector_score, -1, 1) ||
    (value.rerank_score !== null && !isFiniteNumber(value.rerank_score, 0, 1))
  ) {
    invalidStream();
  }
  return {
    label: value.label,
    rank: value.rank,
    vector_rank: value.vector_rank,
    citation_id: value.citation_id,
    document_id: value.document_id,
    source_name: value.source_name,
    snippet: value.snippet,
    section_path: value.section_path,
    page_start: value.page_start,
    page_end: value.page_end,
    page_label: value.page_label,
    vector_score: value.vector_score,
    rerank_score: value.rerank_score,
  };
}

function parseTimings(value: unknown): ChatTimings {
  if (
    !isObject(value) ||
    !hasExactKeys(value, ["embedding_ms", "database_ms", "rerank_ms", "total_ms"]) ||
    !isFiniteNumber(value.embedding_ms, 0) ||
    !isFiniteNumber(value.database_ms, 0) ||
    !isFiniteNumber(value.rerank_ms, 0) ||
    !isFiniteNumber(value.total_ms, 0)
  ) {
    invalidStream();
  }
  return {
    embedding_ms: value.embedding_ms,
    database_ms: value.database_ms,
    rerank_ms: value.rerank_ms,
    total_ms: value.total_ms,
  };
}

function parseMeta(value: JsonObject): ChatMetaEvent {
  if (
    !hasExactKeys(value, ["seq", "trace_id", "knowledge_base_id", "profile", "generation_mode"]) ||
    !isPositiveInteger(value.seq) ||
    typeof value.trace_id !== "string" ||
    value.trace_id.length === 0 ||
    typeof value.knowledge_base_id !== "string" ||
    !UUID_PATTERN.test(value.knowledge_base_id) ||
    (value.generation_mode !== "synthetic_extractive" && value.generation_mode !== "cloud_grounded")
  ) {
    invalidStream();
  }
  const profile = parseProfile(value.profile);
  if (
    (profile.synthetic && value.generation_mode !== "synthetic_extractive") ||
    (!profile.synthetic && value.generation_mode !== "cloud_grounded")
  ) {
    invalidStream();
  }
  return {
    name: "meta",
    seq: value.seq,
    trace_id: value.trace_id,
    knowledge_base_id: value.knowledge_base_id,
    profile,
    generation_mode: value.generation_mode,
  };
}

function parseRetrieval(value: JsonObject): ChatRetrievalEvent {
  if (
    !hasExactKeys(value, [
      "seq",
      "status",
      "rerank_status",
      "degradation_reason",
      "evidence",
      "timings",
    ]) ||
    !isPositiveInteger(value.seq) ||
    (value.status !== "ok" && value.status !== "empty") ||
    !["applied", "degraded", "skipped_empty"].includes(String(value.rerank_status)) ||
    (value.degradation_reason !== null && value.degradation_reason !== "rerank_unavailable") ||
    !Array.isArray(value.evidence)
  ) {
    invalidStream();
  }
  const evidence = value.evidence.map(parseEvidence);
  if (new Set(evidence.map((item) => item.citation_id)).size !== evidence.length) invalidStream();
  if (evidence.some((item, index) => item.rank !== index + 1)) invalidStream();
  if ((value.status === "empty") !== (evidence.length === 0)) invalidStream();
  if ((value.status === "empty") !== (value.rerank_status === "skipped_empty")) invalidStream();
  if ((value.rerank_status === "degraded") !== (value.degradation_reason !== null)) invalidStream();
  return {
    name: "retrieval",
    seq: value.seq,
    status: value.status,
    rerank_status: value.rerank_status as "applied" | "degraded" | "skipped_empty",
    degradation_reason: value.degradation_reason,
    evidence,
    timings: parseTimings(value.timings),
  };
}

function parseDelta(value: JsonObject): ChatDeltaEvent {
  if (
    !hasExactKeys(value, ["seq", "text"]) ||
    !isPositiveInteger(value.seq) ||
    typeof value.text !== "string"
  ) {
    invalidStream();
  }
  return { name: "delta", seq: value.seq, text: value.text };
}

function parseCitations(value: JsonObject): ChatCitationsEvent {
  if (
    !hasExactKeys(value, ["seq", "citations"]) ||
    !isPositiveInteger(value.seq) ||
    !Array.isArray(value.citations)
  ) {
    invalidStream();
  }
  return { name: "citations", seq: value.seq, citations: value.citations.map(parseEvidence) };
}

function parseUsage(value: JsonObject): ChatUsageEvent {
  if (
    !hasExactKeys(value, [
      "seq",
      "model",
      "request_id",
      "input_tokens",
      "output_tokens",
      "total_tokens",
    ]) ||
    !isPositiveInteger(value.seq) ||
    typeof value.model !== "string" ||
    value.model.length === 0 ||
    !isNullableString(value.request_id) ||
    !isNullableNonnegativeInteger(value.input_tokens) ||
    !isNullableNonnegativeInteger(value.output_tokens) ||
    !isNullableNonnegativeInteger(value.total_tokens)
  ) {
    invalidStream();
  }
  return {
    name: "usage",
    seq: value.seq,
    model: value.model,
    request_id: value.request_id,
    input_tokens: value.input_tokens,
    output_tokens: value.output_tokens,
    total_tokens: value.total_tokens,
  };
}

function parseDone(value: JsonObject): ChatDoneEvent {
  if (
    !hasExactKeys(value, ["seq", "status", "answer_mode", "finish_reason"]) ||
    !isPositiveInteger(value.seq) ||
    value.status !== "complete" ||
    !["grounded", "refused", "retrieval_only"].includes(String(value.answer_mode)) ||
    !isNullableString(value.finish_reason)
  ) {
    invalidStream();
  }
  return {
    name: "done",
    seq: value.seq,
    status: "complete",
    answer_mode: value.answer_mode as "grounded" | "refused" | "retrieval_only",
    finish_reason: value.finish_reason,
  };
}

function parseError(value: JsonObject): ChatErrorEvent {
  if (
    !hasExactKeys(value, ["seq", "status", "code", "title", "retryable", "answer_mode"]) ||
    !isPositiveInteger(value.seq) ||
    value.status !== "error" ||
    (value.code !== "CHAT_PROVIDER_UNAVAILABLE" && value.code !== "CHAT_GENERATION_FAILED") ||
    typeof value.title !== "string" ||
    value.title.length === 0 ||
    typeof value.retryable !== "boolean" ||
    value.answer_mode !== "retrieval_only"
  ) {
    invalidStream();
  }
  return {
    name: "error",
    seq: value.seq,
    status: "error",
    code: value.code,
    title: value.title,
    retryable: value.retryable,
    answer_mode: "retrieval_only",
  };
}

function parseEvent(name: string, data: string): ChatStreamEvent {
  if (!EVENT_NAMES.has(name)) invalidStream();
  let value: unknown;
  try {
    value = JSON.parse(data);
  } catch {
    invalidStream();
  }
  if (!isObject(value)) invalidStream();
  switch (name) {
    case "meta":
      return parseMeta(value);
    case "retrieval":
      return parseRetrieval(value);
    case "delta":
      return parseDelta(value);
    case "citations":
      return parseCitations(value);
    case "usage":
      return parseUsage(value);
    case "done":
      return parseDone(value);
    case "error":
      return parseError(value);
    default:
      return invalidStream();
  }
}

function evidenceMatches(left: ChatEvidence, right: ChatEvidence): boolean {
  return JSON.stringify(left) === JSON.stringify(right);
}

export class ChatEventSequence {
  private stage: StreamStage = "expect_meta";
  private expectedSeq = 1;
  private expectedKnowledgeBaseId: string;
  private retrievalEvidence: readonly ChatEvidence[] = [];
  private citations: readonly ChatEvidence[] = [];
  private generatedCharacters = 0;
  private generatedText = "";

  constructor(expectedKnowledgeBaseId: string) {
    this.expectedKnowledgeBaseId = expectedKnowledgeBaseId;
  }

  accept(event: ChatStreamEvent): void {
    if (this.stage === "terminal" || event.seq !== this.expectedSeq) invalidStream();
    this.expectedSeq += 1;

    if (event.name === "meta") {
      if (
        this.stage !== "expect_meta" ||
        event.knowledge_base_id !== this.expectedKnowledgeBaseId
      ) {
        invalidStream();
      }
      this.stage = "expect_retrieval";
      return;
    }
    if (event.name === "retrieval") {
      if (this.stage !== "expect_retrieval") invalidStream();
      this.retrievalEvidence = event.evidence;
      this.stage = "expect_delta_or_error";
      return;
    }
    if (event.name === "delta") {
      if (this.stage !== "expect_delta_or_error" && this.stage !== "accepting_delta")
        invalidStream();
      this.generatedCharacters += event.text.length;
      if (this.generatedCharacters > 64_000) invalidStream();
      this.generatedText += event.text;
      this.stage = "accepting_delta";
      return;
    }
    if (event.name === "citations") {
      if (this.stage !== "accepting_delta") invalidStream();
      const evidenceByLabel = new Map(this.retrievalEvidence.map((item) => [item.label, item]));
      const labels = new Set<string>();
      for (const citation of event.citations) {
        const retrieved = evidenceByLabel.get(citation.label);
        if (
          retrieved === undefined ||
          labels.has(citation.label) ||
          !evidenceMatches(citation, retrieved)
        ) {
          invalidStream();
        }
        labels.add(citation.label);
      }
      const referencedLabels = [
        ...new Set(
          [...this.generatedText.matchAll(/\[S([1-9]\d*)\]/g)].map((match) => `S${match[1]}`),
        ),
      ];
      if (
        referencedLabels.length !== event.citations.length ||
        referencedLabels.some((label, index) => event.citations[index]?.label !== label)
      ) {
        invalidStream();
      }
      this.citations = event.citations;
      this.stage = "expect_usage";
      return;
    }
    if (event.name === "usage") {
      if (this.stage !== "expect_usage") invalidStream();
      this.stage = "expect_done";
      return;
    }
    if (event.name === "done") {
      if (this.stage !== "expect_done") invalidStream();
      if (event.answer_mode === "grounded" && this.citations.length === 0) invalidStream();
      if (event.answer_mode === "refused" && this.citations.length !== 0) invalidStream();
      this.stage = "terminal";
      return;
    }
    if (event.name === "error") {
      if (this.stage !== "expect_delta_or_error") invalidStream();
      this.stage = "terminal";
    }
  }

  finish(): void {
    if (this.stage !== "terminal") invalidStream();
  }
}

function parseSseBlock(block: string): ChatStreamEvent {
  let eventName: string | null = null;
  const dataLines: string[] = [];
  for (const rawLine of block.split("\n")) {
    const line = rawLine.endsWith("\r") ? rawLine.slice(0, -1) : rawLine;
    if (line.length === 0 || line.startsWith(":")) continue;
    const separator = line.indexOf(":");
    const field = separator === -1 ? line : line.slice(0, separator);
    const rawValue = separator === -1 ? "" : line.slice(separator + 1);
    const value = rawValue.startsWith(" ") ? rawValue.slice(1) : rawValue;
    if (field === "event") {
      if (eventName !== null || value.length === 0) invalidStream();
      eventName = value;
    } else if (field === "data") {
      dataLines.push(value);
    } else {
      invalidStream();
    }
  }
  if (eventName === null || dataLines.length === 0) invalidStream();
  return parseEvent(eventName, dataLines.join("\n"));
}

function nextBlock(buffer: string): { block: string; rest: string } | null {
  const match = /\r?\n\r?\n/.exec(buffer);
  if (match?.index === undefined) return null;
  return {
    block: buffer.slice(0, match.index),
    rest: buffer.slice(match.index + match[0].length),
  };
}

export async function consumeChatSse(
  response: Response,
  expectedKnowledgeBaseId: string,
  onEvent: (event: ChatStreamEvent) => void,
): Promise<void> {
  if (!response.headers.get("content-type")?.toLowerCase().startsWith("text/event-stream")) {
    invalidStream();
  }
  if (response.body === null) invalidStream();

  const reader = response.body.getReader();
  const decoder = new TextDecoder();
  const sequence = new ChatEventSequence(expectedKnowledgeBaseId);
  let buffer = "";
  try {
    while (true) {
      const { done, value } = await reader.read();
      if (done) break;
      buffer += decoder.decode(value, { stream: true });
      let extracted = nextBlock(buffer);
      while (extracted !== null) {
        buffer = extracted.rest;
        if (extracted.block.trim().length > 0) {
          const event = parseSseBlock(extracted.block);
          sequence.accept(event);
          onEvent(event);
        }
        extracted = nextBlock(buffer);
      }
    }
    buffer += decoder.decode();
    if (buffer.trim().length > 0) {
      const event = parseSseBlock(buffer);
      sequence.accept(event);
      onEvent(event);
    }
    sequence.finish();
  } finally {
    reader.releaseLock();
  }
}