Socrates/apps/web/lib/llm/analyze/model.ts

// Analyze pass: derive (or refresh) the SysML model from prose + taxonomy.
//
// Differs from generateModel.ts (the seed-time generator) in two ways:
//  1. Inputs are prose + the current taxonomy term list, not a SeedDraft.
//  2. Re-runs *merge* with the existing model: blocks already linked to a
//     taxonomy term are preserved (id, position, properties) so the user's
//     manual canvas work isn't blown away on every Analyze.

import "server-only";
import { defaultGateway, chatJSON, type Message } from "../gateway";
import { loadPrompt } from "../prompts";
import type {
  SysMLModel,
  Block,
  Association,
  Constraint,
  Requirement,
  PropertyType,
  ReviewStatus,
} from "../../sysml/model";

interface LeanProperty {
  name: string;
  type: { kind: "string" | "number" | "boolean" | "enum"; values?: string[] };
}
interface LeanBlock {
  id: string;
  label: string;
  kind: "system" | "actor" | "block";
  linkedTermLabel?: string;
  properties?: LeanProperty[];
  confidence: number;
}
interface LeanAssoc {
  id: string;
  fromBlockId: string;
  toBlockId: string;
  label?: string;
  kind: "association" | "composition" | "aggregation" | "generalization" | "constraintApplies";
  /** Optional concept-name for the relationship itself (T2 integration). */
  linkedTermLabel?: string;
  confidence: number;
}
interface LeanConstraint {
  id: string;
  label: string;
  expression?: string;
  appliesTo?: string[];
  /** Optional concept-name the constraint enforces. */
  linkedTermLabel?: string;
  confidence: number;
}
interface LeanRequirement {
  id: string;
  tag: string;
  text: string;
  satisfiedBy?: string[];
  /** Optional concept this requirement is "about." */
  linkedTermLabel?: string;
  confidence: number;
}
interface LeanModel {
  systemOfInterestId?: string;
  blocks: LeanBlock[];
  associations?: LeanAssoc[];
  constraints?: LeanConstraint[];
  requirements?: LeanRequirement[];
  overallConfidence?: number;
}

const modelJsonSchema = {
  type: "object",
  additionalProperties: false,
  required: ["blocks"],
  properties: {
    systemOfInterestId: { type: "string" },
    blocks: {
      type: "array",
      items: {
        type: "object",
        additionalProperties: false,
        required: ["id", "label", "kind", "confidence"],
        properties: {
          id: { type: "string", minLength: 1 },
          label: { type: "string", minLength: 1 },
          kind: { type: "string", enum: ["system", "actor", "block"] },
          linkedTermLabel: { type: "string" },
          confidence: { type: "number", minimum: 0, maximum: 1 },
          properties: {
            type: "array",
            items: {
              type: "object",
              additionalProperties: false,
              required: ["name", "type"],
              properties: {
                name: { type: "string", minLength: 1 },
                type: {
                  type: "object",
                  additionalProperties: false,
                  required: ["kind"],
                  properties: {
                    kind: { type: "string", enum: ["string", "number", "boolean", "enum"] },
                    values: { type: "array", items: { type: "string" } },
                  },
                },
              },
            },
          },
        },
      },
    },
    associations: {
      type: "array",
      items: {
        type: "object",
        additionalProperties: false,
        required: ["id", "fromBlockId", "toBlockId", "kind", "confidence"],
        properties: {
          id: { type: "string", minLength: 1 },
          fromBlockId: { type: "string", minLength: 1 },
          toBlockId: { type: "string", minLength: 1 },
          label: { type: "string" },
          kind: {
            type: "string",
            enum: ["association", "composition", "aggregation", "generalization", "constraintApplies"],
          },
          linkedTermLabel: { type: "string" },
          confidence: { type: "number", minimum: 0, maximum: 1 },
        },
      },
    },
    constraints: {
      type: "array",
      items: {
        type: "object",
        additionalProperties: false,
        required: ["id", "label", "confidence"],
        properties: {
          id: { type: "string", minLength: 1 },
          label: { type: "string", minLength: 1 },
          expression: { type: "string" },
          appliesTo: { type: "array", items: { type: "string" } },
          linkedTermLabel: { type: "string" },
          confidence: { type: "number", minimum: 0, maximum: 1 },
        },
      },
    },
    requirements: {
      type: "array",
      items: {
        type: "object",
        additionalProperties: false,
        required: ["id", "tag", "text", "confidence"],
        properties: {
          id: { type: "string", minLength: 1 },
          tag: { type: "string", minLength: 1 },
          text: { type: "string", minLength: 1 },
          satisfiedBy: { type: "array", items: { type: "string" } },
          linkedTermLabel: { type: "string" },
          confidence: { type: "number", minimum: 0, maximum: 1 },
        },
      },
    },
    overallConfidence: { type: "number", minimum: 0, maximum: 1 },
  },
} as const;

