// Server-side: ask Socrates to propose a single model change.
//
// Returns { reasoning, ops } as a structured JSON object. The op schema uses
// oneOf per kind (Phase 0 lesson: small models cram all op fields into the
// kind name without it). Caller is responsible for applying the ops; this
// function only generates them.

import "server-only";
import { defaultGateway, chatJSON, type Message } from "./gateway";
import { loadPrompt } from "./prompts";
import type { SysMLModel } from "../sysml/model";
import type { ValidationIssue } from "../sysml/validate";
import type { ModelOp } from "../sync/ops";

// ─── Op-shape JSON Schema (oneOf per kind) ──────────────────────────────

const propertyTypeSchema = {
  type: "object",
  additionalProperties: false,
  required: ["kind"],
  properties: {
    kind: { type: "string", enum: ["string", "number", "boolean", "enum"] },
    values: { type: "array", items: { type: "string" } },
  },
} as const;

const blockShape = {
  type: "object",
  additionalProperties: false,
  required: ["id", "label", "kind"],
  properties: {
    id: { type: "string", minLength: 1 },
    label: { type: "string", minLength: 1 },
    kind: { type: "string", enum: ["system", "actor", "block"] },
    properties: {
      type: "array",
      items: {
        type: "object",
        additionalProperties: false,
        required: ["name", "type"],
        properties: {
          name: { type: "string", minLength: 1 },
          type: propertyTypeSchema,
        },
      },
    },
  },
} as const;

const associationShape = {
  type: "object",
  additionalProperties: false,
  required: ["id", "fromBlockId", "toBlockId", "kind"],
  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"] },
  },
} as const;

const constraintShape = {
  type: "object",
  additionalProperties: false,
  required: ["id", "label"],
  properties: {
    id: { type: "string", minLength: 1 },
    label: { type: "string", minLength: 1 },
    expression: { type: "string" },
    appliesTo: { type: "array", items: { type: "string" } },
  },
} as const;

const requirementShape = {
  type: "object",
  additionalProperties: false,
  required: ["id", "tag", "text"],
  properties: {
    id: { type: "string", minLength: 1 },
    tag: { type: "string", minLength: 1 },
    text: { type: "string", minLength: 1 },
    satisfiedBy: { type: "array", items: { type: "string" } },
  },
} as const;

const proposalJsonSchema = {
  type: "object",
  additionalProperties: false,
  required: ["reasoning", "ops"],
  properties: {
    reasoning: { type: "string", minLength: 1 },
    ops: {
      type: "array",
      items: {
        oneOf: [
          { type: "object", additionalProperties: false, required: ["kind", "block"],
            properties: { kind: { type: "string", enum: ["add-block"] }, block: blockShape } },
          { type: "object", additionalProperties: false, required: ["kind", "blockId"],
            properties: { kind: { type: "string", enum: ["remove-block"] }, blockId: { type: "string", minLength: 1 } } },
          { type: "object", additionalProperties: false, required: ["kind", "association"],
            properties: { kind: { type: "string", enum: ["add-association"] }, association: associationShape } },
          { type: "object", additionalProperties: false, required: ["kind", "associationId"],
            properties: { kind: { type: "string", enum: ["remove-association"] }, associationId: { type: "string", minLength: 1 } } },
          { type: "object", additionalProperties: false, required: ["kind", "constraint"],
            properties: { kind: { type: "string", enum: ["add-constraint"] }, constraint: constraintShape } },
          { type: "object", additionalProperties: false, required: ["kind", "constraintId"],
            properties: { kind: { type: "string", enum: ["remove-constraint"] }, constraintId: { type: "string", minLength: 1 } } },
          { type: "object", additionalProperties: false, required: ["kind", "requirement"],
            properties: { kind: { type: "string", enum: ["add-requirement"] }, requirement: requirementShape } },
          { type: "object", additionalProperties: false, required: ["kind", "requirementId"],
            properties: { kind: { type: "string", enum: ["remove-requirement"] }, requirementId: { type: "string", minLength: 1 } } },
        ],
      },
    },
  },
} as const;

// ─── Public API ──────────────────────────────────────────────────────────

export interface ProposeArgs {
  model: SysMLModel;
  issues?: ValidationIssue[];
  /** Optional user-supplied prompt narrowing what to propose. */
  userPrompt?: string;
}

interface RawProposalOp {
  kind: string;
  [k: string]: unknown;
}

export interface ProposeResult {
  reasoning: string;
  /** Ops in their LLM-emitted shape — caller normalizes to the canonical
   * ModelOp tagged-union (this handles add-block.block.properties needing
   * id/multiplicity defaults, etc.). */
  ops: ModelOp[];
  inputTokens: number;
  outputTokens: number;
  provider: string;
  model: string;
}

