import {
  TypeSubstitution,
  GenericTypeError,
  instantiateInterface,
  appliedInterfaceId,
  valueType,
  type AtomId,
  type ClosedTypeArgument,
  type GenericTypeEnvironment,
  type InterfaceApplicationExpression,
  type InterfaceRevision,
  type ObjectTypeExpression,
  type TypeArgumentExpression,
  type TypeParameter,
  type ValueAliasDefinition,
  type ValueType,
  type ValueTypeExpression,
} from "../capability-model/index.js";
import type {
  InterfaceTypeContext,
  TargetConstraintContext,
  TypeArgumentsContext,
  TypeParametersContext,
  TypeAliasDeclContext,
  ValueTypeContext,
} from "./generated/QuixosCapabilityParser.js";

const literal = (context: { getText(): string }) => JSON.parse(context.getText()) as string;

/** Lexical authoring scope; only `value`/`target` return installed types. */
export class GenericSourceTypes {
  readonly interfaces = new Map<string, InterfaceRevision>();
  readonly definitions = new Map<string, InterfaceRevision>();
  readonly applications = new Map<string, InterfaceRevision>();
  readonly aliases = new Map<string, ValueAliasDefinition>();
  parameters = new Map<string, TypeParameter>();
  self?: AtomId;
  private active: string[] = [];
  private applicationCount = 0;

  constructor(readonly atoms: Map<string, AtomId>) {}

  environment(): GenericTypeEnvironment {
    return {
      arguments: new Map(),
      aliases: this.aliases,
      self: this.self,
      applyInterface: (id, args) => this.apply(id, args).revisionId,
      // Obligations are retained on the application and discharged by workspace
      // validation, where all atom conformances are known (not source order).
      implementsInterface: () => true,
    };
  }

  register(name: string, definition: InterfaceRevision) {
    this.interfaces.set(name, definition);
    this.definitions.set(definition.revisionId, definition);
  }