export interface AnalyzeModelResult {
  model: SysMLModel;
  termIdToBlockId: Map<string, string>;
  inputTokens: number;
  outputTokens: number;
  provider: string;
  model_name: string;
}

export async function analyzeModelFromProse(
  documentText: string,
  terms: Array<{ id: string; label: string; linkedBlockId: string | null }>,
  current: SysMLModel
): Promise<AnalyzeModelResult> {
  const character = loadPrompt("socrates/character.md");
  const generate = loadPrompt("socrates/analyze-model.md");

  const userPayload =
    `Document:\n\n${documentText}\n\n---\n\n` +
    `Taxonomy:\n\`\`\`json\n${JSON.stringify(terms.map(t => t.label), null, 2)}\n\`\`\`\n` +
    `Return only the JSON object conforming to the schema.`;

  const messages: Message[] = [
    { role: "system", content: `${character}\n\n---\n\n${generate}` },
    { role: "user", content: userPayload },
  ];

  const gateway = defaultGateway();
  const { value, result } = await chatJSON<LeanModel>(gateway, messages, {
    temperature: 0.3,
    maxTokens: 2400,
    jsonSchema: { name: "sysml_model", schema: modelJsonSchema as Record<string, unknown> },
    jsonObjectMode: true,
    maxRepairs: 2,
  });

  const merged = mergeIntoCurrent(value, current, terms);

  return {
    ...merged,
    inputTokens: result.inputTokens,
    outputTokens: result.outputTokens,
    provider: gateway.provider,
    model_name: gateway.model,
  };
}

// ─── Lean → canonical, merged with current model (review-aware) ─────────
//
// Re-running the analyzer is non-destructive. For each element kind we
// match against current by a stable key:
//   block       — id (preserved via linkedTermId chain) OR linkedTermId
//   association — `${kind}|${fromBlockId}|${toBlockId}`
//   constraint  — `${linkedTermId ?? "label:"+normalize(label)}`
//   requirement — `${linkedTermId ?? "text:"+normalize(text)}`
//
// New incoming → reviewStatus = "suggested".
// Matched + currently deprecated → flips back to "suggested".
// Existing not in incoming + not pinned + not currently suggested → "deprecated".
// Existing not in incoming + status was "suggested" → quietly dropped.
// Pinned elements left alone regardless.

