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 { 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
  );
}

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 === "openai-image" ||
    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;
  },
  label: string
) {
  const imageUrls = Array.isArray(params.imageUrls)
    ? params.imageUrls
    : extractImageUrlsFromText(String(params.text || ""));
  const videoUrls = Array.isArray(params.videoUrls) ? params.videoUrls : [];
  const normalizedImageUrls = Array.from(
    new Set(
      imageUrls
        .map((url) => String(url || "").trim())
        .filter((url) => /^https?:\/\//i.test(url))
    )
  );
  const normalizedVideoUrls = Array.from(
    new Set(
      videoUrls
        .map((url) => String(url || "").trim())
        .filter((url) => /^https?:\/\//i.test(url))
    )
  );

  if (normalizedImageUrls.length === 0 && normalizedVideoUrls.length === 0) {
    return;
  }

  writeStructuredDataChunk(
    writer,
    "data-canvasMedia",
    {
      messageId: params.messageId,
      imageUrls: normalizedImageUrls,
      videoUrls: normalizedVideoUrls,
      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,
    },
    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;
  }
}