import { z } from "zod";
import { ALL_PHASES } from "@/conversations/phases";

const PhaseSchema = z
    .string()
    .transform((val) => val.toLowerCase())
    .pipe(
        z.enum(ALL_PHASES as [string, ...string[]])
    );

export const PhaseTransitionSchema = z.object({
    from: PhaseSchema,
    to: PhaseSchema,
    message: z.string(),
    timestamp: z.number(),
    agentPubkey: z.string(),
    agentName: z.string(),
    reason: z.string().optional(),
    // Enhanced handoff fields
    summary: z.string().optional(),
    requirements: z.string().optional(),
    artifacts: z.any().optional(),
    goal: z.string().optional(),
});

export const ConversationMetadataSchema = z.record(z.string(), z.unknown());

const ExecutionTimeSchema = z.object({
    totalSeconds: z.number(),
    currentSessionStart: z.number().optional(),
    isActive: z.boolean(),
    lastUpdated: z.number(),
});

// Simplified agent state schema
export const AgentStateSchema = z.object({
    lastProcessedMessageIndex: z.number().int().min(0),
    claudeSessionId: z.string().optional(),
});

const CompletionSchema = z.object({
    agent: z.string(),
    message: z.string(),
    timestamp: z.number().optional(),
});

const OrchestratorTurnSchema = z.object({
    turnId: z.string(),
    timestamp: z.number(),
    phase: PhaseSchema,
    agents: z.array(z.string()),
    completions: z.array(CompletionSchema),
    reason: z.string().optional(),
    isCompleted: z.boolean(),
});

export const SerializedConversationSchema = z.object({
    id: z.string(),
    title: z.string(),
    phase: PhaseSchema,
    history: z.array(z.string()),
    agentStates: z.record(z.string(), AgentStateSchema).optional(), // Map serialized as object
    phaseStartedAt: z.number().optional(),
    metadata: ConversationMetadataSchema,
    phaseTransitions: z.array(PhaseTransitionSchema).default([]),
    orchestratorTurns: z.array(OrchestratorTurnSchema).default([]),
    executionTime: ExecutionTimeSchema.optional(),
});

export const ConversationMetadataFileSchema = z.object({
    id: z.string(),
    title: z.string(),
    createdAt: z.number(),
    updatedAt: z.number(),
    phase: z.string(),
    eventCount: z.number(),
    agentCount: z.number(),
    archived: z.boolean().optional(),
});

export const MetadataFileSchema = z.object({
    conversations: z.array(ConversationMetadataFileSchema),
});

export type SerializedConversation = z.infer<typeof SerializedConversationSchema>;
export type MetadataFile = z.infer<typeof MetadataFileSchema>;