import { createHash } from "node:crypto";
import { validateSummary } from "./context-policy.ts";
import { existsSync, readFileSync } from "node:fs";
import { homedir } from "node:os";
import { TextDecoder } from "node:util";
import { join } from "node:path";
import type {
  BtwAttempt,
  BtwEntry,
  BtwScope,
  BtwScopeV3,
  BtwScopeV4,
  BtwState,
  BtwStateV1,
  BtwStateV2,
  BtwStateV3,
  BtwStateV4,
  BtwSummary,
  BtwSummaryMeter,
  BtwSummarySourceRef,
  BtwThread,
  BtwThreadV4,
  BtwUsageTotals,
  FlatEntry,
  FlatThread,
} from "./threads.ts";

export type StoreFileResult =
  | { kind: "missing" }
  | { kind: "v1"; state: BtwStateV1 }
  | { kind: "v2"; state: BtwStateV2 }
  | { kind: "v3"; state: BtwStateV3 }
  | { kind: "v4"; state: BtwStateV4 }
  | { kind: "v5"; state: BtwState }
  | { kind: "corrupt" }
  | { kind: "future"; version: number };

type JsonObject = Record<string, unknown>;
type ScopeBase<TThread> = Omit<BtwScope, "threads"> & { threads: TThread[] };

export function resolveStorePath(cwd: string, home: string = homedir()): string {
  const key = createHash("sha256").update(cwd).digest("hex").slice(0, 12);
  return join(home, ".pi", "agent", "btw", `threads-${key}.json`);
}

function object(value: unknown): JsonObject | null {
  return value !== null && typeof value === "object" && !Array.isArray(value)
    ? (value as JsonObject)
    : null;
}

function string(value: unknown): value is string {
  return typeof value === "string";
}

function id(value: unknown): value is string {
  return string(value) && value.length > 0;
}

function onlyKeys(value: JsonObject, allowed: readonly string[]): boolean {
  return Object.keys(value).every((key) => allowed.includes(key));
}

function parseGrounding(
  value: unknown,
): { capturedAt: string; model: string; contextInfo: string } | null {
  const item = object(value);
  if (
    !item ||
    !onlyKeys(item, ["capturedAt", "model", "contextInfo"]) ||
    !string(item.capturedAt) ||
    !string(item.model) ||
    !string(item.contextInfo)
  ) {
    return null;
  }
  return {
    capturedAt: item.capturedAt,
    model: item.model,
    contextInfo: item.contextInfo,
  };
}

function parseStringList(value: unknown): string[] | null {
  return Array.isArray(value) && value.every(string) ? [...value] : null;
}

/** Strict parser for the historical flat entry shape used only by v1-v3 and replay. */
export function parseFlatEntry(value: unknown): FlatEntry | null {
  const item = object(value);
  if (
    !item ||
    !onlyKeys(item, ["id", "mode", "question", "answer", "grounding", "toolsUsed", "promoted", "error"]) ||
    !id(item.id) ||
    (item.mode !== "quick" && item.mode !== "deep") ||
    !string(item.question) ||
    !string(item.answer)
  ) {
    return null;
  }

  const grounding = parseGrounding(item.grounding);
  const toolsUsed = item.toolsUsed === undefined ? undefined : parseStringList(item.toolsUsed);
  if (
    !grounding ||
    (item.toolsUsed !== undefined && !toolsUsed) ||
    (item.promoted !== undefined && typeof item.promoted !== "boolean") ||
    (item.error !== undefined && !string(item.error))
  ) {
    return null;
  }

  return {
    id: item.id,
    mode: item.mode,
    question: item.question,
    answer: item.answer,
    grounding,
    ...(toolsUsed ? { toolsUsed } : {}),
    ...(item.promoted === undefined ? {} : { promoted: item.promoted }),
    ...(item.error === undefined ? {} : { error: item.error }),
  };
}

export function parseAttempt(value: unknown): BtwAttempt | null {
  const item = object(value);
  if (
    !item ||
    !onlyKeys(item, ["id", "mode", "answer", "grounding", "toolsUsed", "error"]) ||
    !id(item.id) ||
    (item.mode !== "quick" && item.mode !== "deep") ||
    !string(item.answer)
  ) {
    return null;
  }

  const grounding = parseGrounding(item.grounding);
  const toolsUsed = item.toolsUsed === undefined ? undefined : parseStringList(item.toolsUsed);
  if (!grounding || (item.toolsUsed !== undefined && !toolsUsed) || (item.error !== undefined && !string(item.error))) {
    return null;
  }

  return {
    id: item.id,
    mode: item.mode,
    answer: item.answer,
    grounding,
    ...(toolsUsed ? { toolsUsed } : {}),
    ...(item.error === undefined ? {} : { error: item.error }),
  };
}

