import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
import {
  isBackTab,
  isBackspace,
  isCharacter,
  isDown,
  isEnter,
  isEscape,
  isLeft,
  isRight,
  isTab,
  isUp,
  printableText,
} from "./tui-keys.js";
import { borderedLines, OVERLAY_OPTIONS, truncate } from "./tui-overlay.js";

import type { WorkerRolePreset } from "../workflow/worker-role-presets.js";
import { workerRoleContract } from "../workflow/worker-role-presets.js";
import type {
  EditedWorkerProposal,
  LaunchPreviewDecision,
} from "../workflow/spawn-confirmation.js";
import type {
  ModelSpec,
  WorkerBudget,
  WorkerProposalEditableField,
  WorkerRequest,
} from "../workflow/workflow-types.js";

export interface LaunchPreviewRoleOption {
  readonly preset: WorkerRolePreset;
  readonly defaultModel: ModelSpec;
}

export interface LaunchPreviewModelOptions {
  readonly requests: readonly WorkerRequest[];
  readonly roles: readonly LaunchPreviewRoleOption[];
  readonly models: readonly ModelSpec[];
  readonly difficultyModels: Readonly<Record<"easy" | "medium" | "hard", ModelSpec>>;
  readonly normalizeModel?: (requestKey: string, roleName: string, model: ModelSpec) => ModelSpec;
}

interface EditableProposal {
  request: WorkerRequest;
  readonly originalBudget: WorkerBudget;
  readonly edited: Set<WorkerProposalEditableField>;
}

export type BudgetDimension = keyof WorkerBudget;

const EDITABLE_FIELDS = ["role", "task", "model", "difficulty", "budget"] as const;
const BUDGET_DIMENSIONS = [
  "maxAttempts",
  "maxWallClockMs",
  "maxTokens",
  "maxCostUsd",
  "maxOutputBytes",
] as const satisfies readonly BudgetDimension[];
const DIFFICULTIES = ["easy", "medium", "hard"] as const;

export class LaunchPreviewModel {
  readonly #proposals: EditableProposal[];
  readonly #roles: readonly LaunchPreviewRoleOption[];
  readonly #models: readonly ModelSpec[];
  readonly #difficultyModels: LaunchPreviewModelOptions["difficultyModels"];
  readonly #normalizeModel: (requestKey: string, roleName: string, model: ModelSpec) => ModelSpec;

  constructor(options: LaunchPreviewModelOptions) {
    if (options.requests.length === 0) throw new TypeError("Launch preview requires a proposal");
    if (options.roles.length === 0) throw new TypeError("Launch preview requires role options");
    if (options.models.length === 0) throw new TypeError("Launch preview requires live model options");
    this.#proposals = options.requests.map((request) => ({
      request: cloneRequest(request),
      originalBudget: cloneBudget(request.budget),
      edited: new Set(),
    }));
    this.#roles = [...options.roles];
    this.#models = [...options.models];
    this.#difficultyModels = options.difficultyModels;
    this.#normalizeModel = options.normalizeModel ?? ((_requestKey, _roleName, model) => model);
  }

  get count(): number {
    return this.#proposals.length;
  }

  request(index: number): WorkerRequest {
    return cloneRequest(this.#entry(index).request);
  }

  editedFields(index: number): readonly WorkerProposalEditableField[] {
    const edited = this.#entry(index).edited;
    return EDITABLE_FIELDS.filter((field) => edited.has(field));
  }

  cycleRole(index: number, delta: -1 | 1): void {
    const entry = this.#entry(index);
    const current = this.#roles.findIndex((option) => option.preset.name === entry.request.role?.name);
    const selected = this.#roles[wrap(current < 0 ? 0 : current, delta, this.#roles.length)]!;
    entry.request = {
      ...entry.request,
      role: workerRoleContract(selected.preset),
      difficulty: selected.preset.difficulty,
      capabilities: selected.preset.capabilities,
      model: this.#normalizeModel(entry.request.requestKey, selected.preset.name, selected.defaultModel),
    };
    entry.edited.add("role");
  }

  cycleModel(index: number, delta: -1 | 1): void {
    const entry = this.#entry(index);
    const current = this.#models.findIndex((model) => sameModel(model, entry.request.model));
    entry.request = {
      ...entry.request,
      model: this.#normalizeModel(
        entry.request.requestKey,
        entry.request.role?.name ?? "worker",
        this.#models[wrap(current < 0 ? 0 : current, delta, this.#models.length)]!,
      ),
    };
    entry.edited.add("model");
  }

  cycleDifficulty(index: number, delta: -1 | 1): void {
    const entry = this.#entry(index);
    const current = DIFFICULTIES.indexOf(entry.request.difficulty ?? "medium");
    const difficulty = DIFFICULTIES[wrap(current, delta, DIFFICULTIES.length)]!;
    entry.request = {
      ...entry.request,
      difficulty,
      model: this.#normalizeModel(
        entry.request.requestKey,
        entry.request.role?.name ?? "worker",
        this.#difficultyModels[difficulty],
      ),
    };
    entry.edited.add("difficulty");
  }

  replaceTask(index: number, goal: string): void {
    const normalized = goal.trim();
    if (normalized.length === 0) throw new TypeError("Worker task must not be empty");
    const entry = this.#entry(index);
    entry.request = { ...entry.request, goal: normalized };
    entry.edited.add("task");
  }

  lowerBudget(index: number, dimension: BudgetDimension, value: number): void {
    if (!Number.isFinite(value) || value <= 0) throw new TypeError("Budget value must be positive");
    if (dimension !== "maxCostUsd" && !Number.isSafeInteger(value)) {
      throw new TypeError(`${dimension} must be a positive integer`);
    }
    const entry = this.#entry(index);
    const original = entry.originalBudget[dimension];
    if (original.kind === "limited" && value > original.value) {
      throw new TypeError(`${dimension} may only stay at or below its proposed value ${original.value}`);
    }
    entry.request = {
      ...entry.request,
      budget: { ...entry.request.budget, [dimension]: { kind: "limited", value } },
    };
    entry.edited.add("budget");
  }

  accepted(): LaunchPreviewDecision {
    return {
      kind: "accepted",
      proposals: this.#proposals.map((entry): EditedWorkerProposal => ({
        request: cloneRequest(entry.request),
        humanEditedFields: EDITABLE_FIELDS.filter((field) => entry.edited.has(field)),
      })),
    };
  }

  #entry(index: number): EditableProposal {
    const entry = this.#proposals[index];
    if (!entry) throw new RangeError(`Unknown proposal index ${index}`);
    return entry;
  }
}

