import { createFileRoute } from "@tanstack/react-router";
import { z } from "zod";
import { getLogById } from "../../proxy/store";
import { findProviderByModel } from "../../proxy/providers";
import { registry } from "../../proxy/formats";
import { formatForPath } from "../../proxy/formats";
import { extractModelFromBody } from "../../proxy/schemas";
import {
  PROXY_IDENTITY,
  PRESERVE_HEADERS,
  HEADER_CONTENT_TYPE,
  HEADER_USER_AGENT,
  HEADER_X_PROXY_IDENTITY,
  CONTENT_TYPE_EVENT_STREAM,
  STATUS_FORBIDDEN,
} from "../../proxy/constants";
import {
  buildUpstreamUrl,
  describeApiRoute,
  injectProviderAuth,
  selectUpstreamBase,
  setUpstreamHost,
} from "../../proxy/upstream";

type ReplayRequest = {
  modifiedBody: string;
};

const ReplayRequestSchema = z.object({
  modifiedBody: z.string(),
});

type ReplayResponse = {
  success: boolean;
  error?: string;
  responseStatus?: number;
  responseText?: string;
  inputTokens?: number;
  outputTokens?: number;
  elapsedMs?: number;
  streaming?: boolean;
};

export const Route = createFileRoute("/api/logs/$id/replay")({
  server: {
    handlers: {
      POST: async ({ params, request }: { params: { id: string }; request: Request }) => {
        const id = Number(params.id);
        if (isNaN(id)) {
          return Response.json({ error: "Invalid log ID" }, { status: 400 });
        }

        let replayRequest: ReplayRequest;
        try {
          const raw: unknown = await request.json();
          const parsed = ReplayRequestSchema.safeParse(raw);
          if (!parsed.success) {
            return Response.json({ error: "Invalid request body" }, { status: 400 });
          }
          replayRequest = parsed.data;
        } catch {
          return Response.json({ error: "Invalid JSON body" }, { status: 400 });
        }

        const { modifiedBody } = replayRequest;

        // Get original log for context
        const log = await getLogById(id);
        if (log === null) {
          return Response.json({ error: "Log not found" }, { status: 404 });
        }

        // Extract model from modified body
        const model = extractModelFromBody(modifiedBody);
        if (model === null) {
          return Response.json(
            { error: "Could not extract model from request body" },
            { status: 400 },
          );
        }

        // Find provider
        const matchedProviderConfig = findProviderByModel(model);
        const provider = registry.findProvider(model);
        if (provider === null) {
          return Response.json(
            { error: "No provider found for model" },
            { status: STATUS_FORBIDDEN },
          );
        }

        // Determine API path from original log
        const route = describeApiRoute(log.path);

        // Get format handler
        const formatHandler = formatForPath(route.apiPath);
        if (formatHandler === null) {
          return Response.json({ error: "Unsupported path" }, { status: STATUS_FORBIDDEN });
        }

        // Build upstream URL
        const upstreamBase = selectUpstreamBase(route, matchedProviderConfig);
        const upstreamUrl = buildUpstreamUrl(upstreamBase, route.normalizedPath);

        // Build headers
        const headers = new Headers();
        headers.set(HEADER_USER_AGENT, PROXY_IDENTITY);
        headers.set(HEADER_X_PROXY_IDENTITY, PROXY_IDENTITY);
        headers.set(HEADER_CONTENT_TYPE, "application/json");

        for (const name of PRESERVE_HEADERS) {
          const value = log.rawHeaders?.[name.toLowerCase()];
          if (value !== undefined) {
            headers.set(name, value);
          }
        }
        setUpstreamHost(headers, upstreamBase);
        injectProviderAuth(headers, matchedProviderConfig);

        const startTime = Date.now();

        // Send request
        let upstreamRes: Response;
        try {
          upstreamRes = await fetch(upstreamUrl, {
            method: "POST",
            headers,
            body: modifiedBody,
          });
        } catch (err) {
          return Response.json({
            success: false,
            error: String(err),
          } satisfies ReplayResponse);
        }

        const elapsedMs = Date.now() - startTime;
        const isStream =
          upstreamRes.headers.get(HEADER_CONTENT_TYPE)?.includes(CONTENT_TYPE_EVENT_STREAM) ??
          false;

        if (isStream) {
          // Accumulate streaming response
          const chunks: string[] = [];
          const decoder = new TextDecoder();
          const reader = upstreamRes.body?.getReader();
          if (reader) {
            try {
              while (true) {
                const { done, value } = await reader.read();
                if (done) break;
                chunks.push(decoder.decode(value, { stream: true }));
              }
            } catch {
              // Stream read error
            }
          }
          const fullResponse = chunks.join("");

          // Create a mock log for streaming extraction
          const mockLog: typeof log = { ...log };
          const responseText = formatHandler.extractStream(fullResponse, mockLog, model, false);
          const tokens = formatHandler.extractTokens(responseText);

          return Response.json({
            success: true,
            responseStatus: upstreamRes.status,
            responseText,
            inputTokens: tokens.inputTokens ?? undefined,
            outputTokens: tokens.outputTokens ?? undefined,
            elapsedMs,
            streaming: true,
          } satisfies ReplayResponse);
        } else {
          const responseText = await upstreamRes.text();
          const tokens = formatHandler.extractTokens(responseText);

          return Response.json({
            success: true,
            responseStatus: upstreamRes.status,
            responseText,
            inputTokens: tokens.inputTokens ?? undefined,
            outputTokens: tokens.outputTokens ?? undefined,
            elapsedMs,
            streaming: false,
          } satisfies ReplayResponse);
        }
      },
    },
  },
});