function finiteNonnegative(value: unknown): value is number {
  return typeof value === "number" && Number.isFinite(value) && value >= 0;
}

export function parseUsageTotals(value: unknown): BtwUsageTotals | null {
  const item = object(value);
  const cost = item ? object(item.cost) : null;
  if (!item || !cost || !onlyKeys(item, ["input", "output", "cacheRead", "cacheWrite", "totalTokens", "reasoning", "cacheWrite1h", "cost"]) ||
    !onlyKeys(cost, ["input", "output", "cacheRead", "cacheWrite", "total"]) ||
    ![item.input, item.output, item.cacheRead, item.cacheWrite, item.totalTokens, cost.input, cost.output, cost.cacheRead, cost.cacheWrite, cost.total].every(finiteNonnegative) ||
    (item.reasoning !== undefined && !finiteNonnegative(item.reasoning)) ||
    (item.cacheWrite1h !== undefined && !finiteNonnegative(item.cacheWrite1h))) return null;
  return {
    input: item.input as number, output: item.output as number, cacheRead: item.cacheRead as number, cacheWrite: item.cacheWrite as number, totalTokens: item.totalTokens as number,
    ...(item.reasoning === undefined ? {} : { reasoning: item.reasoning as number }),
    ...(item.cacheWrite1h === undefined ? {} : { cacheWrite1h: item.cacheWrite1h as number }),
    cost: { input: cost.input as number, output: cost.output as number, cacheRead: cost.cacheRead as number, cacheWrite: cost.cacheWrite as number, total: cost.total as number },
  };
}

function parseSummarySource(value: unknown): BtwSummarySourceRef[] | null {
  if (!Array.isArray(value) || !value.length) return null;
  const refs = value.map((item) => {
    const objectItem = object(item);
    return objectItem && onlyKeys(objectItem, ["entryId", "attemptId"]) && id(objectItem.entryId) && id(objectItem.attemptId)
      ? { entryId: objectItem.entryId, attemptId: objectItem.attemptId }
      : null;
  });
  if (refs.some((ref) => !ref)) return null;
  const parsed = refs as BtwSummarySourceRef[];
  return new Set(parsed.map((ref) => ref.entryId)).size === parsed.length ? parsed : null;
}

function parseSummary(value: unknown): BtwSummary | null {
  const item = object(value);
  if (!item || !onlyKeys(item, ["text", "throughEntryId", "source", "sourceHash", "createdAt", "model"]) ||
    !id(item.text) || !id(item.throughEntryId) || !string(item.sourceHash) || !/^[a-f0-9]{64}$/.test(item.sourceHash) || !string(item.createdAt) || !string(item.model)) return null;
  const source = parseSummarySource(item.source);
  if (!source || source.at(-1)?.entryId !== item.throughEntryId) return null;
  return { text: item.text, throughEntryId: item.throughEntryId, source, sourceHash: item.sourceHash, createdAt: item.createdAt, model: item.model };
}

function safeCount(value: unknown): value is number {
  return typeof value === "number" && Number.isSafeInteger(value) && value >= 0;
}

function parseSummaryMeter(value: unknown): BtwSummaryMeter | null {
  const item = object(value);
  if (!item || !onlyKeys(item, ["runIds", "requests", "committed", "stale", "failed", "aborted", "usageKnownRuns", "usage"]) ||
    !Array.isArray(item.runIds) || !item.runIds.every(id) || new Set(item.runIds).size !== item.runIds.length ||
    ![item.requests, item.committed, item.stale, item.failed, item.aborted, item.usageKnownRuns].every(safeCount)) return null;
  const usage = parseUsageTotals(item.usage);
  const requests = item.requests as number;
  const committed = item.committed as number;
  const stale = item.stale as number;
  const failed = item.failed as number;
  const aborted = item.aborted as number;
  const usageKnownRuns = item.usageKnownRuns as number;
  if (!usage || requests !== item.runIds.length || requests !== committed + stale + failed + aborted || usageKnownRuns > requests) return null;
  return { runIds: [...item.runIds] as string[], requests, committed, stale, failed, aborted, usageKnownRuns, usage };
}

function uniqueById<T extends { id: string }>(items: readonly T[]): boolean {
  return new Set(items.map((item) => item.id)).size === items.length;
}

