import http from "node:http";
import { AsyncLocalStorage } from "node:async_hooks";
import { randomUUID } from "node:crypto";
import { create } from "@bufbuild/protobuf";
import { createClient, type Client } from "@connectrpc/connect";
import {
  connectNodeAdapter,
  createConnectTransport,
} from "@connectrpc/connect-node";
import {
  CaminoService,
  CrdtValueSchema,
  ListValueSchema,
  NullValueSchema,
  ObjectValueSchema,
  RefValueSchema,
  ValueSchema,
  type Value,
} from "./camino/api_pb.js";
import {
  HandshakeResponseSchema,
  InvokeResponseSchema,
  PackageRuntime,
  type InvokeRequest,
} from "./quixos/runtime_pb.js";
import { FunctionRefSchema } from "./quixos/refs_pb.js";

declare const objectRefBrand: unique symbol;

export type ObjectRef<ClassId extends string> = string & {
  readonly [objectRefBrand]: ClassId;
};

export type CaminoClient = Client<typeof CaminoService>;

export type { InvokeRequest, Value };

export type DerivedDependency =
  | {
      kind: "object";
      objectId: string;
    }
  | {
      kind: "field";
      objectId: string;
      fieldName: string;
    }
  | {
      kind: "edges";
      objectId: string;
      sourceField?: string;
    };

type DerivedDependencyTracker = {
  record(dependency: DerivedDependency): void;
  dependencies(): DerivedDependency[];
};

const dependencyKey = (dependency: DerivedDependency) => {
  switch (dependency.kind) {
    case "object":
      return `object:${dependency.objectId}`;
    case "field":
      return `field:${dependency.objectId}:${dependency.fieldName}`;
    case "edges":
      return `edges:${dependency.objectId}:${dependency.sourceField ?? ""}`;
  }
};

const createDependencyTracker = (): DerivedDependencyTracker => {
  const dependencies = new Map<string, DerivedDependency>();
  return {
    record(dependency) {
      dependencies.set(dependencyKey(dependency), dependency);
    },
    dependencies() {
      return [...dependencies.values()];
    },
  };
};

const dependencyStorage = new AsyncLocalStorage<DerivedDependencyTracker>();

export const recordDerivedDependency = (dependency: DerivedDependency) => {
  dependencyStorage.getStore()?.record(dependency);
};

export const getActiveDerivedDependencies = (): DerivedDependency[] =>
  dependencyStorage.getStore()?.dependencies() ?? [];

export const runWithDerivedDependencyTracking = async <T>(
  fn: () => Promise<T> | T,
): Promise<{ value: T; dependencies: DerivedDependency[] }> => {
  const tracker = createDependencyTracker();
  const value = await dependencyStorage.run(tracker, fn);
  return {
    value,
    dependencies: tracker.dependencies(),
  };
};

const logDerivedDependencies = (dependencies: DerivedDependency[]) => {
  if (process.env.QUIXOS_DERIVED_DEPS_LOG !== "1") {
    return;
  }
  process.stderr.write(
    `[derived-deps] ${JSON.stringify({ dependencies })}\n`,
  );
};

const logDerivedWatch = (event: string, detail: Record<string, unknown>) => {
  if (process.env.QUIXOS_DERIVED_WATCH_LOG !== "1") {
    return;
  }
  process.stderr.write(
    `[derived-watch] ${JSON.stringify({ event, ...detail })}\n`,
  );
};

type WatchableContext = {
  camino: CaminoClient;
  request: InvokeRequest;
};

type ActiveDerivedWatch<Context extends WatchableContext> = {
  watchId: string;
  context: Context;
  get: RuntimeHandler<Context>;
  controllers: Map<string, AbortController>;
  running: boolean;
  pending: boolean;
  stopped: boolean;
  lastValue: unknown;
};

const dependencyObjectIds = (dependencies: DerivedDependency[]) =>
  new Set(dependencies.map((dependency) => dependency.objectId));

const stopDerivedWatch = <Context extends WatchableContext>(
  watch: ActiveDerivedWatch<Context>,
) => {
  if (watch.stopped) {
    return;
  }
  watch.stopped = true;
  for (const controller of watch.controllers.values()) {
    controller.abort();
  }
  watch.controllers.clear();
  logDerivedWatch("stopped", { watchId: watch.watchId });
};

export type Field<T> = {
  get(): Promise<T | undefined>;
  set(value: T): Promise<void>;
};

export type RuntimeHandler<Context> = (
  context: Context,
) => Promise<unknown> | unknown;

export type FieldOperation<Context> = {
  get?: RuntimeHandler<Context>;
  set?: RuntimeHandler<Context>;
  watchStart?: RuntimeHandler<Context>;
  watchStop?: RuntimeHandler<Context>;
};