function mergeIntoCurrent(
  lean: LeanModel | null | undefined,
  current: SysMLModel,
  terms: Array<{ id: string; label: string; linkedBlockId: string | null }>
): { model: SysMLModel; termIdToBlockId: Map<string, string> } {
  if (!lean || !Array.isArray(lean.blocks)) {
    return { model: current, termIdToBlockId: new Map() };
  }

  const termByLabel = new Map<string, { id: string; linkedBlockId: string | null }>();
  for (const t of terms) termByLabel.set(t.label.toLowerCase(), { id: t.id, linkedBlockId: t.linkedBlockId });
  const resolveTermId = (label: string | undefined): string | undefined => {
    if (!label) return undefined;
    return termByLabel.get(label.toLowerCase())?.id;
  };

  // ─── Blocks ────────────────────────────────────────────────────────────

  const existingByLinkedTerm = new Map<string, Block>();
  const existingById = new Map<string, Block>();
  for (const b of current.blocks) {
    existingById.set(b.id, b);
    if (b.linkedTermId) existingByLinkedTerm.set(b.linkedTermId, b);
  }

  const blocks: Block[] = [];
  const leanIdToCanonical = new Map<string, string>();
  const termIdToBlockId = new Map<string, string>();
  const matchedBlockIds = new Set<string>();

  for (const b of lean.blocks) {
    const linkedTerm = b.linkedTermLabel ? termByLabel.get(b.linkedTermLabel.toLowerCase()) : undefined;

    let existing: Block | undefined;
    if (linkedTerm) existing = existingByLinkedTerm.get(linkedTerm.id);
    if (!existing && existingById.has(b.id)) existing = existingById.get(b.id);

    let canonicalId = b.id;
    let stereotypes: string[] = [b.kind];
    let propertiesSource = b.properties ?? [];
    let nextStatus: ReviewStatus = "suggested";
    let pinned = false;

    if (existing) {
      canonicalId = existing.id;
      stereotypes = existing.stereotypes;
      propertiesSource = mergeProperties(existing.properties, b.properties ?? []);
      pinned = !!existing.reviewPinned;
      // Existing accepted stays accepted; existing deprecated flips back to
      // suggested (analyzer changed its mind); existing suggested stays
      // suggested.
      const cur: ReviewStatus = (existing.reviewStatus ?? "accepted") as ReviewStatus;
      nextStatus = cur === "deprecated" ? "suggested" : cur;
      matchedBlockIds.add(existing.id);
    }

    if (linkedTerm) termIdToBlockId.set(linkedTerm.id, canonicalId);
    leanIdToCanonical.set(b.id, canonicalId);

    blocks.push({
      id: canonicalId,
      label: b.label,
      kind: b.kind,
      stereotypes,
      properties: propertiesSource.map((p, i) => ({
        id: `${canonicalId}_p${i + 1}`,
        name: p.name,
        type: expandPropertyType(p.type),
        multiplicity: "0..1" as const,
      })),
      linkedTermId: linkedTerm?.id,
      reviewStatus: nextStatus,
      reviewPinned: pinned,
    });
  }

  // Existing blocks the analyzer didn't mention.
  for (const b of current.blocks) {
    if (matchedBlockIds.has(b.id)) continue;
    const cur: ReviewStatus = (b.reviewStatus ?? "accepted") as ReviewStatus;
    if (cur === "suggested") continue; // never reviewed + dropped → drop.
    if (b.reviewPinned) {
      blocks.push(b);
      continue;
    }
    blocks.push({ ...b, reviewStatus: "deprecated" });
  }

  // ─── Associations ─────────────────────────────────────────────────────

  const associations = mergeListByKey<Association, LeanAssoc>({
    existing: current.associations,
    incoming: lean.associations ?? [],
    keyExisting: a => `${a.kind}|${a.fromBlockId}|${a.toBlockId}`,
    keyIncoming: a =>
      `${a.kind}|${leanIdToCanonical.get(a.fromBlockId) ?? a.fromBlockId}|${
        leanIdToCanonical.get(a.toBlockId) ?? a.toBlockId
      }`,
    fromIncoming: a => ({
      id: a.id,
      fromBlockId: leanIdToCanonical.get(a.fromBlockId) ?? a.fromBlockId,
      toBlockId: leanIdToCanonical.get(a.toBlockId) ?? a.toBlockId,
      label: a.label ?? "",
      kind: a.kind,
      linkedTermId: resolveTermId(a.linkedTermLabel),
    }),
    mergeFields: (existing, incoming) => ({
      ...existing,
      label: incoming.label ?? existing.label,
      // linkedTermId: prefer fresh suggestion, fall back to existing.
      linkedTermId:
        resolveTermId(incoming.linkedTermLabel) ?? existing.linkedTermId,
    }),
  });

  // ─── Constraints ──────────────────────────────────────────────────────

  const constraints = mergeListByKey<Constraint, LeanConstraint>({
    existing: current.constraints,
    incoming: lean.constraints ?? [],
    keyExisting: c => c.linkedTermId ? `term:${c.linkedTermId}` : `label:${normalizeText(c.label)}`,
    keyIncoming: c => {
      const tid = resolveTermId(c.linkedTermLabel);
      return tid ? `term:${tid}` : `label:${normalizeText(c.label)}`;
    },
    fromIncoming: c => ({
      id: c.id,
      label: c.label,
      expression: c.expression ?? "",
      appliesTo: (c.appliesTo ?? []).map(x => leanIdToCanonical.get(x) ?? x),
      linkedTermId: resolveTermId(c.linkedTermLabel),
    }),
    mergeFields: (existing, incoming) => ({
      ...existing,
      label: incoming.label,
      expression: incoming.expression ?? existing.expression,
      appliesTo: (incoming.appliesTo ?? []).map(x => leanIdToCanonical.get(x) ?? x),
      linkedTermId: resolveTermId(incoming.linkedTermLabel) ?? existing.linkedTermId,
    }),
  });

  // ─── Requirements ─────────────────────────────────────────────────────

  const requirements = mergeListByKey<Requirement, LeanRequirement>({
    existing: current.requirements,
    incoming: lean.requirements ?? [],
    keyExisting: r => r.linkedTermId ? `term:${r.linkedTermId}` : `text:${normalizeText(r.text)}`,
    keyIncoming: r => {
      const tid = resolveTermId(r.linkedTermLabel);
      return tid ? `term:${tid}` : `text:${normalizeText(r.text)}`;
    },
    fromIncoming: r => ({
      id: r.id,
      tag: r.tag,
      text: r.text,
      relations: (r.satisfiedBy ?? []).map(blockId => ({
        kind: "satisfy" as const,
        blockId: leanIdToCanonical.get(blockId) ?? blockId,
      })),
      linkedTermId: resolveTermId(r.linkedTermLabel),
    }),
    mergeFields: (existing, incoming) => ({
      ...existing,
      tag: incoming.tag || existing.tag,
      text: incoming.text || existing.text,
      relations: (incoming.satisfiedBy ?? []).map(blockId => ({
        kind: "satisfy" as const,
        blockId: leanIdToCanonical.get(blockId) ?? blockId,
      })),
      linkedTermId: resolveTermId(incoming.linkedTermLabel) ?? existing.linkedTermId,
    }),
  });

  let soiId = lean.systemOfInterestId ? leanIdToCanonical.get(lean.systemOfInterestId) ?? lean.systemOfInterestId : undefined;
  if (!soiId) {
    const systems = blocks.filter(b => b.kind === "system");
    if (systems.length === 1) soiId = systems[0]!.id;
  }

  return { model: { systemOfInterestId: soiId, blocks, associations, constraints, requirements }, termIdToBlockId };
}