export async function proposeChange(args: ProposeArgs): Promise<ProposeResult> {
  const character = loadPrompt("socrates/character.md");
  const propose = loadPrompt("socrates/propose.md");

  const userPayload =
    `Current model:\n\`\`\`json\n${JSON.stringify(trimModel(args.model), null, 2)}\n\`\`\`\n\n` +
    (args.issues?.length
      ? `Active validation issues:\n\`\`\`json\n${JSON.stringify(args.issues.slice(0, 20), null, 2)}\n\`\`\`\n\n`
      : "") +
    (args.userPrompt
      ? `User has asked:\n${args.userPrompt}\n`
      : `Pick the highest-value change you can make right now.\n`);

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

  const gateway = defaultGateway();
  const { value, result } = await chatJSON<{ reasoning: string; ops: RawProposalOp[] }>(
    gateway,
    messages,
    {
      temperature: 0.3,
      maxTokens: 1024,
      jsonSchema: { name: "proposal", schema: proposalJsonSchema as Record<string, unknown> },
      jsonObjectMode: true,
      maxRepairs: 2,
    }
  );

  const ops = normalizeOps(value?.ops ?? []);
  return {
    reasoning: typeof value?.reasoning === "string" ? value.reasoning : "",
    ops,
    inputTokens: result.inputTokens,
    outputTokens: result.outputTokens,
    provider: gateway.provider,
    model: gateway.model,
  };
}

// ─── Normalize the LLM-shape ops to canonical ModelOp ───────────────────

function normalizeOps(raw: RawProposalOp[]): ModelOp[] {
  const out: ModelOp[] = [];
  for (const op of raw) {
    if (!op || typeof op.kind !== "string") continue;
    switch (op.kind) {
      case "add-block": {
        const block = op.block as { id?: string; label?: string; kind?: string; properties?: Array<{ name: string; type: { kind: string; values?: string[] } }> } | undefined;
        if (!block?.id || !block.label || !isBlockKind(block.kind)) continue;
        out.push({
          kind: "add-block",
          block: {
            id: block.id,
            label: block.label,
            kind: block.kind,
            stereotypes: [block.kind],
            properties: (block.properties ?? []).map((p, i) => ({
              id: `p${i + 1}`,
              name: p.name,
              type: p.type.kind === "enum"
                ? { kind: "enum" as const, values: p.type.values ?? [] }
                : { kind: p.type.kind as "string" | "number" | "boolean" },
              multiplicity: "0..1" as const,
            })),
          },
          tempId: block.id,
        });
        break;
      }
      case "remove-block": {
        if (typeof op.blockId === "string") out.push({ kind: "remove-block", blockId: op.blockId });
        break;
      }
      case "add-association": {
        const a = op.association as { id?: string; fromBlockId?: string; toBlockId?: string; label?: string; kind?: string } | undefined;
        if (!a?.id || !a.fromBlockId || !a.toBlockId || !isAssocKind(a.kind)) continue;
        out.push({
          kind: "add-association",
          association: {
            id: a.id,
            fromBlockId: a.fromBlockId,
            toBlockId: a.toBlockId,
            label: a.label ?? "",
            kind: a.kind,
          },
          tempId: a.id,
        });
        break;
      }
      case "remove-association": {
        if (typeof op.associationId === "string") out.push({ kind: "remove-association", associationId: op.associationId });
        break;
      }
      case "add-constraint": {
        const c = op.constraint as { id?: string; label?: string; expression?: string; appliesTo?: string[] } | undefined;
        if (!c?.id || !c.label) continue;
        out.push({
          kind: "add-constraint",
          constraint: { id: c.id, label: c.label, expression: c.expression ?? "", appliesTo: c.appliesTo ?? [] },
          tempId: c.id,
        });
        break;
      }
      case "remove-constraint": {
        if (typeof op.constraintId === "string") out.push({ kind: "remove-constraint", constraintId: op.constraintId });
        break;
      }
      case "add-requirement": {
        const r = op.requirement as { id?: string; tag?: string; text?: string; satisfiedBy?: string[] } | undefined;
        if (!r?.id || !r.tag || !r.text) continue;
        out.push({
          kind: "add-requirement",
          requirement: {
            id: r.id,
            tag: r.tag,
            text: r.text,
            relations: (r.satisfiedBy ?? []).map(b => ({ kind: "satisfy" as const, blockId: b })),
          },
          tempId: r.id,
        });
        break;
      }
      case "remove-requirement": {
        if (typeof op.requirementId === "string") out.push({ kind: "remove-requirement", requirementId: op.requirementId });
        break;
      }
    }
  }
  return out;
}

function isBlockKind(s: unknown): s is "system" | "actor" | "block" {
  return s === "system" || s === "actor" || s === "block";
}
function isAssocKind(s: unknown): s is "association" | "composition" | "aggregation" | "generalization" {
  return s === "association" || s === "composition" || s === "aggregation" || s === "generalization";
}

function trimModel(model: SysMLModel): unknown {
  return {
    blocks: model.blocks.map(b => ({
      id: b.id,
      label: b.label,
      kind: b.kind,
      properties: b.properties.map(p => p.name),
    })),
    associations: model.associations.map(a => ({
      id: a.id,
      from: a.fromBlockId,
      to: a.toBlockId,
      label: a.label,
      kind: a.kind,
    })),
    constraints: model.constraints.map(c => ({
      id: c.id,
      label: c.label,
      appliesTo: c.appliesTo,
    })),
    requirements: model.requirements.map(r => ({
      id: r.id,
      tag: r.tag,
      text: r.text,
      satisfiedBy: r.relations
        .filter(rel => rel.kind === "satisfy")
        .map(rel => (rel as { kind: "satisfy"; blockId: string }).blockId),
    })),
  };
}