import { z } from "zod";
import {
  defineRecordCodec,
  recordEnvelopeSchema,
  type RecordCodec,
} from "./envelope.js";

export const viewportSchema = z
  .object({
    width: z.number(),
    height: z.number(),
    deviceScaleFactor: z.number(),
  })
  .strict();

export const actionStepSchema = z
  .object({
    id: z.string(),
    kind: z.string(),
    target: z.string().optional(),
    params: z.record(z.string(), z.string()).optional(),
  })
  .strict();

export const networkExpectationSchema = z
  .object({
    method: z.string(),
    urlPattern: z.string(),
    status: z.number().optional(),
  })
  .strict();

export const matrixCellRefSchema = z
  .object({
    id: z.string(),
    routeId: z.string().optional(),
    url: z.string().optional(),
    role: z.string().optional(),
    locale: z.string().optional(),
    viewport: viewportSchema.optional(),
  })
  .strict();

export const actionRefSchema = z
  .object({
    id: z.string(),
  })
  .strict();

export const journeyRefSchema = z
  .object({
    id: z.string(),
  })
  .strict();

export const MATRIX_CELL_V1_VERSION = "matrix-cell/v1" as const;
export const MATRIX_CELL_VERSION = "matrix-cell/v2" as const;

const matrixCellBaseSchema = recordEnvelopeSchema
  .extend({
    routeId: z.string(),
    route: z.string().optional(),
    url: z.string(),
    role: z.string(),
    locale: z.string(),
    viewport: viewportSchema,
    state: z.record(z.string(), z.string()),
    actions: z.array(actionStepSchema),
    seed: z.string(),
    expectedNetwork: z.array(networkExpectationSchema).optional(),
  })
  .strict();

const matrixCellV1Schema = matrixCellBaseSchema.extend({
  schemaVersion: z.literal(MATRIX_CELL_V1_VERSION),
});

const matrixCellV2Schema = matrixCellBaseSchema.extend({
  schemaVersion: z.literal(MATRIX_CELL_VERSION),
  journeys: z.array(journeyRefSchema).min(1),
});

export const matrixCellSchema = z.discriminatedUnion("schemaVersion", [
  matrixCellV1Schema,
  matrixCellV2Schema,
]);

export type JourneyRef = z.infer<typeof journeyRefSchema>;
export type MatrixCell = Omit<z.infer<typeof matrixCellBaseSchema>, "schemaVersion"> & {
  schemaVersion: typeof MATRIX_CELL_V1_VERSION | typeof MATRIX_CELL_VERSION;
  journeys?: JourneyRef[] | undefined;
};

export function matrixCellRef(cell: MatrixCell): MatrixCellRef {
  return {
    id: cell.id,
    routeId: cell.routeId,
    url: cell.url,
    role: cell.role,
    locale: cell.locale,
    viewport: cell.viewport,
  };
}

export const matrixCellCodec: RecordCodec<MatrixCell> = defineRecordCodec({
  schema: matrixCellSchema,
  currentVersion: MATRIX_CELL_VERSION,
  supportedVersions: [MATRIX_CELL_V1_VERSION, MATRIX_CELL_VERSION],
});

const nonBrowserContextKindSchema = z.enum([
  "api",
  "cli",
  "worker",
  "stream",
  "filesystem",
  "package",
  "deployment",
]);

const nonBrowserExecutionContextSchema = z
  .object({
    kind: nonBrowserContextKindSchema,
    surfaceId: z.string(),
    adapterId: z.string(),
    environment: z.record(z.string(), z.string()),
    action: actionRefSchema.optional(),
    seed: z.string(),
  })
  .strict();

const browserExecutionContextSchema = z
  .object({
    kind: z.literal("browser"),
    cell: matrixCellRefSchema,
  })
  .strict();

export type MatrixCellRef = z.infer<typeof matrixCellRefSchema>;
export type ActionRef = z.infer<typeof actionRefSchema>;
export type NonBrowserExecutionContextKind = z.infer<
  typeof nonBrowserContextKindSchema
>;

export type BrowserExecutionContext = {
  kind: "browser";
  cell: MatrixCellRef;
};

export type NonBrowserExecutionContextBase = {
  surfaceId: string;
  adapterId: string;
  environment: Record<string, string>;
  action?: ActionRef | undefined;
  seed: string;
};

export type ApiExecutionContext = NonBrowserExecutionContextBase & { kind: "api" };
export type CliExecutionContext = NonBrowserExecutionContextBase & { kind: "cli" };
export type WorkerExecutionContext = NonBrowserExecutionContextBase & {
  kind: "worker";
};
export type StreamExecutionContext = NonBrowserExecutionContextBase & {
  kind: "stream";
};
export type FilesystemExecutionContext = NonBrowserExecutionContextBase & {
  kind: "filesystem";
};
export type PackageExecutionContext = NonBrowserExecutionContextBase & {
  kind: "package";
};
export type DeploymentExecutionContext = NonBrowserExecutionContextBase & {
  kind: "deployment";
};

export type NonBrowserExecutionContext =
  | ApiExecutionContext
  | CliExecutionContext
  | WorkerExecutionContext
  | StreamExecutionContext
  | FilesystemExecutionContext
  | PackageExecutionContext
  | DeploymentExecutionContext;

export type ExecutionContext = BrowserExecutionContext | NonBrowserExecutionContext;

export const executionContextSchema = z.discriminatedUnion("kind", [
  browserExecutionContextSchema,
  ...nonBrowserContextKindSchema.options.map((kind) =>
    nonBrowserExecutionContextSchema.extend({
      kind: z.literal(kind),
    }),
  ),
]) as z.ZodType<ExecutionContext>;