import {
  extendZodWithOpenApi,
  OpenAPIRegistry,
  OpenApiGeneratorV3,
} from "@asteasolutions/zod-to-openapi";
import { z } from "zod";

extendZodWithOpenApi(z);

const registry = new OpenAPIRegistry();
registry.registerComponent("securitySchemes", "bearerAuth", {
  type: "http",
  scheme: "bearer",
  bearerFormat: "JWT",
});

const SessionId = z.object({ session_id: z.string().max(128) });
const RenderRef = z.object({ render_ref: z.string().min(1) });
const RenderReferenceQuery = z.object({
  session_id: z.string().max(128).optional(),
});
const SessionCreate = z
  .object({
    session_id: z.string(),
    title: z.string().nullable().optional(),
    created_at: z.string(),
    updated_at: z.string(),
    status: z.string(),
    parent_session_id: z.string().nullable().optional(),
  })
  .openapi("AgentSessionCreate");
const SessionSummary = z
  .object({
    id: z.string(),
    title: z.string(),
    created_at: z.string(),
    updated_at: z.string(),
    status: z.string(),
    parent_session_id: z.string().nullable().optional(),
    is_streaming: z.boolean(),
    run_status: z.string().nullable().optional(),
  })
  .openapi("AgentSessionSummary");
const SessionDetail = SessionSummary.extend({
  session_id: z.string(),
  is_title_manually_edited: z.boolean(),
  messages: z.array(z.unknown()),
}).openapi("AgentSessionDetail");
const SessionUpdate = z
  .object({
    id: z.string(),
    title: z.string(),
    updated_at: z.string(),
  })
  .openapi("AgentSessionUpdate");
const SessionFork = z
  .object({
    session_id: z.string(),
  })
  .openapi("AgentSessionFork");
const Problem = z
  .object({
    type: z.string(),
    title: z.string(),
    status: z.number().int(),
    detail: z.string(),
    instance: z.string(),
    code: z.string(),
    trace_id: z.string(),
    errors: z.array(z.unknown()),
  })
  .openapi("ProblemDetails");
const JsonObject = z.record(z.unknown());
const errorResponses = {
  400: {
    description: "Invalid request",
    content: { "application/problem+json": { schema: Problem } },
  },
  401: {
    description: "Authentication required",
    content: { "application/problem+json": { schema: Problem } },
  },
  403: {
    description: "Insufficient permission",
    content: { "application/problem+json": { schema: Problem } },
  },
  404: {
    description: "Resource not found",
    content: { "application/problem+json": { schema: Problem } },
  },
  409: {
    description: "Resource conflict",
    content: { "application/problem+json": { schema: Problem } },
  },
  422: {
    description: "Validation error",
    content: { "application/problem+json": { schema: Problem } },
  },
  500: {
    description: "Internal server error",
    content: { "application/problem+json": { schema: Problem } },
  },
  502: {
    description: "Upstream dependency error",
    content: { "application/problem+json": { schema: Problem } },
  },
  503: {
    description: "Dependency unavailable",
    content: { "application/problem+json": { schema: Problem } },
  },
};
const jsonResponse = (schema: z.ZodTypeAny, description = "Successful response") => ({
  description,
  content: { "application/json": { schema } },
});
type RegisterConfig = {
  summary: string;
  request?: Record<string, unknown>;
  responses: Record<number, unknown>;
};

export const createRouteRegistrar = (targetRegistry: OpenAPIRegistry) => {
  const operations = new Set<string>();

  return (
    path: string,
    method: "get" | "post" | "patch" | "delete",
    config: RegisterConfig,
  ) => {
    const key = `${method.toUpperCase()} ${path}`;
    if (operations.has(key)) {
      throw new Error(`Duplicate Agent OpenAPI operation: ${key}`);
    }
    operations.add(key);
    targetRegistry.registerPath({
      path,
      method,
      operationId: `${method}_${path
        .replace(/^\/api\/v1\/agent\/?/, "")
        .replaceAll(/[{}]/g, "")
        .replaceAll(/[^a-zA-Z0-9]+/g, "_")
        .replaceAll(/^_|_$/g, "")}`,
      tags: ["Agent"],
      security: [{ bearerAuth: [] }],
      ...config,
      responses: { ...config.responses, ...errorResponses },
    });
  };
};

const register = createRouteRegistrar(registry);

