import type { Api, Model } from "@earendil-works/pi-ai";
import { access, readFile } from "node:fs/promises";
import {
  createEditToolDefinition,
  createFindToolDefinition,
  createGrepToolDefinition,
  createLsToolDefinition,
  createReadToolDefinition,
  createWriteToolDefinition,
  defineTool,
  type DefaultResourceLoader,
  type ExtensionContext,
  type SettingsManager,
} from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";
import {
  createChildResources,
  resolveStandaloneChildProjectTrust,
} from "@jbcode/pi-shared/child-session";
import {
  runAgent,
  type AgentOutcome,
  type RunAgentOptions,
  type ThinkingLevel,
  type WorkflowModel,
} from "../_workflows/runner.ts";
import { STRUCTURED_OUTPUT_SYSTEM_INSTRUCTION } from "../_workflows/prompt.ts";
import type {
  AssignmentExecutionError,
  AssignmentExecutionMetadata,
  AssignmentExecutionResult,
  AssignmentRequest,
  AssignmentExecutor,
} from "./engine.ts";
import { GitStateProbe, MutationScope } from "./workspace.ts";

const PROCESS_MAX_ASSIGNMENTS = 32;

interface ScheduledAssignment {
  runId: string;
  runLimit: number;
  exclusive: boolean;
  signal: AbortSignal;
  start: () => void;
  reject: (error: Error) => void;
  state: "queued" | "reserved" | "running" | "settled";
  onAbort: () => void;
}

function abortError(signal: AbortSignal) {
  return signal.reason instanceof Error
    ? signal.reason
    : new Error("Assignment was cancelled");
}

/** Process-wide capacity with synchronous reservation and cancellation-safe release. */
export class ProcessAssignmentScheduler {
  private ceiling = PROCESS_MAX_ASSIGNMENTS;
  private inUse = 0;
  private running = 0;
  private readonly runUsage = new Map<string, number>();
  private readonly exclusiveRuns = new Set<string>();
  private readonly queue: ScheduledAssignment[] = [];

  lowerCeiling(ceiling: number) {
    if (!Number.isInteger(ceiling) || ceiling < 1 || ceiling > PROCESS_MAX_ASSIGNMENTS) {
      throw new Error(`Assignment ceiling must be between 1 and ${PROCESS_MAX_ASSIGNMENTS}`);
    }
    this.ceiling = Math.min(this.ceiling, ceiling);
    this.drain();
  }

  schedule<T>(options: {
    runId: string;
    runLimit: number;
    signal: AbortSignal;
    exclusive?: boolean;
    onReserved?: () => Promise<void>;
    task: (signal: AbortSignal) => Promise<T>;
  }): Promise<T> {
    if (!Number.isInteger(options.runLimit) || options.runLimit < 1) {
      return Promise.reject(new Error("Workflow Run assignment limit must be positive"));
    }
    if (options.signal.aborted) return Promise.reject(abortError(options.signal));

    return new Promise<T>((resolve, reject) => {
      const entry: ScheduledAssignment = {
        runId: options.runId,
        runLimit: options.runLimit,
        signal: options.signal,
        exclusive: options.exclusive ?? false,
        state: "queued",
        reject,
        onAbort: () => {
          if (entry.state !== "queued") return;
          entry.state = "settled";
          const index = this.queue.indexOf(entry);
          if (index >= 0) this.queue.splice(index, 1);
          options.signal.removeEventListener("abort", entry.onAbort);
          reject(abortError(options.signal));
        },
        start: () => {
          queueMicrotask(async () => {
            if (entry.state !== "reserved") return;
            entry.state = "running";
            this.running++;
            try {
              if (options.signal.aborted) throw abortError(options.signal);
              await options.onReserved?.();
              if (options.signal.aborted) throw abortError(options.signal);
              resolve(await options.task(options.signal));
            } catch (error) {
              reject(error);
            } finally {
              this.running--;
              this.release(entry);
            }
          });
        },
      };
      options.signal.addEventListener("abort", entry.onAbort, { once: true });
      this.queue.push(entry);
      this.drain();
    });
  }

  snapshot() {
    return {
      ceiling: this.ceiling,
      inUse: this.inUse,
      running: this.running,
      queued: this.queue.filter((entry) => entry.state === "queued").length,
    };
  }