  /** Check authored applications even when nobody has instantiated this template yet. */
  validateTemplate(definition: InterfaceRevision) {
    const template = definition.template;
    if (!template) return;
    const aliases = new Map((template.aliases ?? []).map((alias) => [alias.id, alias]));
    const parametersById = new Map(
      [...template.parameters, ...(template.aliases ?? []).flatMap((alias) => alias.parameters)].map((parameter) => [
        parameter.id,
        parameter,
      ]),
    );
    const replace = (node: unknown, arguments_: Map<string, TypeArgumentExpression>): unknown => {
      if (!node || typeof node !== "object") return node;
      if ("kind" in node && node.kind === "parameter" && "parameterId" in node) {
        const arg = arguments_.get(String(node.parameterId));
        if (arg) return arg.kind === "value" ? arg.type : arg.target;
      }
      if (Array.isArray(node)) return node.map((entry) => replace(entry, arguments_));
      return Object.fromEntries(Object.entries(node).map(([key, value]) => [key, replace(value, arguments_)]));
    };
    const implies = (
      actual: InterfaceApplicationExpression,
      required: InterfaceApplicationExpression,
      seen = new Set<string>(),
    ): boolean => {
      const key = JSON.stringify(actual);
      if (key === JSON.stringify(required)) return true;
      if (seen.has(key) || seen.size > 128) return false;
      seen.add(key);
      const contract = this.definitions.get(actual.definitionId);
      if (!contract?.template)
        return (contract?.requiredInterfaces ?? []).includes(required.definitionId) && required.arguments.length === 0;
      const args = new Map(
        contract.template.parameters.map((parameter, index) => [parameter.id, actual.arguments[index]]),
      );
      return contract.template.requires.some((parent) =>
        implies(replace(parent, args) as InterfaceApplicationExpression, required, seen),
      );
    };
    const storable = (type: ValueTypeExpression, depth = 0): boolean => {
      if (depth > 128)
        throw new GenericTypeError(
          "type-complexity-limit",
          definition.displayName,
          "Storable alias expansion exceeds depth limit",
        );
      if (type.kind === "parameter") {
        const parameter = parametersById.get(type.parameterId);
        return parameter?.kind === "value" && Boolean(parameter.storable);
      }
      if (type.kind === "list" || type.kind === "optional") return storable(type.value, depth + 1);
      if (type.kind === "alias") {
        const alias = aliases.get(type.definitionId);
        if (!alias || alias.parameters.length !== type.arguments.length) return false;
        return storable(
          replace(
            alias.body,
            new Map(alias.parameters.map((parameter, index) => [parameter.id, type.arguments[index]])),
          ) as ValueTypeExpression,
          depth + 1,
        );
      }
      return type.kind === "scalar" || (type.kind === "builtin" && type.name === "unit");
    };
    let remaining = 10000;
    const visit = (node: unknown, path: string, depth = 0): void => {
      if (depth > 128 || --remaining < 0)
        throw new GenericTypeError("type-complexity-limit", path, "Type exceeds the depth or expansion budget");
      if (!node || typeof node !== "object") return;
      if ("definitionId" in node && "arguments" in node) {
        const application = node as InterfaceApplicationExpression | Extract<ValueTypeExpression, { kind: "alias" }>;
        const alias = "kind" in application && application.kind === "alias";
        const target = alias ? aliases.get(application.definitionId) : this.definitions.get(application.definitionId);
        if (!target)
          throw new GenericTypeError(
            alias ? "unknown-alias" : "unknown-interface",
            path,
            `Unknown definition ${application.definitionId}`,
          );
        if ("template" in target && target.template?.usesSelf) template.usesSelf = true;
        const parameters = "parameters" in target ? target.parameters : (target.template?.parameters ?? []);
        if (parameters.length !== application.arguments.length)
          throw new GenericTypeError(
            "type-arity",
            path,
            `Expected ${parameters.length} type arguments, received ${application.arguments.length}`,
          );
        parameters.forEach((parameter, index) => {
          if (parameter.kind !== application.arguments[index].kind)
            throw new GenericTypeError(
              "parameter-kind",
              path,
              `Expected ${parameter.kind}, received ${application.arguments[index].kind}`,
            );
          const argument = application.arguments[index];
          if (parameter.kind === "value" && parameter.storable && argument.kind === "value" && !storable(argument.type))
            throw new GenericTypeError(
              "non-storable-argument",
              path,
              "Generic application cannot prove its value argument is storable",
            );
          if (parameter.kind === "object" && argument.kind === "object" && argument.target.kind === "parameter") {
            const offered = parametersById.get(argument.target.parameterId);
            const mapping = new Map(parameters.map((p, i) => [p.id, application.arguments[i]]));
            for (const bound of parameter.implements) {
              const required = replace(bound, mapping) as InterfaceApplicationExpression;
              if (offered?.kind !== "object" || !offered.implements.some((evidence) => implies(evidence, required)))
                throw new GenericTypeError(
                  "unsatisfied-bound",
                  path,
                  `Object parameter does not prove ${required.definitionId}`,
                );
            }
          }
        });
      }
      for (const [name, child] of Object.entries(node)) visit(child, `${path}.${name}`, depth + 1);
    };
    visit(template, definition.displayName);
    const active = new Set<string>();
    const done = new Set<string>();
    const checkAlias = (id: string) => {
      if (done.has(id)) return;
      if (active.has(id))
        throw new GenericTypeError("recursive-alias", id, `Recursive value alias: ${[...active, id].join(" -> ")}`);
      active.add(id);
      const walk = (node: unknown): void => {
        if (!node || typeof node !== "object") return;
        if ("kind" in node && node.kind === "alias")
          checkAlias((node as ValueAliasDefinition["body"] & { definitionId: string }).definitionId);
        Object.values(node).forEach(walk);
      };
      walk(aliases.get(id)?.body);
      active.delete(id);
      done.add(id);
    };
    for (const id of aliases.keys()) checkAlias(id);
  }