export function parseEntry(value: unknown): BtwEntry | null {
  const item = object(value);
  if (
    !item ||
    !onlyKeys(item, ["id", "question", "attempts", "promotedAttemptIds"]) ||
    !id(item.id) ||
    !string(item.question) ||
    !Array.isArray(item.attempts)
  ) {
    return null;
  }

  const attempts = item.attempts.map(parseAttempt);
  if (!attempts.length || attempts.some((attempt) => !attempt)) return null;
  const parsedAttempts = attempts as BtwAttempt[];
  if (!uniqueById(parsedAttempts)) return null;

  const promotedAttemptIds = item.promotedAttemptIds === undefined
    ? undefined
    : parseStringList(item.promotedAttemptIds);
  if (
    (item.promotedAttemptIds !== undefined && !promotedAttemptIds) ||
    (promotedAttemptIds &&
      (!uniqueById(promotedAttemptIds.map((promotedId) => ({ id: promotedId }))) ||
        !promotedAttemptIds.every((promotedId) => parsedAttempts.some((attempt) => attempt.id === promotedId))))
  ) {
    return null;
  }

  return {
    id: item.id,
    question: item.question,
    attempts: parsedAttempts,
    ...(promotedAttemptIds ? { promotedAttemptIds } : {}),
  };
}

function parseThreadWith<TEntry extends { id: string }>(
  value: unknown,
  parseEntryValue: (value: unknown) => TEntry | null,
): { id: string; createdAt: string; entries: TEntry[] } | null {
  const item = object(value);
  if (
    !item ||
    !onlyKeys(item, ["id", "createdAt", "entries"]) ||
    !id(item.id) ||
    !string(item.createdAt) ||
    !Array.isArray(item.entries)
  ) {
    return null;
  }
  const entries = item.entries.map(parseEntryValue);
  if (entries.some((entry) => !entry)) return null;
  const parsedEntries = entries as TEntry[];
  return uniqueById(parsedEntries) ? { id: item.id, createdAt: item.createdAt, entries: parsedEntries } : null;
}

function parseFlatThread(value: unknown): FlatThread | null {
  return parseThreadWith(value, parseFlatEntry);
}

export function parseThreadV4(value: unknown): BtwThreadV4 | null {
  return parseThreadWith(value, parseEntry);
}

export function parseThread(value: unknown): BtwThread | null {
  const item = object(value);
  if (!item || !onlyKeys(item, ["id", "createdAt", "entries", "summary", "summaryMeter"]) || !id(item.id) || !string(item.createdAt) || !Array.isArray(item.entries)) return null;
  const entries = item.entries.map(parseEntry);
  if (entries.some((entry) => !entry) || !uniqueById(entries as BtwEntry[])) return null;
  const parsedEntries = entries as BtwEntry[];
  const summary = item.summary === undefined ? undefined : parseSummary(item.summary);
  const summaryMeter = item.summaryMeter === undefined ? undefined : parseSummaryMeter(item.summaryMeter);
  if ((item.summary !== undefined && !summary) || (item.summaryMeter !== undefined && !summaryMeter)) return null;
  const parsed = { id: item.id, createdAt: item.createdAt, entries: parsedEntries, ...(summary ? { summary } : {}), ...(summaryMeter ? { summaryMeter } : {}) };
  return summary && !validateSummary(parsed) ? null : parsed;
}

function parseThreads<TThread extends { id: string }>(
  value: unknown,
  parseThreadValue: (value: unknown) => TThread | null,
): TThread[] | null {
  if (!Array.isArray(value)) return null;
  const threads = value.map(parseThreadValue);
  if (threads.some((thread) => !thread)) return null;
  const parsedThreads = threads as TThread[];
  return uniqueById(parsedThreads) ? parsedThreads : null;
}

export function parseState(value: unknown, version: 1): BtwStateV1 | null;
export function parseState(value: unknown, version: 2): BtwStateV2 | null;
export function parseState(value: unknown, version: 1 | 2): BtwStateV1 | BtwStateV2 | null {
  const item = object(value);
  const allowed = version === 1
    ? ["version", "threads", "activeThreadId"]
    : ["version", "revision", "threads", "activeThreadId"];
  if (!item || !onlyKeys(item, allowed) || item.version !== version) return null;

  const threads = parseThreads(item.threads, parseFlatThread);
  if (
    !threads ||
    (item.activeThreadId !== undefined &&
      (!id(item.activeThreadId) || !threads.some((thread) => thread.id === item.activeThreadId)))
  ) {
    return null;
  }
  const active = item.activeThreadId === undefined ? {} : { activeThreadId: item.activeThreadId };
  if (version === 1) return { version: 1, threads, ...active };
  if (!Number.isSafeInteger(item.revision) || (item.revision as number) < 0) return null;
  return { version: 2, revision: item.revision as number, threads, ...active };
}

