import ShexJTraverser from "@ldo/traverser-shexj";
import type { Annotation } from "shexj";
import { nameFromObject } from "../context/JsonLdContextBuilder.js";
import type { ShapeInterfaceDeclaration } from "./ShapeInterfaceDeclaration.js";
import { getRdfTypesForTripleConstraint } from "../util/getRdfTypesForTripleConstraint.js";
import * as dom from "dts-dom";

// Collected enum alias names (e.g., AuthenticatedAgentId) to emit at end
export const additionalCompactEnumAliases = new Set<string>();

export interface CompactTransformerContext {
  getNameFromIri: (iri: string, rdfType?: string) => string;
}

function commentFromAnnotations(
  annotations?: Annotation[],
): string | undefined {
  const commentAnnotationObject = annotations?.find(
    (annotation) =>
      annotation.predicate === "http://www.w3.org/2000/01/rdf-schema#comment",
  )?.object;
  if (typeof commentAnnotationObject === "string")
    return commentAnnotationObject;
  return commentAnnotationObject?.value;
}

// Helper: classify a dom.Type into categories we care about.
function isObjectLike(t: dom.Type): boolean {
  return (
    (t as dom.ObjectType).kind === "object" ||
    (t as dom.InterfaceDeclaration).kind === "interface"
  );
}

function isPrimitiveLike(t: dom.Type): boolean {
  const kind = (t as any)?.kind;
  if (kind === "name") return true; // named references and intrinsic tokens
  if (kind === "union") {
    return (t as dom.UnionType).members.every(isPrimitiveLike);
  }
  if (kind === "type-parameter") return true;
  // Fallback: treat scalar intrinsic tokens as primitive
  const intrinsicKinds = new Set(["string", "number", "boolean", "undefined"]);
  return intrinsicKinds.has(kind || "");
}

// Property name collision resolution using predicate IRI mapping
const predicateIriByProp = new WeakMap<dom.PropertyDeclaration, string>();
/**
 * resolveCollisions
 * -----------------
 * Purpose: ensure that properties derived from different predicate IRIs but
 * sharing the same local ending (final segment after '/', '#' or ':') become
 * uniquely named TypeScript properties.
 *
 * Strategy (simplified):
 * 1. Group properties by their current (local) name.
 * 2. For any group with more than one property, rename each to
 *    `${secondLast}_${local}` where `secondLast` is the segment immediately
 *    before the local segment in the predicate IRI.
 * 3. If collisions still remain (i.e. two IRIs share both secondLast and local
 *    segments), fall back to using a sanitized form of the full IRI (without the
 *    protocol) as the property name.
 */
