import type { ConsultationRuntimeState } from "../mastra/consultation-tools.ts";
import {
  agentExecutionReceiptSchema,
  consultationAgentPublicEventSchema,
  publicActivityPhaseSchema,
  type AgentExecutionReceipt,
  type ConsultationAgentPublicEvent,
} from "./consultation-agent-events.ts";
import { createVisibleTextTransformer } from "./stream-text-response.ts";

type Chunk = { type?: string; payload?: Record<string, unknown>; data?: unknown };
type ChunkStream = AsyncIterable<unknown> | ReadableStream<unknown>;

async function* readChunks(stream: ChunkStream): AsyncIterable<Chunk> {
  const values = Symbol.asyncIterator in stream
    ? stream as AsyncIterable<unknown>
    : (async function* () {
      const reader = (stream as ReadableStream<unknown>).getReader();
      try {
        while (true) {
          const { done, value } = await reader.read();
          if (done) return;
          yield value;
        }
      } finally {
        reader.releaseLock();
      }
    })();
  for await (const value of values) {
    if (value && typeof value === "object") yield value as Chunk;
  }
}
type Status = "ready" | "degraded" | "blocked";

export const ENSURE_FINAL_RESPONSE_FALLBACK =
  "本次计算已完成，但暂时没有生成可展示的回答。请换一个角度提问，我会基于已完成的计算继续说明。";

export function ensureFinalResponseText(output: string, contractIsReady: boolean) {
  if (!contractIsReady || /\S/.test(output)) return null;
  return ENSURE_FINAL_RESPONSE_FALLBACK;
}

type EventOptions = {
  runId: string;
  requestId: string;
  toolStatus: () => Status;
  receipt: () => AgentExecutionReceipt;
};

function activity(value: unknown): ConsultationAgentPublicEvent | null {
  if (!value || typeof value !== "object") return null;
  const data = value as { phase?: unknown; label?: unknown };
  const phase = publicActivityPhaseSchema.safeParse(data.phase);
  if (!phase.success || typeof data.label !== "string") return null;
  return { type: "activity", phase: phase.data, label: data.label.slice(0, 120) };
}

function safeToolError(error: unknown) {
  if (error instanceof DOMException && error.name === "AbortError") return "cancelled" as const;
  if (error instanceof DOMException && error.name === "TimeoutError") return "timeout" as const;
  return "calculation_failed" as const;
}

function mapChunk(
  chunk: Chunk,
  options: EventOptions,
  startedAt: Map<string, number>,
  jyotishSkillCallIds: Set<string>,
): ConsultationAgentPublicEvent[] {
  const payload = chunk.payload ?? {};
  if (chunk.type === "data-jyotish-activity") {
    const event = activity(chunk.data);
    return event ? [event] : [];
  }
  if (chunk.type === "tool-call") {
    const toolName = payload.toolName;
    const callId = typeof payload.toolCallId === "string" ? payload.toolCallId : "tool";
    if (toolName === "skill" && (payload.args as { name?: unknown } | undefined)?.name === "jyotish-vedic-astrology") {
      jyotishSkillCallIds.add(callId);
      return [{ type: "skill.started", name: "jyotish-vedic-astrology" }];
    }
    if (toolName === "run-jyotish-consultation") {
      startedAt.set(callId, Date.now());
      return [{ type: "tool.started", callId, tool: "run-jyotish-consultation", label: "正在计算个人星盘" }];
    }
  }
  if (chunk.type === "tool-result") {
    const toolName = payload.toolName;
    const callId = typeof payload.toolCallId === "string" ? payload.toolCallId : "tool";
    if (toolName === "skill" && jyotishSkillCallIds.delete(callId)) {
      return [{ type: "skill.completed", name: "jyotish-vedic-astrology" }];
    }
    if (toolName === "run-jyotish-consultation") {
      return [{
        type: "tool.completed", callId, tool: "run-jyotish-consultation", status: options.toolStatus(),
        durationMs: Math.max(0, Date.now() - (startedAt.get(callId) ?? Date.now())),
      }];
    }
  }
  if (chunk.type === "tool-error") {
    const toolName = payload.toolName;
    if (toolName === "run-jyotish-consultation") {
      return [{
        type: "tool.failed",
        callId: typeof payload.toolCallId === "string" ? payload.toolCallId : "tool",
        tool: "run-jyotish-consultation",
        code: safeToolError(payload.error),
      }];
    }
  }
  return [];
}

export async function collectAgentPublicEvents(stream: ChunkStream | Iterable<Chunk>, options: EventOptions) {
  const events: ConsultationAgentPublicEvent[] = [{ type: "run.started", runId: options.runId, requestId: options.requestId }];
  const startedAt = new Map<string, number>();
  const jyotishSkillCallIds = new Set<string>();
  for await (const chunk of stream instanceof ReadableStream || Symbol.asyncIterator in stream ? readChunks(stream as ChunkStream) : stream) {
    events.push(...mapChunk(chunk, options, startedAt, jyotishSkillCallIds));
    if (chunk.type === "text-delta" && typeof chunk.payload?.text === "string") {
      events.push({ type: "answer.delta", text: chunk.payload.text });
    }
  }
  events.push({ type: "run.completed", receipt: agentExecutionReceiptSchema.parse(options.receipt()) });
  return events.map((event) => consultationAgentPublicEventSchema.parse(event));
}