function parseScopeWith<TThread extends { id: string }>(
  value: unknown,
  parseThreadValue: (value: unknown) => TThread | null,
): ScopeBase<TThread> | null {
  const item = object(value);
  if (
    !item ||
    !onlyKeys(item, ["id", "sessionId", "startedAtLeafId", "kind", "createdAt", "threads", "activeThreadId"]) ||
    !id(item.id) ||
    (item.sessionId !== null && !string(item.sessionId)) ||
    (item.startedAtLeafId !== null && !string(item.startedAtLeafId)) ||
    !string(item.createdAt) ||
    (item.kind !== "session" && item.kind !== "tree" && item.kind !== "legacy")
  ) {
    return null;
  }
  const threads = parseThreads(item.threads, parseThreadValue);
  if (
    !threads ||
    (item.activeThreadId !== undefined &&
      (!id(item.activeThreadId) || !threads.some((thread) => thread.id === item.activeThreadId)))
  ) {
    return null;
  }
  return {
    id: item.id,
    sessionId: item.sessionId,
    startedAtLeafId: item.startedAtLeafId,
    kind: item.kind,
    createdAt: item.createdAt,
    threads,
    ...(item.activeThreadId === undefined ? {} : { activeThreadId: item.activeThreadId }),
  };
}

export function parseScopeV3(value: unknown): BtwScopeV3 | null {
  return parseScopeWith(value, parseFlatThread);
}

export function parseScopeV4(value: unknown): BtwScopeV4 | null {
  return parseScopeWith(value, parseThreadV4);
}

export function parseScope(value: unknown): BtwScope | null {
  return parseScopeWith(value, parseThread);
}

function parseScopedState<TScope extends { id: string }>(
  value: unknown,
  version: 3 | 4 | 5,
  parseScopeValue: (value: unknown) => TScope | null,
): { version: 3 | 4 | 5; revision: number; scopes: TScope[] } | null {
  const item = object(value);
  if (
    !item ||
    !onlyKeys(item, ["version", "revision", "scopes"]) ||
    item.version !== version ||
    !Number.isSafeInteger(item.revision) ||
    (item.revision as number) < 0 ||
    !Array.isArray(item.scopes)
  ) {
    return null;
  }
  const scopes = item.scopes.map(parseScopeValue);
  if (scopes.some((scope) => !scope)) return null;
  const parsedScopes = scopes as TScope[];
  return uniqueById(parsedScopes)
    ? { version, revision: item.revision as number, scopes: parsedScopes }
    : null;
}

export function parseStateV3(value: unknown): BtwStateV3 | null {
  const parsed = parseScopedState(value, 3, parseScopeV3);
  return parsed ? { version: 3, revision: parsed.revision, scopes: parsed.scopes } : null;
}

export function parseStateV4(value: unknown): BtwStateV4 | null {
  const parsed = parseScopedState(value, 4, parseScopeV4);
  return parsed ? { version: 4, revision: parsed.revision, scopes: parsed.scopes } : null;
}

export function parseStateV5(value: unknown): BtwState | null {
  const parsed = parseScopedState(value, 5, parseScope);
  return parsed ? { version: 5, revision: parsed.revision, scopes: parsed.scopes } : null;
}

export function loadStateFile(path: string): StoreFileResult {
  if (!existsSync(path)) return { kind: "missing" };
  try {
    const value: unknown = JSON.parse(new TextDecoder("utf-8", { fatal: true }).decode(readFileSync(path)));
    const state = object(value);
    if (!state || !Number.isSafeInteger(state.version)) return { kind: "corrupt" };
    if ((state.version as number) > 5) return { kind: "future", version: state.version as number };
    if (state.version === 1) {
      const parsed = parseState(state, 1);
      return parsed ? { kind: "v1", state: parsed } : { kind: "corrupt" };
    }
    if (state.version === 2) {
      const parsed = parseState(state, 2);
      return parsed ? { kind: "v2", state: parsed } : { kind: "corrupt" };
    }
    if (state.version === 3) {
      const parsed = parseStateV3(state);
      return parsed ? { kind: "v3", state: parsed } : { kind: "corrupt" };
    }
    if (state.version === 4) {
      const parsed = parseStateV4(state);
      return parsed ? { kind: "v4", state: parsed } : { kind: "corrupt" };
    }
    if (state.version === 5) {
      const parsed = parseStateV5(state);
      return parsed ? { kind: "v5", state: parsed } : { kind: "corrupt" };
    }
    return { kind: "corrupt" };
  } catch {
    return { kind: "corrupt" };
  }
}