export const fieldOps = <Context>(
  operation: FieldOperation<Context>,
): FieldOperation<Context> => operation;

export const derivedWithDeps = <Context>(operation: {
  get: RuntimeHandler<Context>;
}): FieldOperation<Context> => {
  type DerivedContext = Context & WatchableContext;
  const watches = new Map<string, ActiveDerivedWatch<DerivedContext>>();

  const rerun = async (watch: ActiveDerivedWatch<DerivedContext>) => {
    if (watch.stopped) {
      return;
    }
    if (watch.running) {
      watch.pending = true;
      return;
    }
    watch.running = true;
    try {
      do {
        watch.pending = false;
        const { value, dependencies } = await runWithDerivedDependencyTracking(
          () => watch.get(watch.context),
        );
        watch.lastValue = value;
        logDerivedDependencies(dependencies);
        logDerivedWatch("evaluated", {
          watchId: watch.watchId,
          value,
          dependencies,
        });
        reconcileSubscriptions(watch, dependencies);
      } while (watch.pending && !watch.stopped);
    } catch (error) {
      logDerivedWatch("error", {
        watchId: watch.watchId,
        error: error instanceof Error ? error.message : String(error),
      });
    } finally {
      watch.running = false;
    }
  };

  const scheduleRerun = (
    watch: ActiveDerivedWatch<DerivedContext>,
    reason: string,
  ) => {
    if (watch.stopped) {
      return;
    }
    logDerivedWatch("scheduled", { watchId: watch.watchId, reason });
    void rerun(watch);
  };

  const startObjectWatch = (
    watch: ActiveDerivedWatch<DerivedContext>,
    objectId: string,
  ) => {
    const controller = new AbortController();
    watch.controllers.set(objectId, controller);
    logDerivedWatch("subscribed", { watchId: watch.watchId, objectId });
    void (async () => {
      try {
        for await (const event of watch.context.camino.watchObject(
          {
            objectId,
            includeSnapshot: false,
          },
          { signal: controller.signal },
        )) {
          if (controller.signal.aborted || watch.stopped) {
            break;
          }
          scheduleRerun(
            watch,
            event.op?.id ? `op:${event.op.id}` : `object:${objectId}`,
          );
        }
      } catch (error) {
        if (!controller.signal.aborted && !watch.stopped) {
          logDerivedWatch("watch-error", {
            watchId: watch.watchId,
            objectId,
            error: error instanceof Error ? error.message : String(error),
          });
        }
      } finally {
        if (watch.controllers.get(objectId) === controller) {
          watch.controllers.delete(objectId);
          logDerivedWatch("unsubscribed", { watchId: watch.watchId, objectId });
        }
      }
    })();
  };

  const reconcileSubscriptions = (
    watch: ActiveDerivedWatch<DerivedContext>,
    dependencies: DerivedDependency[],
  ) => {
    const objectIds = dependencyObjectIds(dependencies);
    for (const [objectId, controller] of watch.controllers.entries()) {
      if (!objectIds.has(objectId)) {
        controller.abort();
        watch.controllers.delete(objectId);
        logDerivedWatch("unsubscribed", { watchId: watch.watchId, objectId });
      }
    }
    for (const objectId of objectIds) {
      if (!watch.controllers.has(objectId)) {
        startObjectWatch(watch, objectId);
      }
    }
  };

  return fieldOps({
    async get(context) {
      const { value, dependencies } = await runWithDerivedDependencyTracking(
        () => operation.get(context),
      );
      logDerivedDependencies(dependencies);
      return value;
    },
    async watchStart(context) {
      const watchId = `watch:${randomUUID()}`;
      const watch: ActiveDerivedWatch<DerivedContext> = {
        watchId,
        context: context as DerivedContext,
        get: operation.get as RuntimeHandler<DerivedContext>,
        controllers: new Map(),
        running: false,
        pending: false,
        stopped: false,
        lastValue: undefined,
      };
      watches.set(watchId, watch);
      logDerivedWatch("started", { watchId });
      await rerun(watch);
      return {
        watch_id: watchId,
      };
    },
    watchStop(context) {
      const input = protoFieldsToJs((context as DerivedContext).request.input);
      const watchId =
        typeof input.watch_id === "string"
          ? input.watch_id
          : typeof input.watchId === "string"
            ? input.watchId
            : "";
      const watch = watches.get(watchId);
      if (!watch) {
        throw new Error(`Unknown derived watch: ${watchId}`);
      }
      stopDerivedWatch(watch);
      watches.delete(watchId);
      return {};
    },
  });
};

export type RuntimeFunctionSpec<ExportName extends string = string> = {
  exportName: ExportName;
  symbol: string;
  operation?: string;
};

const isRecord = (value: unknown): value is Record<string, unknown> =>
  Boolean(value) && typeof value === "object" && !Array.isArray(value);