// ─── Generic merge-with-review helper ──────────────────────────────────

interface ReviewableElement {
  reviewStatus?: ReviewStatus;
  reviewPinned?: boolean;
}

interface MergeArgs<E extends ReviewableElement, I> {
  existing: E[];
  incoming: I[];
  keyExisting: (e: E) => string;
  keyIncoming: (i: I) => string;
  fromIncoming: (i: I) => E;
  mergeFields: (existing: E, incoming: I) => E;
}

function mergeListByKey<E extends ReviewableElement, I>(args: MergeArgs<E, I>): E[] {
  const existingByKey = new Map<string, E>();
  for (const e of args.existing) existingByKey.set(args.keyExisting(e), e);

  const incomingByKey = new Map<string, I>();
  for (const i of args.incoming) {
    const k = args.keyIncoming(i);
    if (!incomingByKey.has(k)) incomingByKey.set(k, i);
  }

  const out: E[] = [];
  const matched = new Set<string>();

  // Phase 1 — incoming.
  for (const [key, inc] of incomingByKey.entries()) {
    const prior = existingByKey.get(key);
    if (!prior) {
      const created = args.fromIncoming(inc);
      out.push({ ...created, reviewStatus: "suggested", reviewPinned: false });
    } else {
      const merged = args.mergeFields(prior, inc);
      const cur: ReviewStatus = (prior.reviewStatus ?? "accepted") as ReviewStatus;
      const nextStatus: ReviewStatus = cur === "deprecated" ? "suggested" : cur;
      out.push({ ...merged, reviewStatus: nextStatus, reviewPinned: !!prior.reviewPinned });
      matched.add(key);
    }
  }

  // Phase 2 — existing not in incoming.
  for (const e of args.existing) {
    const k = args.keyExisting(e);
    if (matched.has(k)) continue;
    const cur: ReviewStatus = (e.reviewStatus ?? "accepted") as ReviewStatus;
    if (cur === "suggested") continue; // unreviewed + dropped → drop.
    if (e.reviewPinned) {
      out.push(e);
      continue;
    }
    out.push({ ...e, reviewStatus: "deprecated" });
  }

  return out;
}

function normalizeText(s: string): string {
  return (s ?? "").trim().toLowerCase().replace(/\s+/g, " ");
}

function mergeProperties(
  current: { name: string }[],
  incoming: LeanProperty[]
): LeanProperty[] {
  const have = new Set(current.map(p => p.name.toLowerCase()));
  const merged: LeanProperty[] = current.map(p => ({
    name: p.name,
    type: { kind: "string" }, // type re-expanded in caller
  }));
  for (const p of incoming) {
    if (!have.has(p.name.toLowerCase())) merged.push(p);
  }
  return merged;
}

function expandPropertyType(type: { kind: string; values?: string[] }): PropertyType {
  if (type.kind === "enum") return { kind: "enum", values: type.values ?? [] };
  if (type.kind === "number") return { kind: "number" };
  if (type.kind === "boolean") return { kind: "boolean" };
  return { kind: "string" };
}