import { createLogger } from "@/lib/observability/logger";
import { JsonToSseTransformStream } from "ai";
import { extractImageUrlsFromText } from "@/lib/image-url";
import type {
  ResolvedChatMediaInput,
  ResolvedImageRef,
  ResolvedTargetAssets,
} from "@/lib/ai/context-memory/types";
import type { CanvasViewOperationPlan } from "@/types/canvas/agent-protocol";
import type { CanvasCapabilityClientCommand } from "@/lib/canvas/capabilities/client-command";
import { generateUUID } from "@/lib/utils";

const chatApiLogger = createLogger("chat.stream");


export type StreamWriter = {
  write: (chunk: any) => void;
};

export function safeWriteStreamChunk(
  writer: StreamWriter,
  chunk: any,
  label: string
): boolean {
  try {
    writer.write(chunk);
    return true;
  } catch (error) {
    chatApiLogger.warn(`[Chat API] Stream write skipped (${label})`, error);
    return false;
  }
}

export function writeAssistantTextChunk(
  writer: StreamWriter,
  text: string,
  label: string,
  messageId = generateUUID()
) {
  safeWriteStreamChunk(
    writer,
    {
      type: "data-appendMessage",
      data: JSON.stringify({
        id: messageId,
        role: "assistant",
        parts: [{ type: "text", text }],
      }),
    },
    label
  );
  return messageId;
}

export function writeStructuredDataChunk(
  writer: StreamWriter,
  type: string,
  data: unknown,
  label: string
) {
  safeWriteStreamChunk(
    writer,
    {
      type,
      data: JSON.stringify(data),
    },
    label
  );
}

function normalizeStreamIdList(value: unknown) {
  return Array.isArray(value)
    ? value.map((id) => String(id || "").trim()).filter(Boolean)
    : [];
}

function normalizeStreamMediaUrls(value: unknown) {
  return Array.isArray(value)
    ? value.map((url) => String(url || "").trim())
    : [];
}

function isValidStreamMediaUrl(url: string) {
  return /^https?:\/\//i.test(url);
}

export function normalizeCanvasMediaStreamPayload(params: {
  imageUrls?: unknown;
  videoUrls?: unknown;
  mediaItemIds?: unknown;
  handledCanvasMediaItemIds?: unknown;
}) {
  const rawImageUrls = normalizeStreamMediaUrls(params.imageUrls);
  const rawVideoUrls = normalizeStreamMediaUrls(params.videoUrls);
  const rawEntries = [
    ...rawImageUrls.map((url, index) => ({
      type: "image" as const,
      url,
      originalIndex: index,
    })),
    ...rawVideoUrls.map((url, index) => ({
      type: "video" as const,
      url,
      originalIndex: rawImageUrls.length + index,
    })),
  ];
  const rawMediaItemIds = normalizeStreamIdList(params.mediaItemIds);
  const validEntries = rawEntries.filter((entry) =>
    isValidStreamMediaUrl(entry.url)
  );

  const idByValidIndex = new Map<number, string>();
  if (rawMediaItemIds.length === rawEntries.length) {
    validEntries.forEach((entry, validIndex) => {
      const id = rawMediaItemIds[entry.originalIndex];
      if (id) idByValidIndex.set(validIndex, id);
    });
  } else if (rawMediaItemIds.length === validEntries.length) {
    rawMediaItemIds.forEach((id, index) => {
      if (id) idByValidIndex.set(index, id);
    });
  }

  const dedupedEntries: Array<{
    type: "image" | "video";
    url: string;
    mediaItemId?: string;
  }> = [];
  const seenMediaKeys = new Set<string>();
  validEntries.forEach((entry, validIndex) => {
    const key = `${entry.type}:${entry.url}`;
    if (seenMediaKeys.has(key)) return;
    seenMediaKeys.add(key);
    dedupedEntries.push({
      type: entry.type,
      url: entry.url,
      mediaItemId: idByValidIndex.get(validIndex),
    });
  });

  if (
    dedupedEntries.length > 0 &&
    dedupedEntries.some((entry) => !entry.mediaItemId) &&
    rawMediaItemIds.length === dedupedEntries.length
  ) {
    dedupedEntries.forEach((entry, index) => {
      entry.mediaItemId = rawMediaItemIds[index];
    });
  }

  const mediaItemIds = dedupedEntries.every((entry) => entry.mediaItemId)
    ? dedupedEntries.map((entry) => entry.mediaItemId!)
    : [];

  return {
    imageUrls: dedupedEntries
      .filter((entry) => entry.type === "image")
      .map((entry) => entry.url),
    videoUrls: dedupedEntries
      .filter((entry) => entry.type === "video")
      .map((entry) => entry.url),
    mediaItemIds,
    handledCanvasMediaItemIds: normalizeStreamIdList(
      params.handledCanvasMediaItemIds
    ),
  };
}