  private drain() {
    while (this.inUse < this.ceiling) {
      const index = this.queue.findIndex((entry, entryIndex) => {
        if (entry.state !== "queued") return false;
        const usage = this.runUsage.get(entry.runId) ?? 0;
        if (entry.exclusive) return usage === 0;
        if (this.exclusiveRuns.has(entry.runId) || usage >= entry.runLimit) {
          return false;
        }
        return !this.queue
          .slice(0, entryIndex)
          .some(
            (queued) =>
              queued.state === "queued" &&
              queued.runId === entry.runId &&
              queued.exclusive,
          );
      });
      if (index < 0) return;
      const entry = this.queue.splice(index, 1)[0];
      if (!entry || entry.signal.aborted) {
        entry?.onAbort();
        continue;
      }
      entry.state = "reserved";
      this.inUse++;
      this.runUsage.set(entry.runId, (this.runUsage.get(entry.runId) ?? 0) + 1);
      if (entry.exclusive) this.exclusiveRuns.add(entry.runId);
      entry.start();
    }
  }

  private release(entry: ScheduledAssignment) {
    if (entry.state === "settled") return;
    entry.state = "settled";
    entry.signal.removeEventListener("abort", entry.onAbort);
    this.inUse--;
    if (entry.exclusive) this.exclusiveRuns.delete(entry.runId);
    const remaining = (this.runUsage.get(entry.runId) ?? 1) - 1;
    if (remaining === 0) this.runUsage.delete(entry.runId);
    else this.runUsage.set(entry.runId, remaining);
    this.drain();
  }
}

export const processAssignmentScheduler = new ProcessAssignmentScheduler();

export function resolveWorkflowModel(
  registry: ExtensionContext["modelRegistry"],
  profile: { provider?: string; model?: string },
  inherited: WorkflowModel | undefined,
): Model<Api> | undefined {
  if (profile.provider && !profile.model) {
    throw new Error(`Agent Profile provider ${profile.provider} requires a model`);
  }
  if (!profile.model) {
    if (!inherited) return undefined;
    const found = registry.find(inherited.provider, inherited.id);
    if (!found) {
      throw new Error(
        `Unknown inherited model ${inherited.provider}/${inherited.id}`,
      );
    }
    return found;
  }
  if (profile.provider) {
    const found = registry.find(profile.provider, profile.model);
    if (!found) throw new Error(`Unknown model ${profile.provider}/${profile.model}`);
    return found;
  }
  const slash = profile.model.indexOf("/");
  if (slash > 0) {
    const provider = profile.model.slice(0, slash);
    const model = profile.model.slice(slash + 1);
    const found = registry.find(provider, model);
    if (!found) throw new Error(`Unknown model ${profile.model}`);
    return found;
  }
  if (inherited) {
    const found = registry.find(inherited.provider, profile.model);
    if (found) return found;
  }
  const matches = registry.getAll().filter((model) => model.id === profile.model);
  if (matches.length === 1) return matches[0];
  if (matches.length > 1) {
    throw new Error(
      `Model ${profile.model} exists in multiple providers; use provider/model`,
    );
  }
  throw new Error(`Unknown model ${profile.model}`);
}

const DISCOVERY_SCHEMA = {
  type: "object",
  additionalProperties: false,
  required: [
    "questions",
    "candidatePaths",
    "languages",
    "tags",
    "riskTags",
  ],
  properties: {
    questions: {
      type: "array",
      maxItems: 20,
      items: {
        type: "object",
        additionalProperties: false,
        required: ["id", "prompt"],
        properties: {
          id: { type: "string", maxLength: 160 },
          prompt: { type: "string", maxLength: 2000 },
        },
      },
    },
    candidatePaths: { type: "array", items: { type: "string" } },
    languages: { type: "array", items: { type: "string" } },
    tags: { type: "array", items: { type: "string" } },
    riskTags: { type: "array", items: { type: "string" } },
  },
} as const;

