import "dotenv/config";
import cors from "cors";
import express from "express";
import path from "node:path";
import { fileURLToPath } from "node:url";
import { z } from "zod";
import { findModel, providerForModel, registry } from "./modelRegistry.js";
import { saveModelConfig } from "./modelConfigStore.js";
import { callProvider, streamProvider } from "./providerClients.js";
import { renderPrompt } from "./template.js";
import { normalizeUsage } from "./tokens.js";

const app = express();
const port = Number(process.env.PORT ?? 8787);

app.use(cors());
app.use(express.json({ limit: "1mb" }));

const runSchema = z.object({
  template: z.string().trim().optional(),
  systemTemplate: z.string().trim().optional(),
  userTemplate: z.string().trim().optional(),
  modelIds: z.array(z.string()).min(1),
  draws: z.number().int().min(1).max(20),
  temperature: z.number().min(0).max(2).default(0.8),
  maxTokens: z.number().int().min(64).max(8192).default(1024)
}).superRefine((data, context) => {
  const userTemplate = data.userTemplate || data.template || "";
  if (!userTemplate.trim()) {
    context.addIssue({
      code: z.ZodIssueCode.custom,
      path: ["userTemplate"],
      message: "User prompt template is required"
    });
  }
});

const composePrompt = (systemPrompt: string, userPrompt: string) =>
  [systemPrompt, userPrompt].filter(Boolean).join("\n\n");

const modelConfigSchema = z.object({
  baseUrl: z.string().trim().optional(),
  model: z.string().trim().optional(),
  apiKey: z.string().optional(),
  enabled: z.boolean().optional(),
  clearApiKey: z.boolean().optional()
});

app.get("/api/health", (_request, response) => {
  response.json({ ok: true });
});

app.get("/api/models", (_request, response) => {
  response.json({ models: registry() });
});

app.get("/api/model-configs", (_request, response) => {
  response.json({ models: registry() });
});

app.put("/api/model-configs/:modelId", (request, response) => {
  const model = findModel(request.params.modelId);
  if (!model) {
    response.status(404).json({ error: "Model not found" });
    return;
  }

  const parsed = modelConfigSchema.safeParse(request.body);
  if (!parsed.success) {
    response.status(400).json({ error: parsed.error.flatten() });
    return;
  }

  saveModelConfig(model.id, parsed.data);
  const updated = findModel(model.id);
  response.json({ model: updated });
});

app.post("/api/runs", async (request, response) => {
  const parsed = runSchema.safeParse(request.body);
  if (!parsed.success) {
    response.status(400).json({ error: parsed.error.flatten() });
    return;
  }

  const {
    systemTemplate = "",
    userTemplate: explicitUserTemplate,
    template: legacyTemplate,
    modelIds,
    draws,
    temperature,
    maxTokens
  } = parsed.data;
  const userTemplate = explicitUserTemplate || legacyTemplate || "";
  const startedAt = new Date();
  const uniqueModelIds = [...new Set(modelIds)];
  const models = uniqueModelIds.map((modelId) => findModel(modelId));
  const missing = uniqueModelIds.filter((modelId, index) => !models[index]);

  if (missing.length) {
    response.status(400).json({ error: `Unknown model: ${missing.join(", ")}` });
    return;
  }

  const jobs = models.flatMap((model) =>
    Array.from({ length: draws }, (_, index) => ({
      model: model!,
      draw: index + 1
    }))
  );

  const results = await Promise.all(
    jobs.map(async ({ model, draw }) => {
      const context = { draw, model, startedAt };
      const systemPrompt = systemTemplate ? renderPrompt(systemTemplate, context) : "";
      const userPrompt = renderPrompt(userTemplate, context);
      const prompt = composePrompt(systemPrompt, userPrompt);
      const provider = providerForModel(model);
      const jobStarted = Date.now();

      try {
        const result = await callProvider({
          model,
          provider,
          systemPrompt,
          prompt: userPrompt,
          temperature,
          maxTokens
        });

        return {
          id: `${model.id}-${draw}-${jobStarted}`,
          ok: true,
          draw,
          prompt,
          systemPrompt,
          userPrompt,
          output: result.output,
          modelId: model.id,
          modelLabel: model.label,
          providerName: model.providerName,
          rawModel: result.rawModel,
          durationMs: Date.now() - jobStarted,
          usage: result.usage
        };
      } catch (error) {
        const message = error instanceof Error ? error.message : "Unknown provider error";

        return {
          id: `${model.id}-${draw}-${jobStarted}`,
          ok: false,
          draw,
          prompt,
          systemPrompt,
          userPrompt,
          output: "",
          error: message,
          modelId: model.id,
          modelLabel: model.label,
          providerName: model.providerName,
          rawModel: model.model,
          durationMs: Date.now() - jobStarted,
          usage: normalizeUsage(prompt, "", null)
        };
      }
    })
  );

  response.json({
    startedAt: startedAt.toISOString(),
    finishedAt: new Date().toISOString(),
    results
  });
});