export function writeProtocolDebugChunk(
  writer: StreamWriter,
  enabled: boolean,
  data: unknown,
  label: string
) {
  if (!enabled) return;
  writeStructuredDataChunk(writer, "data-protocolDebug", data, label);
}

function normalizeImageLoadingProviderType(providerType?: string | null): string {
  const normalized = String(providerType || "").trim().toLowerCase();
  if (normalized === "gemini-image" || normalized === "huanke-gemini") {
    return "gemini";
  }
  if (
    normalized === "alibaba-image" ||
    normalized === "bfl-image" ||
    normalized === "ruxa-image" ||
    normalized === "fal-image" ||
    normalized === "snapgen-image" ||
    normalized === "openai-image" ||
    normalized === "bananaimage" ||
    normalized === "stepfun-image" ||
    normalized === "midjourney" ||
    normalized === "runway" ||
    normalized === "krea" ||
    normalized === "volcano-vision" ||
    normalized === "gemini"
  ) {
    return normalized;
  }
  return "alibaba-image";
}

export function writeImageGenerationStatusChunk(
  writer: StreamWriter,
  params: {
    phase: "queued" | "processing" | "finalizing" | "failed" | "completed";
    providerType?: string | null;
    modelName?: string | null;
    imageCount?: number | null;
    queuePosition?: number | null;
    sameModelRunningCount?: number | null;
    concurrencyLimit?: number | null;
    estimatedWaitMs?: number | null;
  },
  label: string
) {
  writeStructuredDataChunk(
    writer,
    "data-imageGenerationStatus",
    {
      phase: params.phase,
      providerType: normalizeImageLoadingProviderType(params.providerType),
      modelName: String(params.modelName || "").trim() || undefined,
      imageCount:
        Number.isFinite(Number(params.imageCount)) && Number(params.imageCount) > 0
          ? Number(params.imageCount)
          : undefined,
      queuePosition:
        Number.isFinite(Number(params.queuePosition)) &&
        Number(params.queuePosition) >= 0
          ? Number(params.queuePosition)
          : undefined,
      sameModelRunningCount:
        Number.isFinite(Number(params.sameModelRunningCount)) &&
        Number(params.sameModelRunningCount) >= 0
          ? Number(params.sameModelRunningCount)
          : undefined,
      concurrencyLimit:
        Number.isFinite(Number(params.concurrencyLimit)) &&
        Number(params.concurrencyLimit) > 0
          ? Number(params.concurrencyLimit)
          : undefined,
      estimatedWaitMs:
        Number.isFinite(Number(params.estimatedWaitMs)) &&
        Number(params.estimatedWaitMs) >= 0
          ? Number(params.estimatedWaitMs)
          : undefined,
    },
    label
  );
}

export type AgentProgressStreamPayload = {
  stage: "planning" | "clarify" | "executing" | "completed" | "failed";
  currentStep:
    | "lock_target"
    | "analyze"
    | "select_tool"
    | "execute"
    | "finalize";
  title?: string;
  message?: string;
  toolName?: string;
  capability?: string;
};

export function writeAgentProgressChunk(
  writer: StreamWriter,
  payload: AgentProgressStreamPayload,
  label: string
) {
  writeStructuredDataChunk(
    writer,
    "data-agentProgress",
    {
      stage: payload.stage,
      currentStep: payload.currentStep,
      title: String(payload.title || "").trim() || undefined,
      message: String(payload.message || "").trim() || undefined,
      toolName: String(payload.toolName || "").trim() || undefined,
      capability: String(payload.capability || "").trim() || undefined,
    },
    label
  );
}

