import { z } from 'zod';

export const UI_STREAM_HOST_CONTEXT_KEY = 'pi-mcp-adapter/stream';
export const UI_STREAM_REQUEST_META_KEY = 'pi-mcp-adapter/stream-token';
export const UI_STREAM_RESULT_PATCH_METHOD = 'notifications/pi-mcp-adapter/ui-result-patch';
export const SERVER_STREAM_RESULT_PATCH_METHOD = 'notifications/pi-mcp-adapter/result-patch';
export const UI_STREAM_STRUCTURED_CONTENT_KEY = 'pi-mcp-adapter/stream';

export const uiStreamModeSchema = z.enum(['eager', 'stream-first']);
export type UiStreamMode = z.infer<typeof uiStreamModeSchema>;

export const visualizationStreamPhaseSchema = z.enum(['shell', 'narrative', 'structure', 'detail', 'settled']);
export type VisualizationStreamPhase = z.infer<typeof visualizationStreamPhaseSchema>;

export const visualizationStreamFrameTypeSchema = z.enum(['patch', 'checkpoint', 'final']);
export type VisualizationStreamFrameType = z.infer<typeof visualizationStreamFrameTypeSchema>;

export const visualizationStreamStatusSchema = z.enum(['ok', 'error']);
export type VisualizationStreamStatus = z.infer<typeof visualizationStreamStatusSchema>;

const looseRecordSchema = z.record(z.string(), z.unknown());
const looseArraySchema = z.array(z.unknown());

export const uiStreamHostContextSchema = z.object({
  mode: uiStreamModeSchema,
  streamId: z.string().min(1),
  intermediateResultPatches: z.boolean(),
  partialInput: z.boolean(),
});
export type UiStreamHostContext = z.infer<typeof uiStreamHostContextSchema>;

export const visualizationStreamEnvelopeSchema = z.object({
  streamId: z.string().min(1),
  sequence: z.number().int().nonnegative(),
  frameType: visualizationStreamFrameTypeSchema,
  phase: visualizationStreamPhaseSchema,
  status: visualizationStreamStatusSchema,
  message: z.string().optional(),
  spec: looseRecordSchema.optional(),
  checkpoint: looseRecordSchema.optional(),
});
export type VisualizationStreamEnvelope = z.infer<typeof visualizationStreamEnvelopeSchema>;

export const uiStreamCallToolResultSchema = z
  .object({
    content: looseArraySchema.optional(),
    structuredContent: looseRecordSchema.optional(),
    isError: z.boolean().optional(),
    _meta: looseRecordSchema.optional(),
  })
  .passthrough();
export type UiStreamCallToolResult = z.infer<typeof uiStreamCallToolResultSchema>;

export const uiStreamResultPatchNotificationSchema = z.object({
  method: z.literal(UI_STREAM_RESULT_PATCH_METHOD),
  params: uiStreamCallToolResultSchema,
});
export type UiStreamResultPatchNotification = z.infer<typeof uiStreamResultPatchNotificationSchema>;

export const serverStreamResultPatchNotificationSchema = z.object({
  method: z.literal(SERVER_STREAM_RESULT_PATCH_METHOD),
  params: z.object({
    streamToken: z.string().min(1),
    result: uiStreamCallToolResultSchema,
  }),
});
export type ServerStreamResultPatchNotification = z.infer<typeof serverStreamResultPatchNotificationSchema>;

export interface UiStreamSummary {
  streamId: string;
  mode: UiStreamMode;
  frames: number;
  phases: VisualizationStreamPhase[];
  finalStatus?: VisualizationStreamStatus;
  lastMessage?: string;
}

export function getUiStreamHostContext(
  hostContext: Record<string, unknown> | undefined
): UiStreamHostContext | undefined {
  const candidate = hostContext?.[UI_STREAM_HOST_CONTEXT_KEY];
  const parsed = uiStreamHostContextSchema.safeParse(candidate);

  return parsed.success ? parsed.data : undefined;
}

export function getVisualizationStreamEnvelope(structuredContent: unknown): VisualizationStreamEnvelope | undefined {
  if (!structuredContent || typeof structuredContent !== 'object' || Array.isArray(structuredContent)) {
    return undefined;
  }

  const candidate = (structuredContent as Record<string, unknown>)[UI_STREAM_STRUCTURED_CONTENT_KEY];
  const parsed = visualizationStreamEnvelopeSchema.safeParse(candidate);

  return parsed.success ? parsed.data : undefined;
}