export async function openLaunchPreview(
  ctx: ExtensionContext,
  options: LaunchPreviewModelOptions,
): Promise<LaunchPreviewDecision> {
  if (ctx.mode !== "tui" || !ctx.hasUI) {
    throw new Error("Launch preview requires an interactive TUI");
  }
  const model = new LaunchPreviewModel(options);
  return ctx.ui.custom<LaunchPreviewDecision>((tui, _theme, _keys, done) => {
    return new LaunchPreviewComponent(model, done, () => tui.requestRender());
  }, {
    overlay: true,
    overlayOptions: OVERLAY_OPTIONS,
    onHandle: (handle) => handle.focus(),
  });
}

class LaunchPreviewComponent {
  readonly #model: LaunchPreviewModel;
  readonly #done: (decision: LaunchPreviewDecision) => void;
  readonly #invalidate: () => void;
  #worker = 0;
  #field = 0;
  #budgetDimension = 0;
  #editing: "task" | "budget" | undefined;
  #buffer = "";
  #diagnostic = "";

  constructor(
    model: LaunchPreviewModel,
    done: (decision: LaunchPreviewDecision) => void,
    invalidate: () => void,
  ) {
    this.#model = model;
    this.#done = done;
    this.#invalidate = invalidate;
  }

  handleInput(data: string): void {
    if (isEscape(data)) {
      this.#done({ kind: "cancelled" });
      return;
    }
    if (this.#editing) {
      if (isEnter(data)) {
        try {
          if (this.#editing === "task") this.#model.replaceTask(this.#worker, this.#buffer);
          else this.#model.lowerBudget(
            this.#worker,
            BUDGET_DIMENSIONS[this.#budgetDimension]!,
            Number(this.#buffer),
          );
          this.#editing = undefined;
          this.#buffer = "";
          this.#diagnostic = "";
        } catch (error) {
          this.#diagnostic = error instanceof Error ? error.message : String(error);
        }
        this.#invalidate();
        return;
      }
      if (isBackspace(data)) this.#buffer = this.#buffer.slice(0, -1);
      else {
        const text = printableText(data);
        if (text !== undefined) this.#buffer += text;
      }
      this.#invalidate();
      return;
    }
    if (isEnter(data)) {
      this.#done(this.#model.accepted());
      return;
    }
    if (isUp(data)) this.#worker = Math.max(0, this.#worker - 1);
    else if (isDown(data)) this.#worker = Math.min(this.#model.count - 1, this.#worker + 1);
    else if (isTab(data)) this.#field = wrap(this.#field, 1, EDITABLE_FIELDS.length);
    else if (isBackTab(data)) this.#field = wrap(this.#field, -1, EDITABLE_FIELDS.length);
    else if (isCharacter(data, "[") && EDITABLE_FIELDS[this.#field] === "budget") {
      this.#budgetDimension = wrap(this.#budgetDimension, -1, BUDGET_DIMENSIONS.length);
    } else if (isCharacter(data, "]") && EDITABLE_FIELDS[this.#field] === "budget") {
      this.#budgetDimension = wrap(this.#budgetDimension, 1, BUDGET_DIMENSIONS.length);
    } else if (isLeft(data)) this.#cycle(-1);
    else if (isRight(data)) this.#cycle(1);
    else if ((isCharacter(data, "e") || isCharacter(data, "E")) && (
      EDITABLE_FIELDS[this.#field] === "task" || EDITABLE_FIELDS[this.#field] === "budget"
    )) {
      this.#editing = EDITABLE_FIELDS[this.#field] as "task" | "budget";
      this.#buffer = "";
      this.#diagnostic = "Type a replacement value, then Enter. Esc cancels the entire launch.";
    }
    this.#invalidate();
  }

  render(width: number): string[] {
    const request = this.#model.request(this.#worker);
    const field = EDITABLE_FIELDS[this.#field]!;
    const dimension = BUDGET_DIMENSIONS[this.#budgetDimension]!;
    const lines = [
      "↑↓ worker · Tab field · ←→ cycle · [ ] budget dimension · e edit · Enter launch all · Esc cancel all",
      "",
      ...Array.from({ length: this.#model.count }, (_, index) => {
        const candidate = this.#model.request(index);
        const edited = this.#model.editedFields(index).length > 0 ? " edited" : "";
        return `${index === this.#worker ? ">" : " "} ${index + 1}. ${candidate.role?.name ?? "worker"} · ${candidate.goal}${edited}`;
      }),
      "",
      row("role", request.role?.name ?? "worker", field),
      row("task", request.goal, field),
      row("model", renderModel(request.model), field),
      ...(request.modelFallbacks && request.modelFallbacks.length > 0
        ? [`  fallback chain: ${request.modelFallbacks.map(renderModel).join(" → ")}`]
        : []),
      row("difficulty", request.difficulty ?? "medium", field),
      row(`budget.${budgetLabel(dimension)}`, renderLimit(request.budget[dimension]), field),
      `  context: ${request.context ?? "fresh"}${request.forkSource
        ? ` · ${request.forkSource.inheritedEntries} entries · ${request.forkSource.inheritedBytes} bytes · leaf ${request.forkSource.branchPointId}`
        : ""}`,
      `  capabilities: ${request.capabilities.join(", ") || "none"}`,
      ...peerGrantLines(request.capabilities),
      ...(this.#editing ? [`  ${this.#editing}> ${this.#buffer}_`] : []),
      ...(this.#diagnostic ? [`  ${this.#diagnostic}`] : []),
    ];
    return borderedLines(
      `Launch preview · ${this.#model.count} proposed worker${this.#model.count === 1 ? "" : "s"}`,
      lines,
      Math.max(20, width),
    );
  }

  invalidate(): void {}

  #cycle(delta: -1 | 1): void {
    switch (EDITABLE_FIELDS[this.#field]) {
      case "role":
        this.#model.cycleRole(this.#worker, delta);
        return;
      case "model":
        this.#model.cycleModel(this.#worker, delta);
        return;
      case "difficulty":
        this.#model.cycleDifficulty(this.#worker, delta);
        return;
      case "task":
      case "budget":
        this.#diagnostic = "Press e to replace this value";
        return;
    }
  }
}

function cloneRequest(request: WorkerRequest): WorkerRequest {
  return {
    ...request,
    acceptance: [...request.acceptance],
    constraints: [...request.constraints],
    budget: cloneBudget(request.budget),
    model: { ...request.model },
    ...(request.modelFallbacks
      ? { modelFallbacks: request.modelFallbacks.map((model) => ({ ...model })) }
      : {}),
    capabilities: [...request.capabilities],
    workspace: { ...request.workspace },
    contextRefs: request.contextRefs.map((reference) => ({ ...reference })),
    ...(request.role ? { role: { ...request.role } } : {}),
    ...(request.proposal ? {
      proposal: {
        ...request.proposal,
        humanEditedFields: [...request.proposal.humanEditedFields],
      },
    } : {}),
  };
}

function cloneBudget(budget: WorkerBudget): WorkerBudget {
  return {
    maxAttempts: { ...budget.maxAttempts },
    maxWallClockMs: { ...budget.maxWallClockMs },
    maxTokens: { ...budget.maxTokens },
    maxCostUsd: { ...budget.maxCostUsd },
    maxOutputBytes: { ...budget.maxOutputBytes },
  };
}

function renderModel(model: ModelSpec): string {
  return `${model.provider}/${model.id}${model.thinkingLevel ? ` (${model.thinkingLevel})` : ""}`;
}

function peerGrantLines(capabilities: readonly string[]): readonly string[] {
  const grants = capabilities.filter((capability) => capability.startsWith("message_worker:"));
  return grants.length === 0 ? [] : [`  peer grants: ${grants.join(", ")}`];
}

function sameModel(left: ModelSpec, right: ModelSpec): boolean {
  return left.provider === right.provider
    && left.id === right.id;
}

function renderLimit(limit: WorkerBudget[BudgetDimension]): string {
  return limit.kind === "unlimited" ? "unlimited" : String(limit.value);
}

function budgetLabel(dimension: BudgetDimension): string {
  return dimension.replace(/[A-Z]/g, (letter) => `_${letter.toLowerCase()}`);
}

function row(label: string, value: string, selected: string): string {
  return `${selected === label.split(".")[0] ? ">" : " "} ${label}: ${value}`;
}


function wrap(index: number, delta: -1 | 1, length: number): number {
  return (index + delta + length) % length;
}