  apply(id: string, arguments_: readonly ClosedTypeArgument[]): InterfaceRevision {
    const definition = this.definitions.get(id);
    if (!definition) throw new GenericTypeError("unknown-interface", id, "Unknown interface definition");
    const application = {
      definitionId: definition.revisionId,
      source: definition.source,
      arguments: [...arguments_],
      ...(definition.template?.usesSelf ? { self: this.self } : {}),
    };
    const appliedId = definition.template ? appliedInterfaceId(application) : definition.revisionId;
    const cached = this.applications.get(appliedId);
    if (cached) return cached;
    if (this.active.includes(id))
      throw new GenericTypeError(
        "recursive-application",
        id,
        `Expanding recursive interface application: ${[...this.active, id].join(" -> ")}`,
      );
    if (++this.applicationCount > 10000)
      throw new GenericTypeError("type-complexity-limit", id, "Too many interface applications");
    this.active.push(id);
    if (definition.template)
      this.applications.set(appliedId, {
        interfaceId: definition.interfaceId,
        revisionId: appliedId,
        displayName: definition.displayName,
        source: definition.source,
        members: [],
        application,
      });
    try {
      const obligations: NonNullable<InterfaceRevision["argumentRequirements"]> = [];
      const instance = instantiateInterface(definition, arguments_, {
        ...this.environment(),
        implementsInterface: (target, required) => {
          obligations.push({ target, required });
          return true;
        },
      });
      if (instance === definition) return definition;
      instance.argumentRequirements = obligations;
      this.applications.set(instance.revisionId, instance);
      this.definitions.set(instance.revisionId, instance);
      return instance;
    } catch (error) {
      this.applications.delete(appliedId);
      throw error;
    } finally {
      this.active.pop();
    }
  }

  interface(context: InterfaceTypeContext): InterfaceApplicationExpression {
    return this.interfaceByName(context.identifier().getText(), context.typeArguments());
  }

  interfaceByName(name: string, arguments_: TypeArgumentsContext | null): InterfaceApplicationExpression {
    const definition = this.interfaces.get(name);
    if (!definition) throw new GenericTypeError("unknown-interface", name, "Unknown interface");
    return { definitionId: definition.revisionId, arguments: this.arguments(arguments_) };
  }

  arguments(context: TypeArgumentsContext | null): TypeArgumentExpression[] {
    return (context?.typeArgument() ?? []).map((argument): TypeArgumentExpression => {
      if (argument.INTERFACE())
        return {
          kind: "object",
          target: { kind: "application", application: this.interface(argument.interfaceType()!) },
        };
      if (argument.ATOM()) return { kind: "object", target: this.atom(argument.identifier()!.getText()) };
      if (argument.OBJECT()) return { kind: "object", target: this.objectParameter(argument.identifier()!.getText()) };
      const type = argument.valueType()!;
      if (type.getText() === "Self") return { kind: "object", target: { kind: "self" } };
      const parameter = type.identifier() && this.parameters.get(type.identifier()!.getText());
      if (parameter?.kind === "object" && !type.typeArguments())
        return { kind: "object", target: { kind: "parameter", parameterId: parameter.id } };
      return { kind: "value", type: this.expression(type) };
    });
  }

  private atom(name: string): ObjectTypeExpression {
    if (name === "Self") return { kind: "self" };
    const atomId = this.atoms.get(name);
    if (!atomId) throw new GenericTypeError("unknown-atom", name, "Unknown atom");
    return { kind: "atom", atomId };
  }

  private objectParameter(name: string): ObjectTypeExpression {
    if (name === "Self") return { kind: "self" };
    const parameter = this.parameters.get(name);
    if (!parameter || parameter.kind !== "object")
      throw new GenericTypeError("parameter-kind", name, "Expected an object parameter");
    return { kind: "parameter", parameterId: parameter.id };
  }

  targetExpression(context: TargetConstraintContext): ObjectTypeExpression {
    const name = context.identifier().getText();
    if (context.ATOM()) return this.atom(name);
    if (context.OBJECT()) return this.objectParameter(name);
    return { kind: "application", application: this.interfaceByName(name, context.typeArguments()) };
  }

