import type {
  StoreActionContext,
  StoreDefinition,
  StoreInstance,
  StoreInstanceCore,
  StoreMutation,
  StorePersistenceConfig,
  StoreCombinedState,
  StoreFunction,
} from "./types.ts";
export type * from "./types.ts";

const memoryPersistence = new Map<string, string>();

function clone<T>(value: T): T {
  if (typeof structuredClone === "function") {
    try {
      return structuredClone(value);
    } catch {
      /* Proxies are cloned through their JSON-visible state below. */
    }
  }
  return JSON.parse(JSON.stringify(value)) as T;
}

function readonlySnapshot<S extends object>(state: S): Readonly<S> {
  return Object.freeze(clone(state));
}

function storageFor(kind: StorePersistenceConfig<object>["storage"]): Storage | null {
  if (typeof window === "undefined") return null;
  if (kind === "local") return window.localStorage;
  if (kind === "session") return window.sessionStorage;
  return null;
}

function readPersisted<S extends object>(
  key: string,
  config: StorePersistenceConfig<S>,
): Partial<S> | null {
  try {
    const raw =
      config.storage === "memory"
        ? memoryPersistence.get(key)
        : storageFor(config.storage)?.getItem(key);
    if (!raw) return null;
    const parsed = JSON.parse(raw) as { version?: number; state?: unknown };
    const from = Number(parsed.version ?? 0);
    let state = parsed.state;
    if (from !== config.version && config.migrate)
      state = config.migrate(state, from, config.version);
    if (config.validate) return config.validate(state);
    if (!state || typeof state !== "object" || Array.isArray(state)) return null;
    return state as Partial<S>;
  } catch {
    return null;
  }
}

function writePersisted<S extends object>(
  key: string,
  state: S,
  config: StorePersistenceConfig<S>,
): void {
  const picked: Record<string, unknown> = {};
  for (const name of config.include) picked[name] = state[name];
  const raw = JSON.stringify({ version: config.version, state: picked });
  if (config.storage === "memory") memoryPersistence.set(key, raw);
  else storageFor(config.storage)?.setItem(key, raw);
}

export function defineStore<
  S extends object,
  C extends object = Record<string, never>,
  A extends Record<string, StoreFunction> = Record<string, StoreFunction>,
  CS extends object = Record<string, never>,
  SS extends object = Record<string, never>,
>(definition: StoreDefinition<S, C, A, CS, SS>): StoreDefinition<S, C, A, CS, SS> {
  return definition;
}

export interface StoreContainerOptions {
  runtime: "server" | "client";
  request?: unknown;
  routeId?: string;
  hydration?: Record<string, unknown>;
  onMutation?: (mutation: StoreMutation) => void;
}

export class StoreContainer {
  readonly runtime: "server" | "client";
  readonly request?: unknown;
  readonly routeId?: string;
  private readonly instances = new Map<
    string,
    StoreInstance<Record<string, unknown>, Record<string, unknown>, Record<string, StoreFunction>>
  >();
  private readonly hydration: Record<string, unknown>;
  private readonly onMutation?: (mutation: StoreMutation) => void;
  private readonly lastMutations = new Map<string, StoreMutation>();

  constructor(options: StoreContainerOptions) {
    this.runtime = options.runtime;
    this.request = options.request;
    this.routeId = options.routeId;
    this.hydration = options.hydration ?? {};
    this.onMutation = options.onMutation;
  }

  async use<
    S extends object,
    C extends object,
    A extends Record<string, StoreFunction>,
    CS extends object,
    SS extends object,
  >(
    definition: StoreDefinition<S, C, A, CS, SS>,
  ): Promise<StoreInstance<StoreCombinedState<S, CS, SS>, C, A>> {
    const key =
      definition.kind === "page"
        ? `${definition.name}@${this.routeId ?? "default"}`
        : definition.name;
    const existing = this.instances.get(key);
    if (existing) return existing as StoreInstance<StoreCombinedState<S, CS, SS>, C, A>;
    const instance = createStoreInstance(definition, {
      runtime: this.runtime,
      request: this.request,
      routeId: this.routeId,
      hydration: this.hydration[definition.name],
      onMutation: (mutation) => {
        this.lastMutations.set(definition.name, mutation);
        this.onMutation?.(mutation);
      },
    });
    this.instances.set(key, instance);
    await instance.whenReady;
    return instance;
  }