const bytesToBase64 = (value: Uint8Array) =>
  Buffer.from(value).toString("base64");

const base64ToBytes = (value: string) => Buffer.from(value, "base64");

export const jsToProtoValue = (value: unknown): Value => {
  if (value === null || value === undefined) {
    return create(ValueSchema, {
      kind: {
        case: "nullValue",
        value: create(NullValueSchema, {}),
      },
    });
  }
  if (typeof value === "boolean") {
    return create(ValueSchema, {
      kind: { case: "boolValue", value },
    });
  }
  if (typeof value === "number") {
    return create(ValueSchema, {
      kind: { case: "numberValue", value },
    });
  }
  if (typeof value === "bigint") {
    return create(ValueSchema, {
      kind: { case: "integerValue", value: value.toString() },
    });
  }
  if (typeof value === "string") {
    return create(ValueSchema, {
      kind: { case: "stringValue", value },
    });
  }
  if (value instanceof Uint8Array) {
    return create(ValueSchema, {
      kind: { case: "bytesValue", value },
    });
  }
  if (Array.isArray(value)) {
    return create(ValueSchema, {
      kind: {
        case: "listValue",
        value: create(ListValueSchema, {
          values: value.map(jsToProtoValue),
        }),
      },
    });
  }
  if (isRecord(value)) {
    if (typeof value.$caminoRef === "string") {
      return create(ValueSchema, {
        kind: {
          case: "refValue",
          value: create(RefValueSchema, {
            objectId: value.$caminoRef,
          }),
        },
      });
    }
    if (
      typeof value.$caminoBytes === "string" &&
      Object.keys(value).every((key) => key === "$caminoBytes")
    ) {
      return create(ValueSchema, {
        kind: {
          case: "bytesValue",
          value: base64ToBytes(value.$caminoBytes),
        },
      });
    }
    if (
      typeof value.$caminoCrdtType === "string" &&
      typeof value.$caminoCrdtPayload === "string"
    ) {
      return create(ValueSchema, {
        kind: {
          case: "crdtValue",
          value: create(CrdtValueSchema, {
            type: value.$caminoCrdtType,
            encoding:
              typeof value.$caminoCrdtEncoding === "string"
                ? value.$caminoCrdtEncoding
                : "base64",
            payload: base64ToBytes(value.$caminoCrdtPayload),
          }),
        },
      });
    }
    return create(ValueSchema, {
      kind: {
        case: "objectValue",
        value: create(ObjectValueSchema, {
          fields: jsObjectToProtoFields(value),
        }),
      },
    });
  }
  throw new Error(`Unsupported Camino RPC value: ${typeof value}`);
};

export const protoValueToJs = (value: Value | undefined): unknown => {
  switch (value?.kind.case) {
    case "nullValue":
    case undefined:
      return null;
    case "boolValue":
    case "numberValue":
    case "stringValue":
      return value.kind.value;
    case "integerValue":
      return value.kind.value;
    case "bytesValue":
      return bytesToBase64(value.kind.value);
    case "listValue":
      return value.kind.value.values.map(protoValueToJs);
    case "objectValue":
      return protoFieldsToJs(value.kind.value.fields);
    case "refValue":
      return value.kind.value.objectId;
    case "crdtValue":
      return {
        $caminoCrdtType: value.kind.value.type,
        $caminoCrdtEncoding: value.kind.value.encoding,
        $caminoCrdtPayload: bytesToBase64(value.kind.value.payload),
      };
  }
};

export const jsObjectToProtoFields = (value: Record<string, unknown>) =>
  Object.fromEntries(
    Object.entries(value)
      .filter(([, fieldValue]) => fieldValue !== undefined)
      .map(([key, fieldValue]) => [key, jsToProtoValue(fieldValue)]),
  );

export const protoFieldsToJs = (
  fields: Record<string, Value> | undefined,
): Record<string, unknown> =>
  Object.fromEntries(
    Object.entries(fields ?? {}).map(([key, value]) => [
      key,
      protoValueToJs(value),
    ]),
  );

export const objectRef = <ClassId extends string>(objectId: string) =>
  objectId as ObjectRef<ClassId>;

export const createField = <T>(
  camino: CaminoClient,
  objectId: string,
  fieldName: string,
): Field<T> => ({
  async get() {
    recordDerivedDependency({ kind: "field", objectId, fieldName });
    const response = await camino.getObject({ objectId });
    return protoValueToJs(response.object?.fields[fieldName]) as T | undefined;
  },
  async set(value) {
    await camino.setField({
      objectId,
      fieldName,
      value: jsToProtoValue(value),
    });
  },
});