function resolveCollisions(props: dom.PropertyDeclaration[]): void {
  const groups = new Map<string, dom.PropertyDeclaration[]>();
  props.forEach((p) => {
    const name = p.name;
    if (!groups.has(name)) groups.set(name, []);
    groups.get(name)!.push(p);
  });
  groups.forEach((list) => {
    if (list.length < 2) return;

    const predicateIris = new Set(
      list.map((p) => predicateIriByProp.get(p)).filter(Boolean),
    );
    if (predicateIris.size < 2) {
      return;
    }

    // First pass rename using second last segment
    list.forEach((prop) => {
      const iri = predicateIriByProp.get(prop);
      if (!iri) return;
      const segs = iri.split(/[#:\/]/).filter(Boolean);
      if (segs.length < 2) return;
      const local = segs.at(-1)!;
      const secondLast = segs.at(-2)!;
      prop.name = `${secondLast}_${local}`;
    });
    // Detect any remaining duplicates after first pass
    const nameCounts = new Map<string, number>();
    list.forEach((p) =>
      nameCounts.set(p.name, (nameCounts.get(p.name) || 0) + 1),
    );
    list.forEach((p) => {
      if (nameCounts.get(p.name)! > 1) {
        const iri = predicateIriByProp.get(p);
        if (iri) p.name = iri;
      }
    });
  });
}

// Merge duplicate properties without introducing LdSet. If a property appears multiple
// times (e.g., via EXTENDS or grouped expressions) we:
//  - union the types (flattening existing unions)
//  - if one side is Set<T> and the other is plain U, produce Set<T|U>
//  - if both are Set<A>, Set<B> -> Set<A|B>
//  - preserve optional flag if any occurrence optional
function dedupeCompactProperties(
  props: dom.PropertyDeclaration[],
): dom.PropertyDeclaration[] {
  const byName: Record<string, dom.PropertyDeclaration> = {};
  const isSetRef = (t: dom.Type): t is dom.NamedTypeReference =>
    (t as any).kind === "name" && (t as any).name === "Set";
  const getSetInner = (t: dom.Type): dom.Type =>
    isSetRef(t) ? (t as any).typeArguments[0] : t;
  const toSet = (inner: dom.Type): dom.Type =>
    ({ kind: "name", name: "Set", typeArguments: [inner] }) as any;
  const makeUnion = (a: dom.Type, b: dom.Type): dom.Type => {
    const collect = (t: dom.Type, acc: dom.Type[]) => {
      if ((t as any).kind === "union") {
        (t as any).members.forEach((m: dom.Type) => collect(m, acc));
      } else acc.push(t);
    };
    const members: dom.Type[] = [];
    collect(a, members);
    collect(b, members);
    // de-dup via string emission heuristic
    const seen = new Set<string>();
    const filtered: dom.Type[] = [];
    members.forEach((m) => {
      const key =
        (m as any).name ||
        (m as any).value ||
        (m as any).kind + JSON.stringify(m);
      if (!seen.has(key)) {
        seen.add(key);
        filtered.push(m);
      }
    });
    if (filtered.length === 1) return filtered[0];
    return dom.create.union(filtered);
  };
  props.forEach((p) => {
    const existing = byName[p.name];
    if (!existing) {
      byName[p.name] = p;
      return;
    }
    // If predicates differ, keep both (assign numeric suffix to new one)
    const predExisting = predicateIriByProp.get(existing);
    const predNew = predicateIriByProp.get(p);
    if (predExisting && predNew && predExisting !== predNew) {
      const base = p.name;
      let counter = 2;
      while (byName[`${base}${counter}`]) counter++;
      const newName = `${base}${counter}`;
      const clone = dom.create.property(newName, p.type, p.flags);
      clone.jsDocComment = p.jsDocComment;
      if (predNew) predicateIriByProp.set(clone, predNew);
      byName[newName] = clone;
      return;
    }
    const existingSet = isSetRef(existing.type);
    const newSet = isSetRef(p.type);
    let mergedType: dom.Type;
    if (existingSet && newSet) {
      mergedType = toSet(
        makeUnion(getSetInner(existing.type), getSetInner(p.type)),
      );
    } else if (existingSet && !newSet) {
      mergedType = toSet(makeUnion(getSetInner(existing.type), p.type));
    } else if (!existingSet && newSet) {
      mergedType = toSet(makeUnion(existing.type, getSetInner(p.type)));
    } else {
      mergedType = makeUnion(existing.type, p.type);
    }
    const optional =
      existing.flags === dom.DeclarationFlags.Optional ||
      p.flags === dom.DeclarationFlags.Optional
        ? dom.DeclarationFlags.Optional
        : dom.DeclarationFlags.None;
    const merged = dom.create.property(p.name, mergedType, optional);
    merged.jsDocComment =
      existing.jsDocComment && p.jsDocComment
        ? `${existing.jsDocComment} | ${p.jsDocComment}`
        : existing.jsDocComment || p.jsDocComment;
    // Preserve predicate mapping
    const pred = predicateIriByProp.get(existing) || predicateIriByProp.get(p);
    if (pred) predicateIriByProp.set(merged, pred);
    byName[p.name] = merged;
  });
  return Object.values(byName);
}

export const ShexJTypingTransformerCompact = ShexJTraverser.createTransformer<
  {
    Schema: { return: dom.TopLevelDeclaration[] };
    ShapeDecl: { return: dom.InterfaceDeclaration };
    Shape: { return: dom.InterfaceDeclaration };
    EachOf: { return: dom.ObjectType | dom.InterfaceDeclaration };
    TripleConstraint: { return: dom.PropertyDeclaration };
    NodeConstraint: { return: dom.Type };
    ShapeOr: { return: dom.UnionType };
    ShapeAnd: { return: dom.IntersectionType };
    ShapeNot: { return: never };
    ShapeExternal: { return: never };
  },
  CompactTransformerContext
>({
  // Transformer from Schema to interfaces
  Schema: {
    transformer: async (_schema, getTransformedChildren) => {
      const transformedChildren = await getTransformedChildren();
      const interfaces: dom.TopLevelDeclaration[] = [];
      transformedChildren.shapes?.forEach((shape) => {
        if (
          typeof shape !== "string" &&
          (shape as dom.InterfaceDeclaration).kind === "interface"
        ) {
          interfaces.push(shape as dom.InterfaceDeclaration);
        }
      });
      return interfaces;
    },
  },

  // Transformer from ShapeDecl to interface
  ShapeDecl: {
    transformer: async (shapeDecl, getTransformedChildren) => {
      const shapeName = nameFromObject(shapeDecl) || "Shape";
      const { shapeExpr } = await getTransformedChildren();
      if ((shapeExpr as dom.InterfaceDeclaration).kind === "interface") {
        const shapeInterface = shapeExpr as ShapeInterfaceDeclaration;
        shapeInterface.name = shapeName;
        // Preserve shape id for downstream shapeTypes generation
        // (mirrors standard transformer behavior)
        shapeInterface.shapeId = shapeDecl.id;
        if (
          !shapeInterface.members.find(
            (m) => m.kind === "property" && m.name === "id",
          )
        ) {
          shapeInterface.members.unshift(
            dom.create.property(
              "id",
              dom.create.namedTypeReference("IRI"),
              // Root interfaces should have mandatory id
              dom.DeclarationFlags.None,
            ),
          );
        }
        return shapeInterface;
      }
      throw new Error(
        "Unsupported direct shape expression on ShapeDecl for compact format.",
      );
    },
  },

  // Transformer from Shape to interface
  Shape: {
    transformer: async (_shape, getTransformedChildren, setReturnPointer) => {
      const newInterface: ShapeInterfaceDeclaration = dom.create.interface("");
      setReturnPointer(newInterface);
      const transformedChildren = await getTransformedChildren();
      if (
        typeof transformedChildren.expression !== "string" &&
        transformedChildren.expression &&
        ((transformedChildren.expression as dom.ObjectType).kind === "object" ||
          (transformedChildren.expression as dom.InterfaceDeclaration).kind ===
            "interface")
      ) {
        newInterface.members.push(
          ...(transformedChildren.expression as dom.ObjectType).members,
        );
      } else if (
        (transformedChildren.expression as dom.PropertyDeclaration)?.kind ===
        "property"
      ) {
        newInterface.members.push(
          transformedChildren.expression as dom.PropertyDeclaration,
        );
      }
      if (transformedChildren.extends) {
        transformedChildren.extends.forEach((ext) => {
          const extInt = ext as dom.InterfaceDeclaration;
          if (extInt.kind === "interface") {
            const merged = [
              ...extInt.members.filter(
                (m) => !(m.kind === "property" && m.name === "id"),
              ),
              ...newInterface.members,
            ].filter(
              (m): m is dom.PropertyDeclaration => m.kind === "property",
            );
            newInterface.members = dedupeCompactProperties(merged);
          }
        });
      }
      // Final pass: ensure only a single id property
      const idSeen = new Set<number>();
      newInterface.members = newInterface.members.filter((m, idx) => {
        if (m.kind !== "property" || m.name !== "id") return true;
        if (idSeen.size === 0) {
          idSeen.add(idx);
          // normalize id type to IRI
          m.type = dom.create.namedTypeReference("IRI");
          return true;
        }
        return false;
      });
      return newInterface;
    },
  },

  // Transformer from EachOf to object type. EachOf contains the `expressions` array of properties (TripleConstraint)
  EachOf: {
    transformer: async (eachOf, getTransformedChildren, setReturnPointer) => {
      const transformedChildren = await getTransformedChildren();
      const name = nameFromObject(eachOf);

      const objectType = name
        ? dom.create.interface(name)
        : dom.create.objectType([]);
      setReturnPointer(objectType);
      const inputProps: dom.PropertyDeclaration[] = [];
      transformedChildren.expressions.forEach((expr) => {
        if (!expr || typeof expr === "string") return;
        const kind = (expr as any).kind;
        if (kind === "property") {
          inputProps.push(expr as dom.PropertyDeclaration);
        } else if (kind === "object" || kind === "interface") {
          (expr as dom.ObjectType | dom.InterfaceDeclaration).members.forEach(
            (m) => {
              if ((m as any).kind === "property") {
                inputProps.push(m as dom.PropertyDeclaration);
              }
            },
          );
        }
      });
      const deduped = dedupeCompactProperties(inputProps);
      resolveCollisions(deduped);
      objectType.members.push(...deduped);
      return objectType;
    },
  },

  // Transformer from triple constraints to type properties.
  TripleConstraint: {
    transformer: async (
      tripleConstraint,
      getTransformedChildren,
      _setReturnPointer,
      node,
      context,
    ) => {
      const transformedChildren = await getTransformedChildren();
      const rdfTypes = getRdfTypesForTripleConstraint(node);
      const baseName = context.getNameFromIri(
        tripleConstraint.predicate,
        rdfTypes[0],
      );
      const max = tripleConstraint.max;
      const isPlural = max === -1 || (max !== undefined && max !== 1);
      const isOptional = tripleConstraint.min === 0;

      let valueType: dom.Type = dom.type.any;
      if (transformedChildren.valueExpr)
        valueType = transformedChildren.valueExpr as dom.Type;

      if (
        (valueType as dom.InterfaceDeclaration).kind === "interface" &&
        !(valueType as dom.InterfaceDeclaration).name
      ) {
        valueType = dom.create.objectType(
          (valueType as dom.InterfaceDeclaration)
            .members as dom.PropertyDeclaration[],
        );
      }

      // Normalize NodeConstraint returned object forms for IRIs into IRI
      // Heuristic: existing transformer (compact) returns string/number/boolean OR object/interface.
      // We treat any simple string/number/boolean/name as primitive.

      // Determine category
      const objLike = isObjectLike(valueType);
      const isUnion =
        (valueType as unknown as { kind?: string })?.kind === "union";
      const unionMembers: dom.Type[] = isUnion
        ? (valueType as dom.UnionType).members
        : [];
      const unionAllObjLike =
        isUnion && unionMembers.length > 0 && unionMembers.every(isObjectLike);
      const primLike = isPrimitiveLike(valueType);
      if (
        !primLike &&
        !objLike &&
        (valueType as dom.UnionType).kind === "union"
      ) {
        const u = valueType as dom.UnionType;
        const hasObj = u.members.some(isObjectLike);
        const hasPrim = u.members.some(isPrimitiveLike);
        if (isPlural && hasObj && hasPrim) {
          throw new Error(
            `Mixed plural union (object + primitive) not supported for predicate ${tripleConstraint.predicate}`,
          );
        }
      }

      let finalType: dom.Type;
      if (isPlural) {
        if (objLike || unionAllObjLike) {
          if (
            (valueType as dom.InterfaceDeclaration).kind === "interface" &&
            (valueType as dom.InterfaceDeclaration).name
          ) {
            const ifaceName = (valueType as dom.InterfaceDeclaration).name;
            // Dictionary of full object instances keyed by IRI
            finalType = {
              kind: "name",
              name: "Record",
              typeArguments: [
                dom.create.namedTypeReference("IRI"),
                dom.create.namedTypeReference(ifaceName),
              ],
            } as dom.Type;
          } else {
            // Anonymous object or union of anonymous/interface objects
            let valueForRecord: dom.Type = valueType;
            if (unionAllObjLike) {
              // Ensure each union member has id?: IRI if anonymous object
              (valueType as dom.UnionType).members = (
                valueType as dom.UnionType
              ).members.map((m) => {
                if ((m as dom.InterfaceDeclaration).kind === "interface")
                  return m;
                if ((m as dom.ObjectType).kind === "object") {
                  const anonMembers = (
                    m as unknown as { members?: dom.PropertyDeclaration[] }
                  ).members;
                  const hasId = (anonMembers || []).some(
                    (mm) => mm.name === "id",
                  );
                  if (!hasId && anonMembers) {
                    anonMembers.unshift(
                      dom.create.property(
                        "id",
                        dom.create.namedTypeReference("IRI"),
                        dom.DeclarationFlags.None,
                      ),
                    );
                  }
                }
                return m;
              });
              valueForRecord = valueType; // union retained
            } else {
              const anon = valueType as dom.ObjectType;
              const anonMembers = (
                anon as unknown as { members?: dom.PropertyDeclaration[] }
              ).members;
              const hasId = (anonMembers || []).some((m) => m.name === "id");
              if (!hasId && anonMembers) {
                anonMembers.unshift(
                  dom.create.property(
                    "id",
                    dom.create.namedTypeReference("IRI"),
                    dom.DeclarationFlags.None,
                  ),
                );
              }
              valueForRecord = anon as dom.Type;
            }
            finalType = {
              kind: "name",
              name: "Record",
              typeArguments: [
                dom.create.namedTypeReference("IRI"),
                valueForRecord,
              ],
            } as dom.Type;
          }
        } else {
          finalType = {
            kind: "name",
            name: "Set",
            typeArguments: [valueType],
          } as dom.Type;
        }
      } else {
        // Singular
        // If anonymous object or union of object-like types, ensure id: IRI is present (mandatory)
        if (objLike) {
          if ((valueType as dom.ObjectType).kind === "object") {
            const members = (
              valueType as unknown as {
                members?: dom.PropertyDeclaration[];
              }
            ).members;
            const hasId = (members || []).some((m) => m.name === "id");
            if (!hasId && members) {
              members.unshift(
                dom.create.property(
                  "id",
                  dom.create.namedTypeReference("IRI"),
                  dom.DeclarationFlags.None,
                ),
              );
            }
          }
        } else if (isUnion && unionAllObjLike) {
          (valueType as dom.UnionType).members = (
            valueType as dom.UnionType
          ).members.map((m) => {
            if ((m as dom.ObjectType).kind === "object") {
              const mems = (
                m as unknown as { members?: dom.PropertyDeclaration[] }
              ).members;
              const hasId = (mems || []).some((mm) => mm.name === "id");
              if (!hasId && mems) {
                mems.unshift(
                  dom.create.property(
                    "id",
                    dom.create.namedTypeReference("IRI"),
                    dom.DeclarationFlags.None,
                  ),
                );
              }
            }
            return m;
          });
        }
        // Singular: always the interface/object type itself (never Id union)
        if (
          (valueType as dom.InterfaceDeclaration).kind === "interface" &&
          (valueType as dom.InterfaceDeclaration).name
        ) {
          finalType = dom.create.namedTypeReference(
            (valueType as dom.InterfaceDeclaration).name,
          );
        } else {
          finalType = valueType;
        }
      }

      const prop = dom.create.property(
        baseName,
        finalType,
        isOptional ? dom.DeclarationFlags.Optional : dom.DeclarationFlags.None,
      );
      predicateIriByProp.set(prop, tripleConstraint.predicate);
      prop.jsDocComment =
        commentFromAnnotations(tripleConstraint.annotations) || "";
      // Always append original predicate IRI reference (compact format only)
      // If an existing comment is present, add a blank line before the Original IRI line.
      if (prop.jsDocComment) {
        prop.jsDocComment = `${prop.jsDocComment}\n\nOriginal IRI: ${tripleConstraint.predicate}`;
      } else {
        prop.jsDocComment = `Original IRI: ${tripleConstraint.predicate}`;
      }
      return prop;
    },
  },

  // Transformer from node constraint to type
  NodeConstraint: {
    transformer: async (nodeConstraint) => {
      if (nodeConstraint.datatype) {
        switch (nodeConstraint.datatype) {
          case "http://www.w3.org/2001/XMLSchema#boolean":
            return dom.type.boolean;
          case "http://www.w3.org/2001/XMLSchema#byte":
          case "http://www.w3.org/2001/XMLSchema#decimal":
          case "http://www.w3.org/2001/XMLSchema#double":
          case "http://www.w3.org/2001/XMLSchema#float":
          case "http://www.w3.org/2001/XMLSchema#int":
          case "http://www.w3.org/2001/XMLSchema#integer":
          case "http://www.w3.org/2001/XMLSchema#long":
          case "http://www.w3.org/2001/XMLSchema#negativeInteger":
          case "http://www.w3.org/2001/XMLSchema#nonNegativeInteger":
          case "http://www.w3.org/2001/XMLSchema#nonPositiveInteger":
          case "http://www.w3.org/2001/XMLSchema#positiveInteger":
          case "http://www.w3.org/2001/XMLSchema#short":
          case "http://www.w3.org/2001/XMLSchema#unsignedLong":
          case "http://www.w3.org/2001/XMLSchema#unsignedInt":
          case "http://www.w3.org/2001/XMLSchema#unsignedShort":
          case "http://www.w3.org/2001/XMLSchema#unsignedByte":
            return dom.type.number;
          default:
            return dom.type.string; // treat most as string
        }
      }
      if (nodeConstraint.nodeKind) {
        switch (nodeConstraint.nodeKind) {
          case "iri":
            return dom.create.namedTypeReference("IRI");
          case "bnode":
            return dom.type.string; // opaque id as string
          case "nonliteral":
            return dom.create.namedTypeReference("IRI");
          case "literal":
          default:
            return dom.type.string;
        }
      }
      if (nodeConstraint.values) {
        const u = dom.create.union([]);
        nodeConstraint.values.forEach((v) => {
          if (typeof v === "string") u.members.push(dom.type.stringLiteral(v));
        });
        if (!u.members.length) return dom.type.string;
        if (u.members.length === 1) return u.members[0];
        return u;
      }
      return dom.type.any;
    },
  },

  // Transformer from ShapeOr to union type
  ShapeOr: {
    transformer: async (_shapeOr, getTransformedChildren) => {
      const tc = await getTransformedChildren();
      const valid: dom.Type[] = [];
      tc.shapeExprs.forEach((t) => {
        if (typeof t === "object") valid.push(t);
      });
      return dom.create.union(valid);
    },
  },

  // Transformer from ShapeAnd to intersection type
  ShapeAnd: {
    transformer: async (_shapeAnd, getTransformedChildren) => {
      const tc = await getTransformedChildren();
      const valid: dom.Type[] = [];
      tc.shapeExprs.forEach((t) => {
        if (typeof t === "object") valid.push(t);
      });
      return dom.create.intersection(valid);
    },
  },

  // Transformer from ShapeNot to type - not supported.
  ShapeNot: {
    transformer: async () => {
      throw new Error("ShapeNot not supported (compact)");
    },
  },

  // Transformer from ShapeExternal to type - not supported.
  ShapeExternal: {
    transformer: async () => {
      throw new Error("ShapeExternal not supported (compact)");
    },
  },
});