  serialize(): Record<string, unknown> {
    return Object.fromEntries(
      Array.from(this.instances.values(), (instance) => [instance.name, instance.serialize()]),
    );
  }

  inspect(): Array<{
    name: string;
    kind: "global" | "page";
    state: Readonly<Record<string, unknown>>;
    computed: Readonly<Record<string, unknown>>;
    lastAction?: string;
    changed: string[];
    hydrationSource: "server" | "persistence" | "initial";
  }> {
    return Array.from(this.instances.values(), (instance) => {
      const mutation = this.lastMutations.get(instance.name);
      return {
        name: instance.name,
        kind: instance.kind,
        state: instance.snapshot() as Readonly<Record<string, unknown>>,
        computed: Object.freeze(
          Object.fromEntries(
            Reflect.ownKeys(instance.computed).map((key) => [
              String(key),
              Reflect.get(instance.computed, key),
            ]),
          ),
        ),
        ...(mutation ? { lastAction: mutation.action } : {}),
        changed: mutation?.changed ?? [],
        hydrationSource: this.hydration[instance.name] ? "server" : "initial",
      };
    });
  }

  async hotUpdate<
    S extends object,
    C extends object,
    A extends Record<string, StoreFunction>,
    CS extends object,
    SS extends object,
  >(
    definition: StoreDefinition<S, C, A, CS, SS>,
  ): Promise<{ preserved: string[]; reset: string[] }> {
    const entries = [...this.instances.entries()].filter(
      ([, instance]) => instance.name === definition.name,
    );
    const preserved = new Set<string>();
    const reset = new Set<string>();
    for (const [key, previous] of entries) {
      const snapshot = previous.snapshot() as Record<string, unknown>;
      await previous.dispose();
      const freshShared = definition.createSharedState();
      const freshRuntime =
        this.runtime === "client"
          ? (definition.createClientState?.() ?? {})
          : (definition.createServerState?.() ?? {});
      const nextShape = { ...freshShared, ...freshRuntime } as Record<string, unknown>;
      const compatible: Record<string, unknown> = {};
      for (const [name, value] of Object.entries(snapshot)) {
        if (!(name in nextShape)) continue;
        const expected = nextShape[name];
        const same =
          expected === null || value === null
            ? expected === value || expected === null
            : Array.isArray(expected)
              ? Array.isArray(value)
              : typeof expected === typeof value;
        if (same) {
          compatible[name] = value;
          preserved.add(name);
        } else reset.add(name);
      }
      for (const name of Object.keys(nextShape))
        if (!(name in compatible) && name in snapshot) reset.add(name);
      const instance = createStoreInstance(definition, {
        runtime: this.runtime,
        request: this.request,
        routeId: this.routeId,
        hydration: compatible,
        onMutation: (mutation) => {
          this.lastMutations.set(definition.name, mutation);
          this.onMutation?.(mutation);
        },
      });
      this.instances.set(key, instance);
      await instance.whenReady;
    }
    return { preserved: [...preserved], reset: [...reset] };
  }

  async disposePageStores(): Promise<void> {
    for (const [key, instance] of [...this.instances]) {
      if (instance.kind !== "page") continue;
      await instance.dispose();
      this.instances.delete(key);
    }
  }

  async dispose(): Promise<void> {
    for (const instance of this.instances.values()) await instance.dispose();
    this.instances.clear();
  }
}

export function createStoreContainer(options: StoreContainerOptions): StoreContainer {
  return new StoreContainer(options);
}

export function createStoreInstance<
  S extends object,
  C extends object,
  A extends Record<string, StoreFunction>,
  CS extends object,
  SS extends object,