type StreamAgentResponseOptions = EventOptions & {
  state: ConsultationRuntimeState;
  stream: ChunkStream;
  transformText?: (text: string) => string;
  requireTool: boolean;
  retry?: () => Promise<ChunkStream>;
  continueAfterDisconnect?: boolean;
  headers?: HeadersInit;
  onFirstActivity?: () => void | Promise<void>;
  onFirstOutput?: () => void | Promise<void>;
  onComplete?: (output: string, receipt: AgentExecutionReceipt) => void | Promise<void>;
  onError?: (error: unknown, emitted: boolean, output: string) => void | Promise<void>;
  onCancel?: (emitted: boolean) => void | Promise<void>;
};

function contractReady(options: StreamAgentResponseOptions) {
  return options.state.jyotishSkillLoaded
    && (!options.requireTool || (options.state.consultationToolCompleted && options.state.consultationToolCallCount === 1));
}

export function streamAgentResponse(options: StreamAgentResponseOptions) {
  const encoder = new TextEncoder();
  let disconnected = false;
  let settled = false;
  let settling = false;
  let emitted = false;
  let firstActivity = false;
  let firstOutput = false;
  let fullOutput = "";
  const startedAt = new Map<string, number>();
  const jyotishSkillCallIds = new Set<string>();
  const send = (controller: ReadableStreamDefaultController<Uint8Array> | undefined, event: ConsultationAgentPublicEvent) => {
    if (!firstActivity && (event.type === "skill.started" || event.type === "tool.started" || event.type === "activity")) {
      firstActivity = true;
      void Promise.resolve(options.onFirstActivity?.()).catch(() => {});
    }
    if (!disconnected && controller) controller.enqueue(encoder.encode(`${JSON.stringify(consultationAgentPublicEventSchema.parse(event))}\n`));
  };

  async function consumeAttempt(controller: ReadableStreamDefaultController<Uint8Array> | undefined, stream: ChunkStream) {
    const visible = createVisibleTextTransformer(options.transformText ?? ((value) => value));
    let held = "";
    let attemptOutput = "";
    let composingSent = false;
    const outputText = async (text: string) => {
      attemptOutput += text;
      held += text;
      if (!held || !contractReady(options)) return;
      if (!composingSent) {
        composingSent = true;
        send(controller, { type: "activity", phase: "answer-composition", label: "正在组织回答" });
      }
      if (!firstOutput && /\S/.test(held)) {
        firstOutput = true;
        await options.onFirstOutput?.();
      }
      send(controller, { type: "answer.delta", text: held });
      fullOutput += held;
      if (/\S/.test(held)) emitted = true;
      held = "";
    };
    for await (const chunk of readChunks(stream)) {
      for (const event of mapChunk(chunk, options, startedAt, jyotishSkillCallIds)) send(controller, event);
      if (chunk.type === "text-delta" && typeof chunk.payload?.text === "string") {
        await outputText(visible.push(chunk.payload.text));
      }
    }
    await outputText(visible.finish(""));
    return { held, attemptOutput };
  }

  const body = new ReadableStream<Uint8Array>({
    start(controller) {
      void (async () => {
        send(controller, { type: "run.started", runId: options.runId, requestId: options.requestId });
        try {
          const first = await consumeAttempt(controller, options.stream);
          if (!contractReady(options) && options.retry) {
            if (options.state.steps.length < 32) {
              options.state.steps.push({ sequence: options.state.steps.length + 1, kind: "validation", name: "runtime-contract-retry", status: "completed" });
            }
            send(controller, { type: "activity", phase: "loading-method", label: "正在补齐方法与计算步骤" });
            await consumeAttempt(controller, await options.retry());
          }
          if (!contractReady(options)) throw new Error("runtime_contract_incomplete");
          const ensuredFinalResponse = ensureFinalResponseText(fullOutput, contractReady(options));
          if (ensuredFinalResponse) {
            if (options.state.steps.length < 32) {
              options.state.steps.push({ sequence: options.state.steps.length + 1, kind: "validation", name: "ensure-final-response", status: "completed" });
            }
            if (!firstOutput) {
              firstOutput = true;
              await options.onFirstOutput?.();
            }
            send(controller, { type: "answer.delta", text: ensuredFinalResponse });
            fullOutput = ensuredFinalResponse;
            emitted = true;
          }
          if (!/\S/.test(fullOutput)) {
            if (/\S/.test(first.held) || /\S/.test(first.attemptOutput)) throw new Error("runtime_contract_incomplete");
            throw new Error("empty_answer");
          }
          settling = true;
          const receipt = agentExecutionReceiptSchema.parse(options.receipt());
          await options.onComplete?.(fullOutput, receipt);
          settled = true;
          settling = false;
          send(controller, { type: "run.completed", receipt });
          if (!disconnected) controller.close();
        } catch (error) {
          if (settled) return;
          settled = true;
          settling = false;
          try {
            await options.onError?.(error, emitted, fullOutput);
          } catch {}
          const code = error instanceof Error && error.message === "runtime_contract_incomplete"
            ? "runtime_contract_incomplete" as const
            : error instanceof Error && error.message === "empty_answer"
              ? "empty_answer" as const
              : "calculation_failed" as const;
          send(controller, { type: "run.failed", code, message: code === "runtime_contract_incomplete" ? "Agent 未完成必要的方法与计算步骤，本次不会扣点。" : "咨询暂时无法完成，本次不会扣点。" });
          if (!disconnected) controller.close();
        }
      })();
    },
    async cancel() {
      if (settled) return;
      if (settling || options.continueAfterDisconnect) {
        disconnected = true;
        return;
      }
      settled = true;
      await options.onCancel?.(emitted);
    },
  });
  return new Response(body, {
    headers: {
      "cache-control": "no-cache, no-transform",
      "content-type": "application/x-ndjson; charset=utf-8",
      "x-accel-buffering": "no",
      "x-ayanam-mode": "mastra-agentic",
      "x-ayanam-request-id": options.requestId,
      ...options.headers,
    },
  });
}