const DESIGN_SCHEMA = {
  type: "object",
  additionalProperties: false,
  required: [
    "plan",
    "uiImpact",
    "uiImpactEvidence",
    "affectedUiSurfaces",
    "assignments",
    "acceptanceCriteria",
    "verificationCommands",
  ],
  properties: {
    plan: { type: "string" },
    uiImpact: { enum: ["none", "minor", "non-trivial"] },
    uiImpactEvidence: { type: "array", items: { type: "string" } },
    affectedUiSurfaces: { type: "array", items: { type: "string" } },
    acceptanceCriteria: { type: "array", items: { type: "string" } },
    verificationCommands: {
      type: "array",
      items: {
        type: "object",
        additionalProperties: false,
        required: ["executable", "args", "ownedPaths", "kind"],
        properties: {
          executable: { type: "string" },
          args: { type: "array", items: { type: "string" } },
          ownedPaths: { type: "array", items: { type: "string" } },
          kind: {
            enum: ["test", "typecheck", "lint", "build", "browser", "generation"],
          },
        },
      },
    },
    assignments: {
      type: "array",
      minItems: 3,
      items: {
        type: "object",
        additionalProperties: false,
        required: ["role", "profileName", "rationale", "mutates", "ownedPaths"],
        properties: {
          role: { enum: ["prototype", "implement", "verify", "review"] },
          profileName: { type: "string" },
          rationale: { type: "string" },
          mutates: { type: "boolean" },
          ownedPaths: { type: "array", items: { type: "string" } },
        },
      },
    },
  },
} as const;

const PROTOTYPE_SCHEMA = {
  type: "object",
  additionalProperties: false,
  required: [
    "path",
    "route",
    "structuralRationale",
    "responsiveChecks",
    "browserEvidence",
    "launchInstructions",
  ],
  properties: {
    path: { type: "string" },
    route: { type: "string" },
    structuralRationale: { type: "string" },
    responsiveChecks: { type: "array", items: { type: "string" } },
    browserEvidence: { type: "array", items: { type: "string" } },
    launchInstructions: { type: "string" },
  },
} as const;

const IMPLEMENTATION_SCHEMA = {
  type: "object",
  additionalProperties: false,
  required: ["summary", "changedPaths", "focusedTests"],
  properties: {
    summary: { type: "string" },
    changedPaths: { type: "array", items: { type: "string" } },
    focusedTests: { type: "array", items: { type: "string" } },
  },
} as const;

const VERIFICATION_SCHEMA = {
  type: "object",
  additionalProperties: false,
  required: ["succeeded", "focusedBehavioralTests", "summary"],
  properties: {
    succeeded: { type: "boolean" },
    focusedBehavioralTests: { type: "boolean" },
    summary: { type: "string" },
  },
} as const;

const REVIEW_SCHEMA = {
  type: "object",
  additionalProperties: false,
  required: ["specification", "standards"],
  properties: {
    specification: { $ref: "#/$defs/axis" },
    standards: { $ref: "#/$defs/axis" },
  },
  $defs: {
    axis: {
      type: "object",
      additionalProperties: false,
      required: ["verdict", "findings"],
      properties: {
        verdict: { enum: ["approved", "changes_requested"] },
        findings: {
          type: "array",
          items: {
            type: "object",
            additionalProperties: false,
            required: ["severity", "message"],
            properties: {
              severity: { type: "string" },
              message: { type: "string" },
              path: { type: "string" },
            },
          },
        },
      },
    },
  },
} as const;

function schemaForRole(role: AssignmentRequest["role"]) {
  if (role === "discover") return DISCOVERY_SCHEMA;
  if (role === "design") return DESIGN_SCHEMA;
  if (role === "prototype") return PROTOTYPE_SCHEMA;
  if (role === "verify") return VERIFICATION_SCHEMA;
  if (role === "review") return REVIEW_SCHEMA;
  return IMPLEMENTATION_SCHEMA;
}

function pagedTextTool(options: {
  name: "workflow_diff" | "workflow_plan";
  label: string;
  description: string;
  contents: string;
}) {
  const characters = [...options.contents];
  return defineTool({
    name: options.name,
    label: options.label,
    description: options.description,
    parameters: Type.Object({
      offset: Type.Optional(Type.Integer({ minimum: 0 })),
      limit: Type.Optional(Type.Integer({ minimum: 1, maximum: 48 * 1024 })),
    }),
    async execute(_id, params) {
      const offset = params.offset ?? 0;
      const limit = params.limit ?? 48 * 1024;
      return {
        content: [
          {
            type: "text" as const,
            text: characters.slice(offset, offset + limit).join(""),
          },
        ],
        details: { offset, limit, totalCharacters: characters.length },
      };
    },
  });
}

