/**
 * Part N: role→provider routing under Advanced Mode.
 * advancedMode off → single operator client (today's path).
 */
import type {
  DebugSink,
  DrawOpts,
  FeatureCandidate,
  FeatureNode,
  FeatureProbe,
  FeatureSpec,
  GeneratedFile,
  LlmCallRecord,
  LlmClient,
  RepairInput,
  UsageSink,
} from './types';
import type { IRCaseGen, IRImport } from './irScaffold';
import {
  classifyTopology,
  getProviders,
  type RoleProvider,
  type SupervisorModes,
} from './providers';
import { MockLlm } from '../mock/mockLlm';

export type AgentRole = 'operator' | 'supervisor' | 'pool';

function tagSink(
  sink: UsageSink | undefined,
  role: AgentRole,
  provider: string,
  poolIdx?: number,
): UsageSink | undefined {
  if (!sink) return undefined;
  return record => sink({
    ...record,
    role: role === 'pool' ? 'pool' : role,
    provider,
    ...(poolIdx != null ? { poolIdx } : {}),
  });
}

async function buildClient(
  p: RoleProvider,
  onUsage?: UsageSink,
  onDebug?: DebugSink,
  role: AgentRole = 'operator',
  poolIdx?: number,
): Promise<LlmClient> {
  const sink = tagSink(onUsage, role, p.type, poolIdx);
  if (p.type === 'mock') return new MockLlm(sink, onDebug);
  if (p.type === 'openai-compat') {
    const { LocalLlmClient } = await import('./localLlm');
    return new LocalLlmClient(sink, onDebug, {
      baseURL: p.baseURL,
      apiKey: p.apiKey,
      model: p.model,
    });
  }
  const { ClaudeClient } = await import('./llm');
  return new ClaudeClient(sink, onDebug, { apiKey: p.apiKey, model: p.model });
}

async function envClient(onUsage?: UsageSink, onDebug?: DebugSink): Promise<LlmClient> {
  const sink = tagSink(onUsage, 'operator', process.env.MOCK_LLM === '1' ? 'mock'
    : process.env.LOCAL_LLM === '1' ? 'openai-compat' : 'claude');
  if (process.env.MOCK_LLM === '1') return new MockLlm(sink, onDebug);
  if (process.env.LOCAL_LLM === '1') {
    const { LocalLlmClient } = await import('./localLlm');
    return new LocalLlmClient(sink, onDebug);
  }
  const { ClaudeClient } = await import('./llm');
  return new ClaudeClient(sink, onDebug);
}

class RoutingLlm implements LlmClient {
  constructor(
    private operator: LlmClient,
    private operatorB: LlmClient | null,
    private supervisor: LlmClient | null,
    private modes: SupervisorModes,
    private cursor = 0,
  ) {}

  private forFrontEnd(): LlmClient {
    if (this.supervisor && this.modes.resolveIr) return this.supervisor;
    return this.operator;
  }

  private forBuild(): LlmClient {
    // Round-robin across operators when two exist (orchestrated partition uses
    // explicit clients via operatorAt; this covers swarm draws).
    if (this.operatorB) {
      const pick = this.cursor++ % 2 === 0 ? this.operator : this.operatorB;
      return pick;
    }
    return this.operator;
  }

  enumerateFeatures(prompt: string, opts?: DrawOpts) {
    return this.forFrontEnd().enumerateFeatures(prompt, opts);
  }
  probeFeature(candidate: FeatureCandidate, opts?: DrawOpts) {
    return this.forFrontEnd().probeFeature(candidate, opts);
  }
  encodeFeatures(prompt: string, stubFeedback?: string, opts?: DrawOpts) {
    return this.forFrontEnd().encodeFeatures(prompt, stubFeedback, opts);
  }
  repairFeatures(spec: FeatureSpec, offendingIds: string[], problems: string[], opts?: DrawOpts) {
    return this.forFrontEnd().repairFeatures(spec, offendingIds, problems, opts);
  }
  fixStub(
    path: string,
    ctx: { diagnostics: string; callingLines?: string; contractText?: string; checkedIdioms?: string },
    opts?: DrawOpts,
  ) {
    return this.forFrontEnd().fixStub(path, ctx, opts);
  }
  generateTests(spec: FeatureSpec, feature: FeatureNode, feedback?: string, opts?: DrawOpts) {
    return this.forFrontEnd().generateTests(spec, feature, feedback, opts);
  }
  fixTestCases(
    spec: FeatureSpec,
    feature: FeatureNode,
    brokenCases: { name: string; kind: IRCaseGen['kind']; errors: string[] }[],
    importErrors: string[],
    currentImports: IRImport[],
    opts?: DrawOpts,
  ) {
    return this.forFrontEnd().fixTestCases(spec, feature, brokenCases, importErrors, currentImports, opts);
  }
  generateImplementation(
    spec: FeatureSpec,
    feature: FeatureNode,
    frozenTest: string,
    dependencyContext: GeneratedFile[],
    opts?: DrawOpts,
  ) {
    return this.forBuild().generateImplementation(spec, feature, frozenTest, dependencyContext, opts);
  }
  repair(input: RepairInput, opts?: DrawOpts) {
    return this.operator.repair(input, opts);
  }

