import { z } from "zod";

import { DEFAULT_MODEL, resolveModelSlug } from "./model-logic";
import {
  DEFAULT_THREAD_TERMINAL_HEIGHT,
  DEFAULT_RUNTIME_MODE,
  type Project,
  type RuntimeMode,
  type Thread,
} from "./types";

const LEGACY_DEFAULT_MODEL = "gpt-5.2-codex";

const persistedProjectSchema = z.object({
  id: z.string().min(1),
  name: z.string().min(1),
  cwd: z.string().min(1),
  model: z.string().min(1),
  expanded: z.boolean(),
});

const persistedMessageSchema = z.object({
  id: z.string().min(1),
  role: z.enum(["user", "assistant"]),
  text: z.string(),
  attachments: z
    .array(
      z.object({
        type: z.literal("image"),
        id: z.string().min(1),
        name: z.string().min(1),
        mimeType: z.string().min(1),
        sizeBytes: z.number().int().min(1),
      }),
    )
    .optional(),
  createdAt: z.string().min(1),
  streaming: z.boolean(),
});

const persistedThreadSchema = z.object({
  id: z.string().min(1),
  codexThreadId: z.string().min(1).nullable().default(null),
  projectId: z.string().min(1),
  title: z.string().min(1),
  model: z.string().min(1),
  terminalOpen: z.boolean().default(false),
  terminalHeight: z.number().int().min(120).max(4_096).default(
    DEFAULT_THREAD_TERMINAL_HEIGHT,
  ),
  messages: z.array(persistedMessageSchema),
  createdAt: z.string().min(1),
  lastVisitedAt: z.string().min(1).optional(),
  branch: z.string().min(1).nullable().optional(),
  worktreePath: z.string().min(1).nullable().optional(),
});

const persistedStateBodySchema = z.object({
  projects: z.array(persistedProjectSchema),
  threads: z.array(persistedThreadSchema),
  activeThreadId: z.string().min(1).nullable(),
});

const runtimeModeSchema = z.enum(["approval-required", "full-access"]);

export const persistedStateV1Schema = persistedStateBodySchema.extend({
  version: z.literal(1).optional(),
});

export const persistedStateV2Schema = persistedStateBodySchema.extend({
  version: z.literal(2).optional(),
});

export const persistedStateV3Schema = persistedStateBodySchema.extend({
  runtimeMode: runtimeModeSchema.default(DEFAULT_RUNTIME_MODE),
  version: z.literal(3).optional(),
});

export const persistedStateV4Schema = persistedStateBodySchema.extend({
  runtimeMode: runtimeModeSchema.default(DEFAULT_RUNTIME_MODE),
  version: z.literal(4).optional(),
});

export const persistedStateV6Schema = persistedStateBodySchema.extend({
  runtimeMode: runtimeModeSchema.default(DEFAULT_RUNTIME_MODE),
  version: z.literal(6).optional(),
});

export const persistedStateV5Schema = persistedStateBodySchema.extend({
  runtimeMode: runtimeModeSchema.default(DEFAULT_RUNTIME_MODE),
  version: z.literal(5).optional(),
});

const persistedStateSchema = z.union([
  persistedStateV6Schema,
  persistedStateV5Schema,
  persistedStateV4Schema,
  persistedStateV3Schema,
  persistedStateV2Schema,
  persistedStateV1Schema,
]);

export interface PersistedStoreSnapshot {
  projects: Project[];
  threads: Thread[];
  activeThreadId: string | null;
  runtimeMode: RuntimeMode;
}

function maybeMigrateLegacyModel(model: string, isLegacyPayload: boolean): string {
  if (!isLegacyPayload) {
    return model;
  }

  return model === LEGACY_DEFAULT_MODEL ? DEFAULT_MODEL : model;
}

function hydrateProject(
  project: z.infer<typeof persistedProjectSchema>,
  isLegacyPayload: boolean,
): Project {
  return {
    ...project,
    model: resolveModelSlug(maybeMigrateLegacyModel(project.model, isLegacyPayload)),
  };
}

function hydrateThread(
  thread: z.infer<typeof persistedThreadSchema>,
  isLegacyPayload: boolean,
): Thread {
  return {
    id: thread.id,
    codexThreadId: thread.codexThreadId,
    projectId: thread.projectId,
    title: thread.title,
    model: resolveModelSlug(maybeMigrateLegacyModel(thread.model, isLegacyPayload)),
    terminalOpen: thread.terminalOpen ?? false,
    terminalHeight: thread.terminalHeight ?? DEFAULT_THREAD_TERMINAL_HEIGHT,
    session: null,
    messages: thread.messages.map((message) => {
      const hydratedAttachments = message.attachments?.map((attachment) => ({ ...attachment }));
      return {
        id: message.id,
        role: message.role,
        text: message.text,
        ...(hydratedAttachments && hydratedAttachments.length > 0
          ? { attachments: hydratedAttachments }
          : {}),
        createdAt: message.createdAt,
        streaming: false,
      };
    }),
    events: [],
    error: null,
    createdAt: thread.createdAt,
    lastVisitedAt: thread.lastVisitedAt,
    branch: thread.branch ?? null,
    worktreePath: thread.worktreePath ?? null,
  };
}

export function hydratePersistedState(
  raw: string,
  isLegacyPayload: boolean,
): PersistedStoreSnapshot | null {
  let parsedJson: unknown;
  try {
    parsedJson = JSON.parse(raw);
  } catch {
    return null;
  }

  const parsedState = persistedStateSchema.safeParse(parsedJson);
  if (!parsedState.success) {
    return null;
  }

  const projects = parsedState.data.projects.map((project) =>
    hydrateProject(project, isLegacyPayload),
  );
  const projectIds = new Set(projects.map((project) => project.id));
  const threads = parsedState.data.threads
    .map((thread) => hydrateThread(thread, isLegacyPayload))
    .filter((thread) => projectIds.has(thread.projectId));
  const hasActiveThread = Boolean(
    parsedState.data.activeThreadId &&
    threads.some((thread) => thread.id === parsedState.data.activeThreadId),
  );

  return {
    projects,
    threads,
    activeThreadId: hasActiveThread ? parsedState.data.activeThreadId : (threads[0]?.id ?? null),
    runtimeMode:
      "runtimeMode" in parsedState.data ? parsedState.data.runtimeMode : DEFAULT_RUNTIME_MODE,
  };
}

export function toPersistedState(
  state: PersistedStoreSnapshot,
): z.infer<typeof persistedStateV6Schema> {
  return {
    version: 6,
    projects: state.projects,
    threads: state.threads.map((thread) => ({
      id: thread.id,
      codexThreadId: thread.codexThreadId,
      projectId: thread.projectId,
      title: thread.title,
      model: thread.model,
      terminalOpen: thread.terminalOpen,
      terminalHeight: thread.terminalHeight,
      messages: thread.messages.map((message) => ({
        id: message.id,
        role: message.role,
        text: message.text,
        ...(message.attachments && message.attachments.length > 0
          ? {
              attachments: message.attachments.map((attachment) => ({
                type: attachment.type,
                id: attachment.id,
                name: attachment.name,
                mimeType: attachment.mimeType,
                sizeBytes: attachment.sizeBytes,
              })),
            }
          : {}),
        createdAt: message.createdAt,
        streaming: message.streaming,
      })),
      createdAt: thread.createdAt,
      ...(thread.lastVisitedAt ? { lastVisitedAt: thread.lastVisitedAt } : {}),
      branch: thread.branch,
      worktreePath: thread.worktreePath,
    })),
    activeThreadId: state.activeThreadId,
    runtimeMode: state.runtimeMode,
  };
}