function readCompletePagedText(
  transcript: AssignmentExecutionMetadata["transcript"],
  toolName: "workflow_diff" | "workflow_plan",
  totalCharacters: number,
) {
  if (totalCharacters === 0) return true;
  const entries = transcript.filter(
    (
      entry,
    ): entry is {
      role: string;
      name?: string;
      text: string;
      toolCallId?: string;
      isError?: boolean;
    } =>
      typeof entry === "object" &&
      entry !== null &&
      "role" in entry &&
      typeof entry.role === "string" &&
      "text" in entry &&
      typeof entry.text === "string",
  );
  const successfulCallIds = new Set(
    entries
      .filter(
        (entry) =>
          entry.role === "toolResult" &&
          entry.name === toolName &&
          !entry.isError &&
          typeof entry.toolCallId === "string",
      )
      .map((entry) => entry.toolCallId!),
  );
  const intervals = entries
    .filter(
      (entry) =>
        entry.role === "tool" &&
        typeof entry.toolCallId === "string" &&
        successfulCallIds.has(entry.toolCallId) &&
        entry.name === toolName,
    )
    .map((entry) => {
      try {
        const input = JSON.parse(entry.text) as { offset?: unknown; limit?: unknown };
        const offset = input.offset === undefined ? 0 : input.offset;
        const limit = input.limit === undefined ? 48 * 1024 : input.limit;
        if (
          typeof offset !== "number" ||
          !Number.isInteger(offset) ||
          offset < 0 ||
          typeof limit !== "number" ||
          !Number.isInteger(limit) ||
          limit < 1 ||
          limit > 48 * 1024
        ) {
          return [0, 0] as const;
        }
        return [offset, Math.min(totalCharacters, offset + limit)] as const;
      } catch {
        return [0, 0] as const;
      }
    })
    .sort((left, right) => left[0] - right[0]);
  let covered = 0;
  for (const [start, end] of intervals) {
    if (start > covered) return false;
    covered = Math.max(covered, end);
    if (covered >= totalCharacters) return true;
  }
  return false;
}

function assignmentPrompt(request: AssignmentRequest) {
  const plannerCandidates = request.plannerCandidates?.map((entry) => ({
    role: entry.role,
    candidates: entry.candidates
      .filter((candidate) => candidate.eligible)
      .map((candidate) => ({
        name: candidate.profileName,
        digest: candidate.profileDigest,
        specificity: candidate.specificity,
      })),
  }));
  return [
    `Assignment role: ${request.role}`,
    `Workflow request: ${request.request}`,
    `Repository facts: ${JSON.stringify(request.facts)}`,
    `Accepted decisions: ${JSON.stringify(request.decisions.map((item) => item.value))}`,
    ...(plannerCandidates
      ? [`Eligible planner candidates: ${JSON.stringify(plannerCandidates)}`]
      : []),
    ...(request.approvedPlan
      ? request.role === "verify" || request.role === "review"
        ? [
            `Approved plan: ${[...request.approvedPlan].length} characters available through workflow_plan. Read the complete plan.`,
          ]
        : [`Approved plan:\n${request.approvedPlan}`]
      : []),
    ...(request.verificationEvidence
      ? [
          `Persisted verification evidence: ${JSON.stringify(request.verificationEvidence)}`,
        ]
      : []),
    ...(request.reviewDiff
      ? [
          `Review diff: ${[...request.reviewDiff].length} characters available through workflow_diff. Review the complete scope.`,
        ]
      : []),
    "Follow the active Agent Profile instructions. Return only through structured_output.",
  ].join("\n\n");
}

interface ChildResources {
  loader: DefaultResourceLoader;
  settingsManager: SettingsManager;
}

export interface PiAssignmentExecutorOptions {
  modelRegistry: ExtensionContext["modelRegistry"];
  parentCwd: string;
  parentProjectTrusted: boolean;
  parentModel?: WorkflowModel;
  parentThinkingLevel: ThinkingLevel;
  concurrencyCeiling: number;
  runConcurrency: number;
  scheduler?: ProcessAssignmentScheduler;
  runAgent?: (options: RunAgentOptions) => Promise<AgentOutcome>;
  createResources?: (options: {
    cwd: string;
    projectTrusted: boolean;
    appendSystemPrompt: string[];
  }) => Promise<ChildResources>;
  resolveProjectTrust?: typeof resolveStandaloneChildProjectTrust;
}