  expression(context: ValueTypeContext): ValueTypeExpression {
    if (context.scalarType())
      return {
        kind: "scalar",
        name: context.scalarType()!.getText() as Extract<ValueType, { kind: "scalar" }>["name"],
      };
    if (context.UNIT()) return valueType.unit;
    if (context.WATCH_HANDLE()) return valueType.watchHandle;
    if (context.MESSAGE()) return valueType.message(literal(context.stringLiteral()!));
    if (context.OPTIONAL() || context.LIST())
      return { kind: context.LIST() ? "list" : "optional", value: this.expression(context.valueType()!) };
    if (context.RECORD()) {
      const fields = context
        .recordField()
        .map((field) => [field.identifier().getText(), this.expression(field.valueType())] as const);
      if (new Set(fields.map(([name]) => name)).size !== fields.length)
        throw new GenericTypeError("duplicate-field", "record", "Duplicate record field");
      return { kind: "record", fields: Object.fromEntries(fields) };
    }
    const name = context.identifier()!.getText();
    if (context.ATOM_REF()) return { kind: "object-ref", expectation: this.atom(name) };
    if (context.INTERFACE_REF())
      return {
        kind: "object-ref",
        expectation: { kind: "application", application: this.interfaceByName(name, context.typeArguments()) },
      };
    if (context.REF()) return { kind: "object-ref", expectation: this.objectParameter(name) };
    const parameter = this.parameters.get(name);
    if (parameter) {
      if (parameter.kind !== "value" || context.typeArguments())
        throw new GenericTypeError("parameter-kind", name, "Expected a value parameter; object parameters need ref<T>");
      return { kind: "parameter", parameterId: parameter.id };
    }
    if (!this.aliases.has(name)) throw new GenericTypeError("unknown-type", name, "Unknown value type");
    return { kind: "alias", definitionId: name, arguments: this.arguments(context.typeArguments()) };
  }

  value(context: ValueTypeContext): ValueType {
    return new TypeSubstitution(this.environment()).value(this.expression(context));
  }

  target(context: TargetConstraintContext) {
    return new TypeSubstitution(this.environment()).object(this.targetExpression(context));
  }

  declareParameters(context: TypeParametersContext | null, owner: string): TypeParameter[] {
    const entries = context?.typeParameter() ?? [];
    const result: TypeParameter[] = entries.map((parameter, index) => {
      const name = parameter.identifier().getText();
      if (name === "Self" || this.parameters.has(name))
        throw new GenericTypeError("duplicate-parameter", owner, `Duplicate or reserved parameter ${name}`);
      const declaration: TypeParameter = parameter.VALUE()
        ? {
            id: `${owner}/parameter/${index}`,
            name,
            kind: "value",
            ...(parameter.STORABLE() ? { storable: true } : {}),
          }
        : { id: `${owner}/parameter/${index}`, name, kind: "object", implements: [] };
      this.parameters.set(name, declaration);
      return declaration;
    });
    entries.forEach((entry, index) => {
      const parameter = result[index];
      if (parameter.kind === "object")
        parameter.implements = entry.interfaceType().map((bound) => this.interface(bound));
    });
    return result;
  }

  declareAliases(contexts: readonly TypeAliasDeclContext[]) {
    for (const context of contexts) {
      const name = context.identifier().getText();
      if (this.aliases.has(name)) throw new GenericTypeError("duplicate-alias", name, "Duplicate type alias");
      this.aliases.set(name, { id: name, parameters: [], body: valueType.unit });
    }
    for (const context of contexts) {
      const name = context.identifier().getText();
      const previous = this.parameters;
      this.parameters = new Map();
      try {
        this.aliases.set(name, {
          id: name,
          parameters: this.declareParameters(context.typeParameters(), JSON.stringify(["alias", name])),
          body: this.expression(context.valueType()),
        });
      } finally {
        this.parameters = previous;
      }
    }
  }
}