  get supervisorClient(): LlmClient | null {
    return this.supervisor;
  }

  operatorAt(i: number): LlmClient {
    return i === 1 && this.operatorB ? this.operatorB : this.operator;
  }

  get modesSnapshot(): SupervisorModes {
    return { ...this.modes };
  }
}

export interface LlmBundle {
  client: LlmClient;
  supervisor: LlmClient | null;
  operators: LlmClient[];
  modes: SupervisorModes;
  topology: ReturnType<typeof classifyTopology>;
}

export async function makeLlmBundle(
  onUsage?: UsageSink,
  onDebug?: DebugSink,
): Promise<LlmBundle> {
  const providers = getProviders();
  const topo = classifyTopology(providers);

  // Solo / advanced off: one client
  if (!providers.advancedMode || topo.kind === 'solo' || topo.kind === 'invalid') {
    const op = providers.operators[0];
    const client = op?.connected
      ? await buildClient(op, onUsage, onDebug, 'operator', 0)
      : await envClient(onUsage, onDebug);
    return {
      client,
      supervisor: null,
      operators: [client],
      modes: { resolveIr: false, resolveAllFailures: false, resolveEscalationsOnly: false },
      topology: topo.kind === 'invalid' ? { kind: 'solo' } : topo,
    };
  }

  const opClients: LlmClient[] = [];
  for (let i = 0; i < providers.operators.length; i++) {
    const o = providers.operators[i];
    if (o?.connected) opClients.push(await buildClient(o, onUsage, onDebug, 'operator', i));
  }
  if (!opClients.length) opClients.push(await envClient(onUsage, onDebug));

  let supervisor: LlmClient | null = null;
  const modes = providers.supervisorModes;
  if (providers.supervisor?.connected
    && (modes.resolveIr || modes.resolveEscalationsOnly || modes.resolveAllFailures)) {
    supervisor = await buildClient(providers.supervisor, onUsage, onDebug, 'supervisor');
  }

  const client = (supervisor || opClients.length > 1)
    ? new RoutingLlm(opClients[0], opClients[1] ?? null, supervisor, modes)
    : opClients[0];

  return { client, supervisor, operators: opClients, modes, topology: topo };
}

export function isRoutingLlm(c: LlmClient): c is RoutingLlm {
  return c instanceof RoutingLlm;
}

/**
 * Partition feature ids into two disjoint topo-respecting halves for the
 * Orchestrator (N5). Roots alternate; dependents follow their root's owner.
 */
export function partitionFeatureIds(
  features: { id: string; deps: string[]; parentId: string | null }[],
): [string[], string[]] {
  const byId = new Map(features.map(f => [f.id, f]));
  const roots = features.filter(f => !f.parentId && f.deps.every(d => !byId.has(d) || byId.get(d)!.parentId));
  // Prefer features with no deps as partition seeds
  const seeds = features.filter(f => f.deps.length === 0);
  const use = seeds.length ? seeds : features;
  const owner = new Map<string, 0 | 1>();
  use.forEach((f, i) => owner.set(f.id, (i % 2) as 0 | 1));
  // Propagate: a feature joins the partition of its first dep if any owned
  for (const f of features) {
    if (owner.has(f.id)) continue;
    const depOwner = f.deps.map(d => owner.get(d)).find(x => x != null);
    owner.set(f.id, (depOwner ?? (owner.size % 2)) as 0 | 1);
  }
  const a: string[] = [];
  const b: string[] = [];
  for (const f of features) {
    (owner.get(f.id) === 0 ? a : b).push(f.id);
  }
  if (!a.length && b.length) a.push(b.shift()!);
  if (!b.length && a.length > 1) b.push(a.pop()!);
  return [a, b];
}