export function writeCanvasMediaChunk(
  writer: StreamWriter,
  params: {
    messageId: string;
    text?: string;
    imageUrls?: string[];
    videoUrls?: string[];
    prompt?: string | null;
    sourceImageUrls?: string[] | null;
    agentPresetId?: string | null;
    agentPresetName?: string | null;
    agentCategoryId?: string | null;
    agentCategoryName?: string | null;
    agentExecutionSnapshotId?: string | null;
    mediaItemIds?: string[] | null;
    handledCanvasMediaItemIds?: string[] | null;
  },
  label: string
) {
  const inputImageUrls = Array.isArray(params.imageUrls)
    ? params.imageUrls
    : extractImageUrlsFromText(String(params.text || ""));
  const normalizedPayload = normalizeCanvasMediaStreamPayload({
    imageUrls: inputImageUrls,
    videoUrls: params.videoUrls,
    mediaItemIds: params.mediaItemIds,
    handledCanvasMediaItemIds: params.handledCanvasMediaItemIds,
  });

  if (
    normalizedPayload.imageUrls.length === 0 &&
    normalizedPayload.videoUrls.length === 0
  ) {
    return;
  }

  writeStructuredDataChunk(
    writer,
    "data-canvasMedia",
    {
      messageId: params.messageId,
      imageUrls: normalizedPayload.imageUrls,
      videoUrls: normalizedPayload.videoUrls,
      prompt: String(params.prompt || "").trim() || undefined,
      sourceImageUrls: Array.isArray(params.sourceImageUrls)
        ? params.sourceImageUrls
            .map((url) => String(url || "").trim())
            .filter((url) => /^https?:\/\//i.test(url))
        : [],
      agentPresetId: String(params.agentPresetId || "").trim() || undefined,
      agentPresetName: String(params.agentPresetName || "").trim() || undefined,
      agentCategoryId: String(params.agentCategoryId || "").trim() || undefined,
      agentCategoryName: String(params.agentCategoryName || "").trim() || undefined,
      agentExecutionSnapshotId:
        String(params.agentExecutionSnapshotId || "").trim() || undefined,
      mediaItemIds: normalizedPayload.mediaItemIds,
      handledCanvasMediaItemIds: normalizedPayload.handledCanvasMediaItemIds,
    },
    label
  );
}

export function writeCanvasOperationChunk(
  writer: StreamWriter,
  params: {
    projectId: string;
    operationId: string;
    operation: string;
    status?: "applied" | "failed" | "shadowed";
    createdNodeIds: string[];
    createdNodeIdsByMediaItemId?: Record<string, string>;
    updatedNodeIds?: string[];
    createdEdgeIds?: string[];
    canvasRevision?: string;
    anchorNodeId?: string | null;
    handledCanvasMediaItemIds?: string[];
    mediaItemIds?: string[];
    errorCode?: string;
    errorMessage?: string;
    fallbackToCanvasMedia?: boolean;
    message?: string;
  },
  label: string
) {
  writeStructuredDataChunk(
    writer,
    "data-canvasOperation",
    {
      projectId: params.projectId,
      operationId: params.operationId,
      operation: params.operation,
      status: params.status || "applied",
      createdNodeIds: params.createdNodeIds || [],
      createdNodeIdsByMediaItemId: params.createdNodeIdsByMediaItemId || {},
      updatedNodeIds: params.updatedNodeIds || [],
      createdEdgeIds: params.createdEdgeIds || [],
      canvasRevision: params.canvasRevision,
      anchorNodeId: params.anchorNodeId || undefined,
      handledCanvasMediaItemIds:
        params.status && params.status !== "applied"
          ? []
          : normalizeStreamIdList(params.handledCanvasMediaItemIds),
      mediaItemIds: normalizeStreamIdList(params.mediaItemIds),
      errorCode: params.errorCode,
      errorMessage: params.errorMessage,
      fallbackToCanvasMedia: params.fallbackToCanvasMedia,
      message: params.message,
    },
    label
  );
}

export function writeCanvasViewOperationChunk(
  writer: StreamWriter,
  params: CanvasViewOperationPlan,
  label: string
) {
  writeStructuredDataChunk(
    writer,
    "data-canvasViewOperation",
    {
      version: 1,
      viewOperationId: params.viewOperationId,
      operation: params.operation,
      projectId: params.projectId,
      nodeIds: params.nodeIds || [],
      anchorNodeId: params.anchorNodeId,
      zoom: params.zoom,
      padding: params.padding,
      durationMs: params.durationMs,
      persist: false,
      center: params.center,
      reason: params.reason,
    },
    label
  );
}

export function writeCanvasCapabilityCommandChunk(
  writer: StreamWriter,
  command: CanvasCapabilityClientCommand,
  label: string
) {
  writeStructuredDataChunk(
    writer,
    "data-canvasCapabilityCommand",
    command,
    label
  );
}

type AgentTargetLockRefPreview = {
  url: string;
  title?: string;
  nodeId?: string;
  assetId?: string;
  source: ResolvedImageRef["source"];
  role: ResolvedImageRef["role"];
  isSelected?: boolean;
  positionLabel?: string;
};

export type AgentTargetLockStreamPayload = {
  stage: "locked" | "clarify" | "unresolved";
  resolutionMode: ResolvedTargetAssets["resolutionMode"];
  confidence: number;
  reason: string;
  needsClarification: boolean;
  clarificationMessage?: string;
  targetCount: number;
  candidateCount: number;
  targets: AgentTargetLockRefPreview[];
  candidates: AgentTargetLockRefPreview[];
};

function sortTargetLockRefsForCanvasOrdering(refs: ResolvedImageRef[]) {
  return [...refs].sort((left, right) => {
    const leftY = Number(left.y ?? 0);
    const rightY = Number(right.y ?? 0);
    if (Math.abs(leftY - rightY) > 12) {
      return leftY - rightY;
    }
    return Number(left.x ?? 0) - Number(right.x ?? 0);
  });
}

function getTargetLockPositionLabel(
  ref: ResolvedImageRef,
  refs: ResolvedImageRef[]
): string {
  if (refs.length <= 1) return "中间";
  const xs = refs.map((item) => Number(item.x ?? 0));
  const ys = refs.map((item) => Number(item.y ?? 0));
  const minX = Math.min(...xs);
  const maxX = Math.max(...xs);
  const minY = Math.min(...ys);
  const maxY = Math.max(...ys);
  const spanX = Math.max(1, maxX - minX);
  const spanY = Math.max(1, maxY - minY);
  const relX = (Number(ref.x ?? 0) - minX) / spanX;
  const relY = (Number(ref.y ?? 0) - minY) / spanY;

  const horizontal = relX <= 0.33 ? "左" : relX >= 0.66 ? "右" : "中";
  const vertical = relY <= 0.33 ? "上" : relY >= 0.66 ? "下" : "中";

  if (horizontal === "中" && vertical === "中") return "中间";
  return `${vertical}${horizontal}`;
}

function serializeAgentTargetLockRefs(refs: ResolvedImageRef[]) {
  const orderedRefs = sortTargetLockRefsForCanvasOrdering(refs).slice(0, 4);
  return orderedRefs.map((ref) => ({
    url: ref.url,
    title: String(ref.title || "").trim() || undefined,
    nodeId: ref.nodeId,
    assetId: ref.assetId,
    source: ref.source,
    role: ref.role,
    isSelected: ref.isSelected === true,
    positionLabel: getTargetLockPositionLabel(ref, orderedRefs),
  })) satisfies AgentTargetLockRefPreview[];
}

export function buildAgentTargetLockPayload(params: {
  requestTargetResolution: ResolvedTargetAssets;
  requestResolvedMediaInput: ResolvedChatMediaInput;
  needsClarification: boolean;
  clarificationMessage?: string;
}): AgentTargetLockStreamPayload {
  const stage = params.needsClarification
    ? "clarify"
    : params.requestTargetResolution.targetRefs.length > 0
      ? "locked"
      : "unresolved";

  return {
    stage,
    resolutionMode: params.requestTargetResolution.resolutionMode,
    confidence: Number(params.requestTargetResolution.confidence || 0),
    reason: params.requestTargetResolution.reason,
    needsClarification: params.needsClarification,
    clarificationMessage: params.clarificationMessage?.trim() || undefined,
    targetCount: params.requestTargetResolution.targetRefs.length,
    candidateCount: params.requestTargetResolution.candidateRefs.length,
    targets: serializeAgentTargetLockRefs(params.requestTargetResolution.targetRefs),
    candidates: serializeAgentTargetLockRefs(
      params.requestTargetResolution.candidateRefs.length > 0
        ? params.requestTargetResolution.candidateRefs
        : params.requestResolvedMediaInput.imageRefs
    ),
  };
}

export function writeAgentTargetLockChunk(
  writer: StreamWriter,
  payload: AgentTargetLockStreamPayload,
  label: string
) {
  writeStructuredDataChunk(writer, "data-agentTargetLock", payload, label);
}

export function createSseResponse(stream: ReadableStream) {
  return new Response(stream.pipeThrough(new JsonToSseTransformStream()), {
    headers: {
      "Content-Type": "text/event-stream; charset=utf-8",
      "Cache-Control": "no-cache, no-transform",
      Connection: "keep-alive",
      "X-Accel-Buffering": "no",
    },
  });
}

export function startStreamKeepalive(
  writer: StreamWriter,
  label: string,
  intervalMs = 15_000
) {
  const timer = setInterval(() => {
    safeWriteStreamChunk(
      writer,
      {
        type: "data-keepalive",
        data: JSON.stringify({ ts: Date.now(), label }),
      },
      `${label}:keepalive`
    );
  }, intervalMs);

  return () => clearInterval(timer);
}

export async function runBoundedParallel<T>(
  items: T[],
  workerCount: number,
  worker: (item: T, index: number) => Promise<void>
) {
  const maxWorkers = Math.max(1, Math.min(workerCount, items.length));
  let cursor = 0;
  let firstError: unknown = null;

  const workers = Array.from({ length: maxWorkers }, async () => {
    while (cursor < items.length) {
      const index = cursor;
      cursor += 1;
      if (firstError) {
        return;
      }

      try {
        await worker(items[index]!, index);
      } catch (error) {
        firstError = firstError || error;
      }
    }
  });

  await Promise.all(workers);

  if (firstError) {
    throw firstError;
  }
}