export class PiAssignmentExecutor implements AssignmentExecutor {
  readonly reservesCapacity = true;
  private readonly options: PiAssignmentExecutorOptions;
  private readonly scheduler: ProcessAssignmentScheduler;

  constructor(options: PiAssignmentExecutorOptions) {
    this.options = options;
    this.scheduler = options.scheduler ?? processAssignmentScheduler;
    this.scheduler.lowerCeiling(options.concurrencyCeiling);
  }

  async run(request: AssignmentRequest): Promise<AssignmentExecutionResult> {
    const model = resolveWorkflowModel(
      this.options.modelRegistry,
      request.profile,
      this.options.parentModel,
    ) as WorkflowModel | undefined;
    const thinkingLevel =
      (request.profile.effort as ThinkingLevel | undefined) ??
      this.options.parentThinkingLevel;
    const resolveTrust =
      this.options.resolveProjectTrust ?? resolveStandaloneChildProjectTrust;
    const projectTrusted = resolveTrust({
      parentCwd: this.options.parentCwd,
      childCwd: request.cwd,
      parentTrusted: this.options.parentProjectTrusted,
    });
    const createResources = this.options.createResources ?? createChildResources;
    const resources = await createResources({
      cwd: request.cwd,
      projectTrusted,
      appendSystemPrompt: [
        request.profile.instructions,
        STRUCTURED_OUTPUT_SYSTEM_INSTRUCTION,
      ],
    });
    const execute = this.options.runAgent ?? runAgent;
    const mutationScope = request.mutates
      ? await MutationScope.create(request.cwd, request.ownedPaths ?? [])
      : undefined;
    const gitProbe =
      mutationScope || request.captureGitState ? new GitStateProbe() : undefined;
    const evidenceTools = [
      ...(request.reviewDiff
        ? [
            pagedTextTool({
              name: "workflow_diff",
              label: "Workflow Diff",
              description:
                "Read the complete review diff in bounded character ranges.",
              contents: request.reviewDiff,
            }),
          ]
        : []),
      ...(request.approvedPlan &&
      (request.role === "verify" || request.role === "review")
        ? [
            pagedTextTool({
              name: "workflow_plan",
              label: "Workflow Plan",
              description:
                "Read the complete approved plan in bounded character ranges.",
              contents: request.approvedPlan,
            }),
          ]
        : []),
    ];
    const allowedReadTools = new Set(
      request.allowedReadTools ?? ["read", "grep", "find", "ls"],
    );
    const auditedReadTools = [
      createReadToolDefinition(request.cwd),
      createGrepToolDefinition(request.cwd),
      createFindToolDefinition(request.cwd),
      createLsToolDefinition(request.cwd),
    ].filter((tool) => allowedReadTools.has(tool.name));
    const mutationTools = mutationScope
      ? [
          createEditToolDefinition(request.cwd, {
            operations: {
              async readFile(absolutePath) {
                return readFile(await mutationScope.assertWritePath(absolutePath));
              },
              async writeFile(absolutePath, content) {
                await mutationScope.writeFile(absolutePath, content);
              },
              async access(absolutePath) {
                await access(await mutationScope.assertWritePath(absolutePath));
              },
            },
          }),
          createWriteToolDefinition(request.cwd, {
            operations: {
              async writeFile(absolutePath, content) {
                await mutationScope.writeFile(absolutePath, content);
              },
              async mkdir(directory) {
                await mutationScope.mkdir(directory);
              },
            },
          }),
        ]
      : [];
    const removalTools =
      mutationScope && request.role === "remediate"
        ? [
            defineTool({
              name: "workflow_remove",
              label: "Remove Owned Prototype Path",
              description:
                "Remove a prototype-only file or directory inside Assignment ownership.",
              parameters: Type.Object({ path: Type.String() }),
              async execute(_id, params) {
                await mutationScope.remove(params.path);
                return {
                  content: [
                    { type: "text" as const, text: `Removed ${params.path}` },
                  ],
                  details: { path: params.path },
                };
              },
            }),
          ]
        : [];
    return this.scheduler.schedule({
      runId: request.runId,
      runLimit: request.runConcurrency ?? this.options.runConcurrency,
      signal: request.signal,
      exclusive: request.mutates,
      onReserved: request.onReserved,
      task: async (signal) => {
        let created: AssignmentExecutionMetadata | undefined;
        const gitBefore = gitProbe
          ? await gitProbe.capture(mutationScope?.root ?? request.cwd)
          : undefined;
        const outcome = await execute({
          prompt: assignmentPrompt(request),
          schema: schemaForRole(request.role),
          model,
          thinkingLevel,
          cwd: request.cwd,
          loader: resources.loader,
          settingsManager: resources.settingsManager,
          modelRegistry: this.options.modelRegistry,
          signal,
          tools: request.mutates
            ? [
                ...(request.profile.tools ?? []).filter((tool) =>
                  (["edit", "write"].includes(tool) ||
                    allowedReadTools.has(tool)),
                ),
                ...evidenceTools.map((tool) => tool.name),
                ...removalTools.map((tool) => tool.name),
                "structured_output",
              ]
            : [
                ...(request.profile.tools ?? []).filter((tool) =>
                  allowedReadTools.has(tool),
                ),
                ...evidenceTools.map((tool) => tool.name),
                "structured_output",
              ],
          customTools: [
            ...auditedReadTools,
            ...mutationTools,
            ...removalTools,
            ...evidenceTools,
          ],
          onSessionCreated(metadata) {
            created = {
              childSessionId: metadata.childSessionId,
              provider: metadata.provider,
              model: metadata.model,
              effort: metadata.thinkingLevel,
              contextCapacity: undefined,
              transcript: [],
            };
          },
        });
        const gitAfter = gitProbe
          ? await gitProbe.capture(mutationScope?.root ?? request.cwd)
          : undefined;
        let safetyError: Error | undefined;
        if (
          request.role === "review" &&
          request.reviewDiff !== undefined &&
          !readCompletePagedText(
            outcome.transcript,
            "workflow_diff",
            [...request.reviewDiff].length,
          )
        ) {
          safetyError = new Error(
            "Review Assignment did not read the complete workflow diff",
          );
        }
        if (
          !safetyError &&
          request.approvedPlan !== undefined &&
          (request.role === "verify" || request.role === "review") &&
          !readCompletePagedText(
            outcome.transcript,
            "workflow_plan",
            [...request.approvedPlan].length,
          )
        ) {
          safetyError = new Error(
            "Assignment did not read the complete approved workflow plan",
          );
        }
        if (!safetyError && gitProbe && gitBefore && gitAfter) {
          try {
            gitProbe.assertIdentity(gitBefore, gitAfter);
            const touched = gitProbe.changedSince(gitBefore, gitAfter);
            const outOfScope = mutationScope
              ? touched.filter((changed) => !mutationScope.owns(changed))
              : touched;
            if (outOfScope.length > 0) {
              throw new Error(
                request.captureGitState && !mutationScope
                  ? `Read-only Assignment changed the repository: ${outOfScope.join(", ")}`
                  : `Assignment changed paths outside ownership: ${outOfScope.join(", ")}`,
              );
            }
          } catch (error) {
            safetyError = error instanceof Error ? error : new Error(String(error));
          }
        }
        const metadata: AssignmentExecutionMetadata = {
          childSessionId: outcome.childSessionId ?? created?.childSessionId,
          provider: outcome.provider ?? created?.provider,
          model: outcome.model ?? created?.model,
          effort: outcome.thinkingLevel ?? created?.effort ?? thinkingLevel,
          contextCapacity: outcome.contextWindow,
          transcript: outcome.transcript,
          gitBefore,
          gitAfter,
        };
        if (safetyError || !outcome.ok) {
          const error = new Error(
            safetyError?.message ?? outcome.error ?? "Pi Assignment failed",
          ) as AssignmentExecutionError;
          error.assignmentMetadata = metadata;
          throw error;
        }
        return {
          kind: "workflow-assignment-result",
          output: outcome.structured,
          metadata,
        };
      },
    });
  }
}