app.post("/api/runs/stream", async (request, response) => {
  const parsed = runSchema.safeParse(request.body);
  if (!parsed.success) {
    response.status(400).json({ error: parsed.error.flatten() });
    return;
  }

  const {
    systemTemplate = "",
    userTemplate: explicitUserTemplate,
    template: legacyTemplate,
    modelIds,
    draws,
    temperature,
    maxTokens
  } = parsed.data;
  const userTemplate = explicitUserTemplate || legacyTemplate || "";
  const startedAt = new Date();
  const uniqueModelIds = [...new Set(modelIds)];
  const models = uniqueModelIds.map((modelId) => findModel(modelId));
  const missing = uniqueModelIds.filter((modelId, index) => !models[index]);

  if (missing.length) {
    response.status(400).json({ error: `Unknown model: ${missing.join(", ")}` });
    return;
  }

  response.setHeader("Content-Type", "application/x-ndjson; charset=utf-8");
  response.setHeader("Cache-Control", "no-cache, no-transform");
  response.setHeader("Connection", "keep-alive");
  response.setHeader("X-Accel-Buffering", "no");
  response.flushHeaders();

  const controller = new AbortController();
  response.on("close", () => {
    if (!response.writableEnded) {
      controller.abort();
    }
  });

  const writeEvent = (event: Record<string, unknown>) => {
    if (!response.destroyed && !response.writableEnded) {
      response.write(`${JSON.stringify(event)}\n`);
    }
  };

  const jobs = models.flatMap((model) =>
    Array.from({ length: draws }, (_, index) => ({
      model: model!,
      draw: index + 1
    }))
  );

  writeEvent({
    type: "run_start",
    startedAt: startedAt.toISOString(),
    totalJobs: jobs.length
  });

  await Promise.all(
    jobs.map(async ({ model, draw }, jobIndex) => {
      const context = { draw, model, startedAt };
      const systemPrompt = systemTemplate ? renderPrompt(systemTemplate, context) : "";
      const userPrompt = renderPrompt(userTemplate, context);
      const prompt = composePrompt(systemPrompt, userPrompt);
      const provider = providerForModel(model);
      const jobStarted = Date.now();
      const id = `${model.id}-${draw}-${startedAt.getTime()}-${jobIndex}`;

      writeEvent({
        type: "card_start",
        id,
        draw,
        prompt,
        systemPrompt,
        userPrompt,
        modelId: model.id,
        modelLabel: model.label,
        providerName: model.providerName,
        rawModel: model.model
      });

      try {
        const result = await streamProvider({
          model,
          provider,
          systemPrompt,
          prompt: userPrompt,
          temperature,
          maxTokens,
          signal: controller.signal,
          onDelta: (delta) => {
            writeEvent({
              type: "delta",
              id,
              channel: delta.channel,
              text: delta.text
            });
          }
        });

        writeEvent({
          type: "card_end",
          id,
          ok: true,
          output: result.output,
          reasoning: result.reasoning ?? "",
          rawModel: result.rawModel,
          durationMs: Date.now() - jobStarted,
          usage: result.usage
        });
      } catch (error) {
        const message = error instanceof Error ? error.message : "Unknown provider error";

        writeEvent({
          type: "card_end",
          id,
          ok: false,
          output: "",
          reasoning: "",
          error: message,
          rawModel: model.model,
          durationMs: Date.now() - jobStarted,
          usage: normalizeUsage(prompt, "", null)
        });
      }
    })
  );

  writeEvent({
    type: "run_end",
    startedAt: startedAt.toISOString(),
    finishedAt: new Date().toISOString()
  });
  response.end();
});

const dirname = path.dirname(fileURLToPath(import.meta.url));
const distPath = path.resolve(dirname, "../dist");

app.use(express.static(distPath));
app.use((request, response, next) => {
  if (request.path.startsWith("/api")) {
    next();
    return;
  }

  const indexPath = path.join(distPath, "index.html");
  response.sendFile(indexPath, (error) => {
    if (error) {
      next();
    }
  });
});

app.listen(port, () => {
  console.log(`Face to API server listening on http://127.0.0.1:${port}`);
});