import { z } from "zod/v4";
import { Identifier } from "../id.js";
import { Examples } from "../examples.js";

// Basic types we need
export const MachinePresetName = z.enum([
  "ns3-micro",
  "ns3-mini",
  "ns3-small",
  "ns3-medium",
  "ns3-large",
]);

export const ScreenResolutionPresetName = z.enum([
  "1280x720",
  "1080x1920",
  "1440x2560",
  "2160x3840",
]);

export const VideoCodecName = z.enum(["h264", "h265", "vp8", "vp9", "av1"]);

export const AudioCodecName = z.enum(["opus", "aac", "mp3", "pcm"]);

export const RetryOptions = z.object({
  maxAttempts: z.number().int().optional(),
  factor: z.number().optional(),
  minTimeoutInMs: z.number().int().optional(),
  maxTimeoutInMs: z.number().int().optional(),
  randomize: z.boolean().optional(),
  outOfMemory: z
    .object({
      machine: MachinePresetName.optional(),
    })
    .optional(),
});

export namespace Session {
  export const Info = z
    .object({
      id: Identifier.schema("session").meta({
        description: "The unique identifier of the session",
        example: Examples.Session.id,
      }),
      description: z.string().optional().meta({
        description: "Description of what the session does",
        example: Examples.Session.description,
      }),
      retry: RetryOptions.optional().meta({
        description: "Retry configuration for failed session attempts",
        example: Examples.Session.retry,
      }),
      queue: z
        .object({
          name: z.string().optional(),
          concurrencyLimit: z.number().optional(),
        })
        .optional()
        .meta({
          description: "Queue configuration for concurrent session execution",
          example: Examples.Session.queue,
        }),
      machine: MachinePresetName.optional().meta({
        description: "Machine preset for the session",
        example: Examples.Session.machine,
      }),
      maxDuration: z.number().optional().meta({
        description: "Maximum duration in seconds the session can run",
        example: Examples.Session.maxDuration,
      }),
      relayUrl: z.url().optional().meta({
        description: "Nestri relay URL for streaming and signaling",
        example: Examples.Session.relayUrl,
      }),
      resolution: ScreenResolutionPresetName.optional().meta({
        description: "Display or stream resolution in 'WxH' format",
        example: Examples.Session.resolution,
      }),
      framerate: z.number().optional().meta({
        description: "Display or stream framerate",
        example: Examples.Session.framerate,
      }),
      room: z.string().optional().meta({
        description: "Room name or identifier for multiplayer sessions",
        example: Examples.Session.room,
      }),
      gpuVendor: z.string().optional().meta({
        description: "Preferred GPU vendor",
        example: Examples.Session.gpuVendor,
      }),
      gpuName: z.string().optional().meta({
        description: "Preferred GPU name",
        example: Examples.Session.gpuName,
      }),
      gpuIndex: z.number().optional().meta({
        description: "Specific GPU index to use",
        example: Examples.Session.gpuIndex,
      }),
      gpuCardPath: z.string().optional().meta({
        description: "Force specific GPU card by /dev/dri path",
        example: Examples.Session.gpuCardPath,
      }),
      videoCodec: VideoCodecName.optional().meta({
        description: "Preferred video codec for streaming",
        example: Examples.Session.videoCodec,
      }),
      videoEncoder: z.string().optional().meta({
        description: "Override video encoder implementation",
        example: Examples.Session.videoEncoder,
      }),
      videoRateControl: z
        .enum(["cbr", "vbr", "cq", "lossless"])
        .optional()
        .meta({
          description: "Video rate control method",
          example: Examples.Session.videoRateControl,
        }),
      videoCqp: z.number().min(1).max(51).optional().meta({
        description: "Constant Quantization Parameter for quality control",
        example: Examples.Session.videoCqp,
      }),
      videoBitrate: z.number().min(1).optional().meta({
        description: "Target video bitrate in kbps",
        example: Examples.Session.videoBitrate,
      }),
      videoBitrateMax: z.number().min(1).optional().meta({
        description: "Maximum video bitrate in kbps",
        example: Examples.Session.videoBitrateMax,
      }),
      videoEncoderType: z.enum(["hardware", "software"]).optional().meta({
        description: "Encoder type: hardware or software",
        example: Examples.Session.videoEncoderType,
      }),
      videoBitDepth: z
        .union([z.literal(8), z.literal(10)])
        .optional()
        .meta({
          description: "Video bit depth (8 or 10)",
          example: Examples.Session.videoBitDepth,
        }),
      audioCaptureMethod: z
        .enum(["pipewire", "pulse", "alsa"])
        .optional()
        .meta({
          description: "Audio capture method",
          example: Examples.Session.audioCaptureMethod,
        }),
      audioCodec: AudioCodecName.optional().meta({
        description: "Preferred audio codec for streaming",
        example: Examples.Session.audioCodec,
      }),
      audioEncoder: z.string().optional().meta({
        description: "Override audio encoder binary",
        example: Examples.Session.audioEncoder,
      }),
      audioRateControl: z.enum(["cbr", "vbr"]).optional().meta({
        description: "Audio rate control method",
        example: Examples.Session.audioRateControl,
      }),
      audioBitrate: z.number().min(1).optional().meta({
        description: "Target audio bitrate in kbps",
        example: Examples.Session.audioBitrate,
      }),
      audioBitrateMax: z.number().optional().meta({
        description: "Maximum audio bitrate in kbps",
        example: Examples.Session.audioBitrateMax,
      }),
      zeroCopy: z.boolean().optional().meta({
        description: "Whether to use DMA-BUF for zero-copy video pipelines",
        example: Examples.Session.zeroCopy,
      }),
      verbose: z.boolean().optional().meta({
        description: "Enable verbose output for debugging",
        example: Examples.Session.verbose,
      }),
    })
    .meta({
      ref: "Session",
      description:
        "Represents a session configuration with streaming and execution options",
      example: Examples.Session,
    });

  export type Info = z.infer<typeof Info>;

  // Session interface that matches the SDK session object
  export interface Session {
    /** Optional description of the session */
    description?: string;

    /** Start a session without waiting for the result */
    start: (options?: any, requestOptions?: any) => Promise<any>;

    /** Start a session and wait for the result */
    startAndWait: (options?: any) => any;
  }

  // Session configuration type that includes the run function
  export type Config = Omit<Info, "id"> & {
    /** The function that executes when the session is started */
    run: (ctx: any) => Promise<any>;
  };
}