>(
  definition: StoreDefinition<S, C, A, CS, SS>,
  options: Omit<StoreContainerOptions, "hydration"> & { hydration?: unknown },
): StoreInstance<StoreCombinedState<S, CS, SS>, C, A> {
  type State = StoreCombinedState<S, CS, SS>;
  const createInitialState = (): State =>
    ({
      ...definition.createSharedState(),
      ...(options.runtime === "client"
        ? definition.createClientState?.()
        : definition.createServerState?.()),
    }) as State;
  const initial = createInitialState();
  const persisted =
    options.runtime === "client" && definition.persist
      ? readPersisted(`wrnexus:store:${definition.name}`, definition.persist)
      : null;
  const hydrated =
    options.hydration && typeof options.hydration === "object"
      ? (options.hydration as Partial<State>)
      : null;
  const raw = Object.assign(initial, persisted ?? {}, hydrated ?? {});
  const listeners = new Set<(snapshot: Readonly<State>, mutation?: StoreMutation) => void>();
  let currentAction = "direct";
  let internalMutation = false;

  const mutableState = new Proxy(raw, {
    set(target, property, value) {
      if (!internalMutation) {
        throw new TypeError(
          `WRN-STORE-READONLY: ${definition.name}.${String(property)} must be changed by a store action.`,
        );
      }
      if (Object.is(Reflect.get(target, property), value)) return true;
      const before = readonlySnapshot(target);
      Reflect.set(target, property, value);
      const after = readonlySnapshot(target);
      const mutation: StoreMutation = {
        store: definition.name,
        action: currentAction,
        changed: [String(property)],
        before: before as Readonly<Record<string, unknown>>,
        after: after as Readonly<Record<string, unknown>>,
        timestamp: Date.now(),
      };
      if (definition.persist && options.runtime === "client") {
        writePersisted(`wrnexus:store:${definition.name}`, target, definition.persist);
      }
      options.onMutation?.(mutation);
      for (const listener of listeners) listener(after, mutation);
      return true;
    },
    deleteProperty(target, property) {
      if (!internalMutation) {
        throw new TypeError(
          `WRN-STORE-READONLY: ${definition.name}.${String(property)} must be changed by a store action.`,
        );
      }
      return Reflect.deleteProperty(target, property);
    },
  });

  const publicState = new Proxy({} as State, {
    get(_target, property) {
      return Reflect.get(mutableState, property);
    },
    set(_target, property) {
      throw new TypeError(
        `WRN-STORE-READONLY: ${definition.name}.${String(property)} must be changed by a store action.`,
      );
    },
    deleteProperty(_target, property) {
      throw new TypeError(
        `WRN-STORE-READONLY: ${definition.name}.${String(property)} must be changed by a store action.`,
      );
    },
    ownKeys() {
      return Reflect.ownKeys(mutableState);
    },
    getOwnPropertyDescriptor(_target, property) {
      if (!(property in mutableState)) return undefined;
      return {
        enumerable: true,
        configurable: true,
        value: Reflect.get(mutableState, property),
        writable: false,
      };
    },
    has(_target, property) {
      return property in mutableState;
    },
  });

  const mutate = <T>(action: string, callback: () => T): T => {
    const previousAction = currentAction;
    const previousMutation = internalMutation;
    currentAction = action;
    internalMutation = true;
    try {
      return callback();
    } finally {
      currentAction = previousAction;
      internalMutation = previousMutation;
    }
  };

  const reset = () =>
    mutate("$reset", () => {
      const next = createInitialState();
      for (const key of Object.keys(mutableState)) {
        if (!(key in next)) Reflect.deleteProperty(mutableState, key);
      }
      Object.assign(mutableState, next);
    });

  const context: StoreActionContext<State> = {
    state: mutableState,
    snapshot: () => readonlySnapshot(mutableState),
    reset,
    runtime: options.runtime,
    request: options.request,
    routeId: options.routeId,
  };

  const actions = {} as A;
  for (const [name, rawDefinitions] of Object.entries(definition.actions ?? {})) {
    const definitions = Array.isArray(rawDefinitions) ? rawDefinitions : [rawDefinitions];
    const selected =
      definitions.find((entry) => entry.runtime === options.runtime) ??
      definitions.find((entry) => entry.runtime === "shared") ??
      definitions.find((entry) => entry.runtime === "legacy");
    if (!selected) continue;
    Reflect.set(actions, name, async (...args: unknown[]) => {
      currentAction = name;
      internalMutation = true;
      try {
        const handler = selected.handler as unknown as (
          context: StoreActionContext<State>,
          ...args: unknown[]
        ) => unknown;
        return await handler(context, ...args);
      } finally {
        currentAction = "direct";
        internalMutation = false;
      }
    });
  }

  const computed = new Proxy({} as C, {
    get(_target, property) {
      const fn = definition.computed?.[property as keyof C];
      return fn ? fn(publicState) : undefined;
    },
    set() {
      throw new TypeError("Computed store values are readonly");
    },
    ownKeys() {
      return Reflect.ownKeys(definition.computed ?? {});
    },
    getOwnPropertyDescriptor() {
      return { enumerable: true, configurable: true };
    },
  });

  const lifecycleContext = {
    state: mutableState,
    runtime: options.runtime,
    request: options.request,
    routeId: options.routeId,
  };
  const runLifecycle = async (
    name: string,
    hook?: (context: typeof lifecycleContext) => void | Promise<void>,
  ): Promise<void> => {
    if (!hook) return;
    const previousAction = currentAction;
    const previousMutation = internalMutation;
    currentAction = name;
    internalMutation = true;
    try {
      await hook(lifecycleContext);
    } finally {
      currentAction = previousAction;
      internalMutation = previousMutation;
    }
  };
  const initHook =
    options.runtime === "server"
      ? definition.lifecycle?.serverInit
      : definition.lifecycle?.clientInit;
  const ready = runLifecycle(
    options.runtime === "server" ? "$serverInit" : "$clientInit",
    initHook,
  ).then(async () => {
    if (options.runtime === "client" && options.hydration && definition.lifecycle?.hydrate) {
      await runLifecycle("$hydrate", definition.lifecycle.hydrate);
    }
  });

  const core: StoreInstanceCore<State, C, A> = {
    name: definition.name,
    kind: definition.kind,
    state: publicState,
    computed,
    actions,
    whenReady: ready,
    reset,
    snapshot: () => readonlySnapshot(mutableState),
    hydrate(value) {
      mutate("$hydrate", () => Object.assign(mutableState, value));
    },
    serialize() {
      const result: Partial<State> = {};
      const serverKeys = new Set(Object.keys(definition.createServerState?.() ?? {}));
      for (const [key, value] of Object.entries(mutableState)) {
        if (!serverKeys.has(key)) Reflect.set(result, key, clone(value));
      }
      return result;
    },
    subscribe(listener) {
      listeners.add(listener);
      return () => listeners.delete(listener);
    },
    async dispose() {
      await ready;
      await runLifecycle("$dispose", definition.lifecycle?.dispose);
      listeners.clear();
    },
  };

  return new Proxy(core as StoreInstance<State, C, A>, {
    get(target, property, receiver) {
      if (Reflect.has(target, property)) return Reflect.get(target, property, receiver);
      if (property in actions) return Reflect.get(actions, property);
      if (property in (definition.computed ?? {})) return Reflect.get(computed, property);
      if (property in mutableState) return Reflect.get(publicState, property);
      return undefined;
    },
    set(_target, property) {
      throw new TypeError(
        `WRN-STORE-READONLY: ${definition.name}.${String(property)} must be changed by a store action.`,
      );
    },
    ownKeys(target) {
      return [
        ...new Set([
          ...Reflect.ownKeys(target),
          ...Reflect.ownKeys(mutableState),
          ...Reflect.ownKeys(definition.computed ?? {}),
          ...Reflect.ownKeys(actions),
        ]),
      ];
    },
    getOwnPropertyDescriptor(target, property) {
      return (
        Reflect.getOwnPropertyDescriptor(target, property) ?? {
          enumerable: true,
          configurable: true,
        }
      );
    },
  });
}