register("/api/v1/agent/models", "get", {
  summary: "List available agent models",
  responses: { 200: jsonResponse(JsonObject) },
});
register("/api/v1/agent/sessions", "get", {
  summary: "List agent sessions",
  responses: { 200: jsonResponse(z.object({ sessions: z.array(SessionSummary) })) },
});
register("/api/v1/agent/sessions", "post", {
  summary: "Create an agent session",
  request: {
    body: {
      content: {
        "application/json": {
          schema: z.object({
            session_id: z.string().max(128).optional(),
            parent_session_id: z.string().max(128).optional(),
          }),
        },
      },
    },
  },
  responses: {
    200: jsonResponse(SessionCreate, "Existing session returned"),
    201: jsonResponse(SessionCreate, "Session created"),
  },
});
register("/api/v1/agent/sessions/{session_id}", "get", {
  summary: "Get an agent session",
  request: { params: SessionId },
  responses: { 200: jsonResponse(SessionDetail) },
});
register("/api/v1/agent/sessions/{session_id}", "patch", {
  summary: "Update an agent session",
  request: {
    params: SessionId,
    body: {
      content: {
        "application/json": {
          schema: z.object({
            title: z.string().min(1).max(120),
            is_title_manually_edited: z.boolean().optional(),
          }),
        },
      },
    },
  },
  responses: { 200: jsonResponse(SessionUpdate) },
});
register("/api/v1/agent/sessions/{session_id}", "delete", {
  summary: "Delete an agent session",
  request: { params: SessionId },
  responses: { 204: { description: "Session deleted" } },
});
register("/api/v1/agent/sessions/{session_id}/runs", "post", {
  summary: "Run an agent session",
  request: {
    params: SessionId,
    body: {
      content: {
        "application/json": {
          schema: z.object({
            message: z.string().min(1).max(10000),
            model: z.string().optional(),
            approval_mode: z.enum(["request", "always"]).optional(),
          }),
        },
      },
    },
  },
  responses: {
    200: {
      description: "Agent event stream",
      content: { "text/event-stream": { schema: z.string() } },
    },
  },
});
register(
  "/api/v1/agent/sessions/{session_id}/runs/current/events",
  "get",
  {
    summary: "Resume the current agent event stream",
    request: { params: SessionId },
    responses: {
      200: {
        description: "Agent event stream",
        content: { "text/event-stream": { schema: z.string() } },
      },
    },
  },
);
register("/api/v1/agent/sessions/{session_id}/runs/current", "delete", {
  summary: "Abort the current agent run",
  request: { params: SessionId },
  responses: { 202: jsonResponse(JsonObject), 204: { description: "No active run" } },
});
register(
  "/api/v1/agent/sessions/{session_id}/permission-responses",
  "post",
  {
    summary: "Reply to an agent permission request",
    request: {
      params: SessionId,
      body: {
        content: {
          "application/json": {
            schema: z.object({
              request_id: z.string(),
              reply: z.enum(["once", "always", "reject"]),
              message: z.string().max(1000).optional(),
            }),
          },
        },
      },
    },
    responses: { 202: jsonResponse(JsonObject) },
  },
);
register(
  "/api/v1/agent/sessions/{session_id}/question-responses",
  "post",
  {
    summary: "Reply to or reject an agent question",
    request: {
      params: SessionId,
      body: {
        content: {
          "application/json": {
            schema: z.object({
              request_id: z.string(),
              action: z.enum(["reply", "reject"]).default("reply"),
              answers: z.array(z.array(z.string().max(2000))).optional(),
            }),
          },
        },
      },
    },
    responses: { 202: jsonResponse(JsonObject) },
  },
);
register("/api/v1/agent/sessions/{session_id}/forks", "post", {
  summary: "Fork an agent session",
  request: {
    params: SessionId,
    body: {
      content: {
        "application/json": {
          schema: z.object({
            keep_message_count: z.number().int().nonnegative(),
          }),
        },
      },
    },
  },
  responses: { 200: jsonResponse(SessionFork) },
});
register("/api/v1/agent/render-references/{render_ref}", "get", {
  summary: "Resolve a render reference",
  request: { params: RenderRef, query: RenderReferenceQuery },
  responses: { 200: jsonResponse(JsonObject) },
});

export const generateAgentOpenApi = () => {
  const generator = new OpenApiGeneratorV3(registry.definitions);
  return generator.generateDocument({
    openapi: "3.0.3",
    info: {
      title: "TJWater Agent API",
      version: "1.0.0",
      description: "Public REST API for TJWater Agent sessions and runs",
    },
  });
};