const createDependencyTrackingCaminoClient = (
  camino: CaminoClient,
): CaminoClient =>
  new Proxy(camino, {
    get(target, property, receiver) {
      if (property === "getObject") {
        return async (request: { objectId: string }, ...args: unknown[]) => {
          if (request.objectId) {
            recordDerivedDependency({
              kind: "object",
              objectId: request.objectId,
            });
          }
          return await (
            Reflect.get(target, property, receiver) as (
              request: { objectId: string },
              ...args: unknown[]
            ) => Promise<unknown>
          )(request, ...args);
        };
      }
      if (property === "listEdges") {
        return async (
          request: { objectId: string; sourceField?: string },
          ...args: unknown[]
        ) => {
          if (request.objectId) {
            recordDerivedDependency({
              kind: "edges",
              objectId: request.objectId,
              ...(request.sourceField ? { sourceField: request.sourceField } : {}),
            });
          }
          return await (
            Reflect.get(target, property, receiver) as (
              request: { objectId: string; sourceField?: string },
              ...args: unknown[]
            ) => Promise<unknown>
          )(request, ...args);
        };
      }
      return Reflect.get(target, property, receiver);
    },
  });

export const serveQuixosPackageRuntime = <
  Context,
  Spec extends readonly RuntimeFunctionSpec[],
>(params: {
  packageNamespace: string;
  packageName: string;
  runtimeProtocolVersion: string;
  functions: Spec;
  implementation: {
    [Name in Spec[number]["exportName"]]:
      | RuntimeHandler<Context>
      | FieldOperation<Context>;
  };
  createContext: (camino: CaminoClient, request: InvokeRequest) => Context;
}) => {
  const host = process.env.QUIXOS_RUNTIME_HOST ?? "127.0.0.1";
  const port = Number(process.env.QUIXOS_RUNTIME_PORT ?? "0");
  const caminoUrl = process.env.CAMINO_URL ?? "http://127.0.0.1:7310";
  const camino = createClient(
    CaminoService,
    createConnectTransport({ baseUrl: caminoUrl, httpVersion: "1.1" }),
  );
  const trackingCamino = createDependencyTrackingCaminoClient(camino);

  if (!Number.isInteger(port) || port <= 0) {
    throw new Error("QUIXOS_RUNTIME_PORT must be a positive integer");
  }

  const functionRefs = params.functions.map((entry) =>
    create(FunctionRefSchema, {
      packageNamespace: params.packageNamespace,
      packageName: params.packageName,
      symbol: entry.symbol,
      operation: entry.operation ?? "",
    }),
  );

  const handler = connectNodeAdapter({
    routes: (router) => {
      router.service(PackageRuntime, {
        async handshake() {
          return create(HandshakeResponseSchema, {
            packageNamespace: params.packageNamespace,
            packageName: params.packageName,
            runtimeProtocolVersion: params.runtimeProtocolVersion,
            functions: functionRefs,
          });
        },

        async invoke(request) {
          try {
            const symbol = request.function?.symbol ?? "";
            const spec = params.functions.find(
              (entry) =>
                entry.symbol === symbol &&
                (entry.operation ?? "") === (request.function?.operation ?? ""),
            );
            if (!spec) {
              throw new Error(
                `Unknown ${params.packageName} function: ${symbol}${
                  request.function?.operation
                    ? `.${request.function.operation}`
                    : ""
                }`,
              );
            }
            const implementationExport = (
              params.implementation as Record<
                string,
                RuntimeHandler<Context> | FieldOperation<Context>
              >
            )[spec.exportName];
            const runtimeFunction =
              spec.operation && typeof implementationExport === "object"
                ? implementationExport[spec.operation as keyof FieldOperation<Context>]
                : implementationExport;
            if (!runtimeFunction) {
              throw new Error(
                `Missing ${params.packageName} implementation export: ${spec.exportName}${
                  spec.operation ? `.${spec.operation}` : ""
                }`,
              );
            }
            if (typeof runtimeFunction !== "function") {
              throw new Error(
                `${params.packageName} implementation export is not callable: ${spec.exportName}${
                  spec.operation ? `.${spec.operation}` : ""
                }`,
              );
            }
            return create(InvokeResponseSchema, {
              ok: true,
              result: jsToProtoValue(
                await runtimeFunction(params.createContext(trackingCamino, request)),
              ),
            });
          } catch (error) {
            return create(InvokeResponseSchema, {
              ok: false,
              result: jsToProtoValue(null),
              error: error instanceof Error ? error.message : String(error),
            });
          }
        },
      });
    },
  });

  const server = http.createServer((request, response) => {
    void handler(request, response);
  });
  const close = () => server.close(() => process.exit(0));

  server.listen(port, host, () => {
    process.stderr.write(
      `${params.packageName} listening on ${host}:${port}\n`,
    );
  });
  process.on("SIGTERM", close);
  process.on("SIGINT", close);
};