import { logLlmDiagnostics } from "../utils/tokens";
import { LMStudioClient } from "@lmstudio/sdk";
import {
  QueryPlan,
  WorkerRole,
  DynamicWorkerSpec,
  AdaptiveGapPlan,
  CrawledSource,
  AgentMessage,
  StatusFn,
} from "../types";
import { DIMENSIONS, detectGaps, gapFillQueries } from "./dimensions";
import {
  DepthProfile,
  AI_PLANNING_MAX_TOKENS,
  AI_PLANNING_TEMPERATURE,
  AI_PLANNING_TIMEOUT_MS,
  AI_MIN_ACCEPTABLE_QUERIES,
  AI_DECOMPOSITION_MAX_TOKENS,
  AI_DECOMPOSITION_TEMPERATURE,
  AI_DECOMPOSITION_TIMEOUT_MS,
  AI_FINDINGS_SUMMARY_MAX_TOKENS,
  AI_FINDINGS_SUMMARY_TEMPERATURE,
  FINDINGS_SUMMARY_SOURCE_CHARS,
  DECOMPOSITION_MIN_WORKERS,
  QUERY_LINE_MIN_LEN,
  QUERY_LINE_MAX_LEN,
  SYSTEM_INSTRUCTIONS,
} from "../constants";

async function callLoadedModel(
  prompt: string,
  maxTokens: number = AI_PLANNING_MAX_TOKENS,
  temperature: number = AI_PLANNING_TEMPERATURE,
  timeoutMs: number = AI_PLANNING_TIMEOUT_MS,
): Promise<string | null> {
  try {
    const client = new LMStudioClient();
    const models = await Promise.race<Awaited<ReturnType<typeof client.llm.listLoaded>>>([
      client.llm.listLoaded(),
      new Promise<never>((_, reject) =>
        setTimeout(() => reject(new Error("timeout")), timeoutMs),
      ),
    ]);
    if (!Array.isArray(models) || models.length === 0) return null;
    const model = await client.llm.model(models[0].identifier);
    const stream = model.respond(
      [
        { role: "system", content: SYSTEM_INSTRUCTIONS },
        { role: "user", content: prompt },
      ],
      {
        maxTokens,
        temperature,
      },
    );
    let result = "";
    for await (const chunk of stream) result += chunk.content ?? "";
    return result.trim() || null;
  } catch {
    return null;
  }
}

function parseLines(raw: string, maxLines: number = 6): ReadonlyArray<string> {
  return raw
    .split(/\n/)
    .map((line) => line.replace(/^\d+[.)]\s*|^[-*•]\s*/, "").trim())
    .filter(
      (line) =>
        line.length > QUERY_LINE_MIN_LEN && line.length < QUERY_LINE_MAX_LEN,
    )
    .filter((line, idx, arr) => arr.indexOf(line) === idx)
    .slice(0, maxLines);
}

const VALID_ROLES: ReadonlyArray<WorkerRole> = [
  "breadth",
  "depth",
  "recency",
  "academic",
  "critical",
  "statistical",
  "regulatory",
  "technical",
  "primary",
  "comparative",
];

function makeDecompositionPrompt(
  topic: string,
  focusAreas: ReadonlyArray<string>,
  profile: DepthProfile,
): string {
  const focus = focusAreas.length
    ? `\nFocus areas: ${focusAreas.join(", ")}`
    : "";
  return `You are a research decomposition system. Given a research topic, output a JSON array of specialized worker agents.
Topic: "${topic}"${focus}
Each worker needs:
"role": one of "breadth", "depth", "recency", "academic", "critical", "statistical", "regulatory", "technical", "primary", "comparative"
"label": descriptive name (e.g., "Clinical Evidence Researcher", "Policy Critic")
"queries": array of ${Math.min(profile.maxQueriesPerWorker, 6)}-${profile.maxQueriesPerWorker} specific search queries for this worker
"budgetWeight": number 0.1-0.4 (must sum to ~1.0 across all workers)
"followLinks": true/false (true for depth/academic workers)
"preferredTiers": optional array of "academic","government","reference","news","professional","general"
Rules:
Output ${DECOMPOSITION_MIN_WORKERS} to ${profile.maxDecompositionWorkers} workers
Tailor the workers to THIS specific topic - not generic roles
Queries must be highly specific to the topic and each worker's assignment
Generate MORE queries for broader or more complex topics
Budget weights must roughly sum to 1.0
Output ONLY valid JSON, no other text
JSON:`;
}

async function aiDecompose(
  topic: string,
  focusAreas: ReadonlyArray<string>,
  status: StatusFn,
  profile: DepthProfile,
): Promise<ReadonlyArray<DynamicWorkerSpec> | null> {
  const raw = await callLoadedModel(
    makeDecompositionPrompt(topic, focusAreas, profile),
    AI_DECOMPOSITION_MAX_TOKENS,
    AI_DECOMPOSITION_TEMPERATURE,
    AI_DECOMPOSITION_TIMEOUT_MS,
  );
  if (!raw) return null;
  try {
    const jsonStr = raw.replace(/```json\s*|```\s*/g, "").trim();
    const parsed = JSON.parse(jsonStr);
    if (!Array.isArray(parsed) || parsed.length < DECOMPOSITION_MIN_WORKERS)
      return null;
    const specs: DynamicWorkerSpec[] = [];
    for (const item of parsed.slice(0, profile.maxDecompositionWorkers)) {
      const role = VALID_ROLES.includes(item.role) ? item.role : "breadth";
      const queries = Array.isArray(item.queries)
        ? item.queries
            .filter((q: unknown) => typeof q === "string" && q.length > 3)
            .slice(0, profile.maxQueriesPerWorker)
        : [];
      if (queries.length < 2) continue;
      specs.push({
        role: role as WorkerRole,
        label:
          typeof item.label === "string"
            ? item.label.slice(0, 60)
            : `${role} worker`,
        queries,
        budgetWeight:
          typeof item.budgetWeight === "number"
            ? Math.max(0.05, Math.min(0.5, item.budgetWeight))
            : 0.2,
        followLinks: item.followLinks === true,
        preferredTiers: Array.isArray(item.preferredTiers)
          ? item.preferredTiers
          : undefined,
      });
    }
    if (specs.length < DECOMPOSITION_MIN_WORKERS) return null;
    const totalWeight = specs.reduce((sum, s) => sum + s.budgetWeight, 0);
    const normalised = specs.map((s) => ({
      ...s,
      budgetWeight: s.budgetWeight / totalWeight,
    }));
    status(`AI decomposed topic into ${normalised.length} specialised workers`);
    return normalised;
  } catch {
    return null;
  }
}

function makeRolePlanPrompt(
  role: WorkerRole,
  topic: string,
  focusAreas: ReadonlyArray<string>,
  profile: DepthProfile,
): string {
  const roleDescriptions: Readonly<Record<WorkerRole, string>> = {
    breadth: "broad coverage - many different angles, facts, and sub-topics",
    depth: "deep dive - mechanisms, how it works, technical detail, and evidence",
    recency: "recent developments - 2024-2026+ news, updates, and latest research",
    academic: "academic and scientific sources - peer-reviewed studies, journals, authoritative papers",
    critical: "critical analysis - limitations, counterarguments, criticism, controversy, drawbacks",
    statistical: "statistics and data - numbers, percentages, datasets, surveys, market sizes, quantitative evidence",
    regulatory: "regulatory and policy - laws, regulations, government policies, compliance, standards, guidelines",
    technical: "technical deep-dive - implementation details, specifications, architecture, engineering approaches",
    primary: "primary sources - original reports, official statements, first-hand accounts, press releases, white papers",
    comparative: "comparative analysis - vs alternatives, head-to-head comparisons, benchmarks, trade-offs, pros and cons",
  };
  const focus = focusAreas.length
    ? `\nFocus especially on: ${focusAreas.join(", ")}`
    : "";
  return `You are a research planning assistant. Generate search queries for a specialised research agent.
Topic: "${topic}"${focus}
This agent's role: ${roleDescriptions[role]}
Generate exactly ${profile.maxQueriesPerWorker} highly specific, diverse search queries for this role.
Rules:
Each query must be different from the others
Use natural language (as a human would type into a search engine)
Be specific to the role - ${roleDescriptions[role]}
Vary query structure: some factual, some comparative, some recent
Return ONLY the queries, one per line, no numbering, no extra text
Queries:`;
}

const ROLE_DIMENSIONS: Readonly<Record<WorkerRole, ReadonlyArray<string>>> = {
  breadth: ["overview", "applications", "history", "economics"],
  depth: ["mechanism", "evidence", "expert"],
  recency: ["current", "future"],
  academic: ["evidence", "expert", "mechanism"],
  critical: ["challenges", "controversy", "comparison"],
  statistical: ["evidence", "economics", "overview"],
  regulatory: ["challenges", "controversy", "current"],
  technical: ["mechanism", "applications", "evidence"],
  primary: ["evidence", "expert", "history"],
  comparative: ["comparison", "challenges", "applications"],
};

function dimensionFallbackQueries(
  role: WorkerRole,
  topic: string,
  focusAreas: ReadonlyArray<string>,
  maxQueries: number,
): ReadonlyArray<string> {
  const dimIds = ROLE_DIMENSIONS[role];
  const dims = DIMENSIONS.filter((d) => dimIds.includes(d.id));
  const queries: string[] = [];
  const shortTopic = topic;
  for (const dim of dims) {
    for (const q of dim.queries(shortTopic)) {
      if (!queries.includes(q)) queries.push(q);
    }
  }
  for (const area of focusAreas) {
    const q = `${shortTopic} ${area}`;
    if (!queries.includes(q)) queries.push(q);
  }
  return queries.slice(0, maxQueries);
}

export async function buildQueryPlan(
  topic: string,
  focusAreas: ReadonlyArray<string>,
  useAI: boolean,
  status: StatusFn,
  profile: DepthProfile,
): Promise<QueryPlan> {
  const CORE_ROLES: ReadonlyArray<WorkerRole> = [
    "breadth",
    "depth",
    "recency",
    "academic",
    "critical",
  ];
  const EXTENDED_ROLES: ReadonlyArray<WorkerRole> = [
    "statistical",
    "regulatory",
    "technical",
    "primary",
    "comparative",
  ];
  let roles: ReadonlyArray<WorkerRole>;
  if (profile.depthRounds >= 10) {
    roles = [...CORE_ROLES, ...EXTENDED_ROLES];
  } else if (profile.depthRounds >= 5) {
    roles = [...CORE_ROLES, "technical", "comparative", "statistical"];
  } else {
    roles = CORE_ROLES;
  }
  const queriesByRole: Partial<Record<WorkerRole, ReadonlyArray<string>>> = {};
  let usedAI = false;
  let dynamicSpecs: ReadonlyArray<DynamicWorkerSpec> | undefined;
  if (useAI) {
    status("AI task decomposition - analysing topic for specialised workers…");
    const specs = await aiDecompose(topic, focusAreas, status, profile);
    if (specs && specs.length >= DECOMPOSITION_MIN_WORKERS) {
      dynamicSpecs = specs;
      usedAI = true;
      for (const spec of specs) {
        queriesByRole[spec.role] = spec.queries;
      }
    } else {
      status("AI planning queries for each swarm worker…");
    }
    const uncoveredRoles = roles.filter((r) => !queriesByRole[r]?.length);
    if (uncoveredRoles.length > 0) {
      const results = await Promise.allSettled(
        uncoveredRoles.map(async (role) => ({
          role,
          queries: await callLoadedModel(
            makeRolePlanPrompt(role, topic, focusAreas, profile),
          ),
        })),
      );
      for (const result of results) {
        if (result.status !== "fulfilled") continue;
        const { role, queries: raw } = result.value;
        if (!raw) continue;
        const parsed = parseLines(raw, profile.maxQueriesPerWorker);
        if (parsed.length >= AI_MIN_ACCEPTABLE_QUERIES) {
          queriesByRole[role] = parsed;
          usedAI = true;
        }
      }
    }
    if (usedAI) {
      status(
        `AI generated queries for ${Object.keys(queriesByRole).length} worker role(s)`,
      );
    } else {
      status("AI unavailable, using dimension-based query planning");
    }
  }
  for (const role of roles) {
    if (!queriesByRole[role] || queriesByRole[role]!.length === 0) {
      queriesByRole[role] = dimensionFallbackQueries(
        role,
        topic,
        focusAreas,
        profile.maxQueriesPerWorker,
      );
    }
  }
  return {
    queriesByRole: queriesByRole as Record<WorkerRole, ReadonlyArray<string>>,
    usedAI,
    topicKeywords: extractKeywords(topic),
    dynamicSpecs,
  };
  const topicKeywords = extractKeywords(topic);

  // PRE-SEARCH PLANNING (If AI is enabled)
  if (useAI) {
    const planPrompt = `Topic: ${topic}\nFocus: ${focusAreas.join(", ")}\n
    Generate a JSON object with:
    1. "subQuestions": 3-5 specific questions to answer.
    2. "likelySourceTypes": Array of strings (e.g., "academic", "news", "reference").
    3. "stopConditions": Array of strings (e.g., "Found 3 sources confirming Q3 revenue").`;
    
    logLlmDiagnostics("buildQueryPlan-PreSearch", planPrompt);
    
    // Call LLM here and parse the JSON to inject into your existing query generation logic.
    // This prevents immediate blind fan-out.
  }

  return {
    queriesByRole: queriesByRole as Record<WorkerRole, ReadonlyArray<string>>,
    usedAI,
    topicKeywords,
    dynamicSpecs,
  };
} // <--- THIS BRACE CLOSES THE FUNCTION

export async function summariseFindings(
  sources: ReadonlyArray<CrawledSource>,
  topic: string,
  useAI: boolean,
  status: StatusFn,
): Promise<ReadonlyArray<AgentMessage>> {
  if (!useAI || sources.length === 0) return [];
  const sourceSummaries = sources
    .slice(0, 20)
    .map(
      (s, i) =>
        `[${i + 1}] ${s.workerLabel}: ${s.title} - ${s.text.slice(0, FINDINGS_SUMMARY_SOURCE_CHARS)}`,
    )
    .join("\n\n");
  const prompt = `You are a research coordinator. A team of research agents collected these sources on "${topic}":
${sourceSummaries}
Summarise:
The 3-5 most important findings discovered so far (one line each)
3-5 specific questions or angles that were NOT covered and need follow-up
Output format:
FINDINGS:
finding 1
finding 2
...
FOLLOW_UP:
question 1
question 2
...`;
  const raw = await callLoadedModel(
    prompt,
    AI_FINDINGS_SUMMARY_MAX_TOKENS,
    AI_FINDINGS_SUMMARY_TEMPERATURE,
  );
  if (!raw) return [];
  const findings: string[] = [];
  const followUps: string[] = [];
  let section: "findings" | "followup" | null = null;
  for (const line of raw.split("\n")) {
    const trimmed = line.trim();
    if (/^FINDINGS:/i.test(trimmed)) {
      section = "findings";
      continue;
    }
    if (/^FOLLOW.?UP:/i.test(trimmed)) {
      section = "followup";
      continue;
    }
    const item = trimmed.replace(/^[-*•]\s*/, "").trim();
    if (item.length < 5) continue;
    if (section === "findings") findings.push(item);
    if (section === "followup") followUps.push(item);
  }
  if (findings.length === 0 && followUps.length === 0) return [];
  status(
    `AI summarised ${findings.length} key findings, ${followUps.length} follow-up suggestions`,
  );
  return [
    {
      fromWorker: "round-coordinator",
      keyFindings: findings,
      suggestedFollowUps: followUps,
    },
  ];
}

const GAP_ROLE_MAP: Readonly<
  Record<
    string,
    {
      role: WorkerRole;
      followLinks: boolean;
      tiers?: ReadonlyArray<import("../types").SourceTier>;
    }
  >
> = {
  overview: { role: "breadth", followLinks: false },
  mechanism: { role: "technical", followLinks: true },
  history: { role: "breadth", followLinks: false },
  current: { role: "recency", followLinks: false },
  applications: { role: "breadth", followLinks: false },
  challenges: { role: "critical", followLinks: false },
  comparison: { role: "comparative", followLinks: false },
  evidence: {
    role: "academic",
    followLinks: true,
    tiers: ["academic", "government", "reference"],
  },
  expert: { role: "primary", followLinks: true, tiers: ["academic", "news"] },
  future: { role: "recency", followLinks: false },
  controversy: { role: "critical", followLinks: false },
  economics: { role: "statistical", followLinks: false },
};

export async function buildAdaptiveGapFill(
  topic: string,
  coveredIds: ReadonlyArray<string>,
  priorMessages: ReadonlyArray<AgentMessage>,
  useAI: boolean,
  status: StatusFn,
  profile: DepthProfile,
): Promise<ReadonlyArray<AdaptiveGapPlan>> {
  const gaps = detectGaps(coveredIds);
  if (gaps.length === 0) {
    status("All research dimensions covered - no gap queries needed");
    return [];
  }
  status(`Gaps: ${gaps.map((g) => g.label).join(", ")}`);
  const byRole = new Map<
    WorkerRole,
    {
      dimIds: string[];
      dimLabels: string[];
      queries: string[];
      followLinks: boolean;
      tiers?: ReadonlyArray<import("../types").SourceTier>;
    }
  >();
  const shortTopic = topic;
  for (const gap of gaps) {
    const mapping = GAP_ROLE_MAP[gap.id] ?? {
      role: "breadth" as WorkerRole,
      followLinks: false,
    };
    const existing = byRole.get(mapping.role) ?? {
      dimIds: [],
      dimLabels: [],
      queries: [],
      followLinks: mapping.followLinks,
      tiers: mapping.tiers,
    };
    existing.dimIds.push(gap.id);
    existing.dimLabels.push(gap.label);
    existing.queries.push(...gap.queries(shortTopic));
    byRole.set(mapping.role, existing);
  }
  if (useAI) {
    const followUpContext = priorMessages
      .flatMap((m) => m.suggestedFollowUps)
      .slice(0, 6);
    const entries = Array.from(byRole.entries());
    const aiResults = await Promise.allSettled(
      entries.map(async ([role, group]) => {
        const queryCount = Math.min(
          group.dimLabels.length * 3,
          profile.maxGapFillQueries,
        );
        const prompt = `You are a research assistant. A research session on "${topic}" is missing these angles:
${group.dimLabels.join(", ")}
${followUpContext.length > 0 ? `Previous round suggested exploring:\n${followUpContext.join("\n")}\n` : ""}
Generate ${queryCount} specific search queries to fill these gaps.
The queries should be best suited for a ${role} research agent.
Make queries diverse - cover different angles and phrasings.
Return ONLY the queries, one per line.
Queries:`;
        return { role, raw: await callLoadedModel(prompt) };
      }),
    );
    for (const result of aiResults) {
      if (result.status !== "fulfilled" || !result.value.raw) continue;
      const { role, raw } = result.value;
      const parsed = parseLines(raw, profile.maxGapFillQueries);
      if (parsed.length >= 2) {
        const group = byRole.get(role);
        if (group) group.queries = [...parsed];
      }
    }
  }
  const plans: AdaptiveGapPlan[] = [];
  for (const [role, group] of byRole) {
    // Determine if we have API keys available (passed via profile or config context)
    // Note: Since planner doesn't have direct cfg access, we rely on the fact that 
    // extraEngines in the profile will be used by the orchestrator for gap tasks.
    // We just need to ensure the gap task signals that it's a gap fill.
    plans.push({
      role,
      label: `Gap-fill: ${group.dimLabels.slice(0, 3).join(", ")}`,
      queries: group.queries.slice(0, profile.maxGapFillQueries),
      followLinks: group.followLinks,
      preferredTiers: group.tiers,
    });
  }
  status(`${plans.length} adaptive gap-fill worker(s) planned`);
  return plans;
}

const STOP_WORDS = new Set([
  "the",
  "a",
  "an",
  "is",
  "in",
  "of",
  "and",
  "or",
  "for",
  "to",
  "how",
  "what",
  "why",
  "when",
  "does",
  "with",
  "from",
  "that",
  "could",
  "which",
  "about",
  "their",
  "this",
  "these",
  "those",
  "would",
  "should",
  "current",
  "hypothetical",
  "scenarios",
  "lead",
]);

function extractKeywords(topic: string): ReadonlyArray<string> {
  return topic
    .toLowerCase()
    .replace(/[^a-z0-9\s]/g, "")
    .split(/\s+/)
    .filter((w) => w.length > 2 && !STOP_WORDS.has(w))
    .slice(0, 8);
}

function shortenTopic(topic: string): string {
  const colonIdx = topic.indexOf(":");
  const dashIdx = topic.indexOf(" - ");
  const sepIdx = colonIdx > 3 ? colonIdx : dashIdx > 3 ? dashIdx : -1;
  let core: string;
  if (sepIdx > 3 && sepIdx < topic.length * 0.6) {
    core = topic.slice(0, sepIdx).trim();
  } else {
    core = topic;
  }
  const words = core
    .replace(/[,;()]/g, " ")
    .split(/\s+/)
    .filter((w) => w.length > 1 && !STOP_WORDS.has(w.toLowerCase()));
  let result = "";
  let count = 0;
  for (const w of words) {
    if (count >= 6 || result.length + w.length > 58) break;
    result += (result ? " " : "") + w;
    count++;
  }
  if (result.length < 5) {
    result = topic.split(/\s+/).slice(0, 5).join(" ");
  }
  return result;
}

import { logLlmDiagnostics } from "../utils/tokens";
import { LMStudioClient } from "@lmstudio/sdk";
import {
  QueryPlan,
  WorkerRole,
  DynamicWorkerSpec,
  AdaptiveGapPlan,
  CrawledSource,
  AgentMessage,
  StatusFn,
} from "../types";
import { DIMENSIONS, detectGaps, gapFillQueries } from "./dimensions";
import {
  DepthProfile,
  AI_PLANNING_MAX_TOKENS,
  AI_PLANNING_TEMPERATURE,
  AI_PLANNING_TIMEOUT_MS,
  AI_MIN_ACCEPTABLE_QUERIES,
  AI_DECOMPOSITION_MAX_TOKENS,
  AI_DECOMPOSITION_TEMPERATURE,
  AI_DECOMPOSITION_TIMEOUT_MS,
  AI_FINDINGS_SUMMARY_MAX_TOKENS,
  AI_FINDINGS_SUMMARY_TEMPERATURE,
  FINDINGS_SUMMARY_SOURCE_CHARS,
  DECOMPOSITION_MIN_WORKERS,
  QUERY_LINE_MIN_LEN,
  QUERY_LINE_MAX_LEN,
  SYSTEM_INSTRUCTIONS,
} from "../constants";

async function callLoadedModel(
  prompt: string,
  maxTokens: number = AI_PLANNING_MAX_TOKENS,
  temperature: number = AI_PLANNING_TEMPERATURE,
  timeoutMs: number = AI_PLANNING_TIMEOUT_MS,
): Promise<string | null> {
  try {
    const client = new LMStudioClient();
    const models = await Promise.race<Awaited<ReturnType<typeof client.llm.listLoaded>>>([
      client.llm.listLoaded(),
      new Promise<never>((_, reject) =>
        setTimeout(() => reject(new Error("timeout")), timeoutMs),
      ),
    ]);
    if (!Array.isArray(models) || models.length === 0) return null;
    const model = await client.llm.model(models[0].identifier);
    const stream = model.respond(
      [
        { role: "system", content: SYSTEM_INSTRUCTIONS },
        { role: "user", content: prompt },
      ],
      {
        maxTokens,
        temperature,
      },
    );
    let result = "";
    for await (const chunk of stream) result += chunk.content ?? "";
    return result.trim() || null;
  } catch {
    return null;
  }
}

function parseLines(raw: string, maxLines: number = 6): ReadonlyArray<string> {
  return raw
    .split(/\n/)
    .map((line) => line.replace(/^\d+[.)]\s*|^[-*•]\s*/, "").trim())
    .filter(
      (line) =>
        line.length > QUERY_LINE_MIN_LEN && line.length < QUERY_LINE_MAX_LEN,
    )
    .filter((line, idx, arr) => arr.indexOf(line) === idx)
    .slice(0, maxLines);
}

const VALID_ROLES: ReadonlyArray<WorkerRole> = [
  "breadth",
  "depth",
  "recency",
  "academic",
  "critical",
  "statistical",
  "regulatory",
  "technical",
  "primary",
  "comparative",
];

function makeDecompositionPrompt(
  topic: string,
  focusAreas: ReadonlyArray<string>,
  profile: DepthProfile,
): string {
  const focus = focusAreas.length
    ? `\nFocus areas: ${focusAreas.join(", ")}`
    : "";
  return `You are a research decomposition system. Given a research topic, output a JSON array of specialized worker agents.
Topic: "${topic}"${focus}
Each worker needs:
"role": one of "breadth", "depth", "recency", "academic", "critical", "statistical", "regulatory", "technical", "primary", "comparative"
"label": descriptive name (e.g., "Clinical Evidence Researcher", "Policy Critic")
"queries": array of ${Math.min(profile.maxQueriesPerWorker, 6)}-${profile.maxQueriesPerWorker} specific search queries for this worker
"budgetWeight": number 0.1-0.4 (must sum to ~1.0 across all workers)
"followLinks": true/false (true for depth/academic workers)
"preferredTiers": optional array of "academic","government","reference","news","professional","general"
Rules:
Output ${DECOMPOSITION_MIN_WORKERS} to ${profile.maxDecompositionWorkers} workers
Tailor the workers to THIS specific topic - not generic roles
Queries must be highly specific to the topic and each worker's assignment
Generate MORE queries for broader or more complex topics
Budget weights must roughly sum to 1.0
Output ONLY valid JSON, no other text
JSON:`;
}

async function aiDecompose(
  topic: string,
  focusAreas: ReadonlyArray<string>,
  status: StatusFn,
  profile: DepthProfile,
): Promise<ReadonlyArray<DynamicWorkerSpec> | null> {
  const raw = await callLoadedModel(
    makeDecompositionPrompt(topic, focusAreas, profile),
    AI_DECOMPOSITION_MAX_TOKENS,
    AI_DECOMPOSITION_TEMPERATURE,
    AI_DECOMPOSITION_TIMEOUT_MS,
  );
  if (!raw) return null;
  try {
    const jsonStr = raw.replace(/```json\s*|```\s*/g, "").trim();
    const parsed = JSON.parse(jsonStr);
    if (!Array.isArray(parsed) || parsed.length < DECOMPOSITION_MIN_WORKERS)
      return null;
    const specs: DynamicWorkerSpec[] = [];
    for (const item of parsed.slice(0, profile.maxDecompositionWorkers)) {
      const role = VALID_ROLES.includes(item.role) ? item.role : "breadth";
      const queries = Array.isArray(item.queries)
        ? item.queries
            .filter((q: unknown) => typeof q === "string" && q.length > 3)
            .slice(0, profile.maxQueriesPerWorker)
        : [];
      if (queries.length < 2) continue;
      specs.push({
        role: role as WorkerRole,
        label:
          typeof item.label === "string"
            ? item.label.slice(0, 60)
            : `${role} worker`,
        queries,
        budgetWeight:
          typeof item.budgetWeight === "number"
            ? Math.max(0.05, Math.min(0.5, item.budgetWeight))
            : 0.2,
        followLinks: item.followLinks === true,
        preferredTiers: Array.isArray(item.preferredTiers)
          ? item.preferredTiers
          : undefined,
      });
    }
    if (specs.length < DECOMPOSITION_MIN_WORKERS) return null;
    const totalWeight = specs.reduce((sum, s) => sum + s.budgetWeight, 0);
    const normalised = specs.map((s) => ({
      ...s,
      budgetWeight: s.budgetWeight / totalWeight,
    }));
    status(`AI decomposed topic into ${normalised.length} specialised workers`);
    return normalised;
  } catch {
    return null;
  }
}

function makeRolePlanPrompt(
  role: WorkerRole,
  topic: string,
  focusAreas: ReadonlyArray<string>,
  profile: DepthProfile,
): string {
  const roleDescriptions: Readonly<Record<WorkerRole, string>> = {
    breadth: "broad coverage - many different angles, facts, and sub-topics",
    depth: "deep dive - mechanisms, how it works, technical detail, and evidence",
    recency: "recent developments - 2024-2026+ news, updates, and latest research",
    academic: "academic and scientific sources - peer-reviewed studies, journals, authoritative papers",
    critical: "critical analysis - limitations, counterarguments, criticism, controversy, drawbacks",
    statistical: "statistics and data - numbers, percentages, datasets, surveys, market sizes, quantitative evidence",
    regulatory: "regulatory and policy - laws, regulations, government policies, compliance, standards, guidelines",
    technical: "technical deep-dive - implementation details, specifications, architecture, engineering approaches",
    primary: "primary sources - original reports, official statements, first-hand accounts, press releases, white papers",
    comparative: "comparative analysis - vs alternatives, head-to-head comparisons, benchmarks, trade-offs, pros and cons",
  };
  const focus = focusAreas.length
    ? `\nFocus especially on: ${focusAreas.join(", ")}`
    : "";
  return `You are a research planning assistant. Generate search queries for a specialised research agent.
Topic: "${topic}"${focus}
This agent's role: ${roleDescriptions[role]}
Generate exactly ${profile.maxQueriesPerWorker} highly specific, diverse search queries for this role.
Rules:
Each query must be different from the others
Use natural language (as a human would type into a search engine)
Be specific to the role - ${roleDescriptions[role]}
Vary query structure: some factual, some comparative, some recent
Return ONLY the queries, one per line, no numbering, no extra text
Queries:`;
}

const ROLE_DIMENSIONS: Readonly<Record<WorkerRole, ReadonlyArray<string>>> = {
  breadth: ["overview", "applications", "history", "economics"],
  depth: ["mechanism", "evidence", "expert"],
  recency: ["current", "future"],
  academic: ["evidence", "expert", "mechanism"],
  critical: ["challenges", "controversy", "comparison"],
  statistical: ["evidence", "economics", "overview"],
  regulatory: ["challenges", "controversy", "current"],
  technical: ["mechanism", "applications", "evidence"],
  primary: ["evidence", "expert", "history"],
  comparative: ["comparison", "challenges", "applications"],
};

function dimensionFallbackQueries(
  role: WorkerRole,
  topic: string,
  focusAreas: ReadonlyArray<string>,
  maxQueries: number,
): ReadonlyArray<string> {
  const dimIds = ROLE_DIMENSIONS[role];
  const dims = DIMENSIONS.filter((d) => dimIds.includes(d.id));
  const queries: string[] = [];
  const shortTopic = topic;
  for (const dim of dims) {
    for (const q of dim.queries(shortTopic)) {
      if (!queries.includes(q)) queries.push(q);
    }
  }
  for (const area of focusAreas) {
    const q = `${shortTopic} ${area}`;
    if (!queries.includes(q)) queries.push(q);
  }
  return queries.slice(0, maxQueries);
}

export async function buildQueryPlan(
  topic: string,
  focusAreas: ReadonlyArray<string>,
  useAI: boolean,
  status: StatusFn,
  profile: DepthProfile,
): Promise<QueryPlan> {
  const CORE_ROLES: ReadonlyArray<WorkerRole> = [
    "breadth",
    "depth",
    "recency",
    "academic",
    "critical",
  ];
  const EXTENDED_ROLES: ReadonlyArray<WorkerRole> = [
    "statistical",
    "regulatory",
    "technical",
    "primary",
    "comparative",
  ];
  let roles: ReadonlyArray<WorkerRole>;
  if (profile.depthRounds >= 10) {
    roles = [...CORE_ROLES, ...EXTENDED_ROLES];
  } else if (profile.depthRounds >= 5) {
    roles = [...CORE_ROLES, "technical", "comparative", "statistical"];
  } else {
    roles = CORE_ROLES;
  }
  const queriesByRole: Partial<Record<WorkerRole, ReadonlyArray<string>>> = {};
  let usedAI = false;
  let dynamicSpecs: ReadonlyArray<DynamicWorkerSpec> | undefined;
  if (useAI) {
    status("AI task decomposition - analysing topic for specialised workers…");
    const specs = await aiDecompose(topic, focusAreas, status, profile);
    if (specs && specs.length >= DECOMPOSITION_MIN_WORKERS) {
      dynamicSpecs = specs;
      usedAI = true;
      for (const spec of specs) {
        queriesByRole[spec.role] = spec.queries;
      }
    } else {
      status("AI planning queries for each swarm worker…");
    }
    const uncoveredRoles = roles.filter((r) => !queriesByRole[r]?.length);
    if (uncoveredRoles.length > 0) {
      const results = await Promise.allSettled(
        uncoveredRoles.map(async (role) => ({
          role,
          queries: await callLoadedModel(
            makeRolePlanPrompt(role, topic, focusAreas, profile),
          ),
        })),
      );
      for (const result of results) {
        if (result.status !== "fulfilled") continue;
        const { role, queries: raw } = result.value;
        if (!raw) continue;
        const parsed = parseLines(raw, profile.maxQueriesPerWorker);
        if (parsed.length >= AI_MIN_ACCEPTABLE_QUERIES) {
          queriesByRole[role] = parsed;
          usedAI = true;
        }
      }
    }
    if (usedAI) {
      status(
        `AI generated queries for ${Object.keys(queriesByRole).length} worker role(s)`,
      );
    } else {
      status("AI unavailable, using dimension-based query planning");
    }
  }
  for (const role of roles) {
    if (!queriesByRole[role] || queriesByRole[role]!.length === 0) {
      queriesByRole[role] = dimensionFallbackQueries(
        role,
        topic,
        focusAreas,
        profile.maxQueriesPerWorker,
      );
    }
  }
  return {
    queriesByRole: queriesByRole as Record<WorkerRole, ReadonlyArray<string>>,
    usedAI,
    topicKeywords: extractKeywords(topic),
    dynamicSpecs,
  };
  const topicKeywords = extractKeywords(topic);

  // PRE-SEARCH PLANNING (If AI is enabled)
  if (useAI) {
    const planPrompt = `Topic: ${topic}\nFocus: ${focusAreas.join(", ")}\n
    Generate a JSON object with:
    1. "subQuestions": 3-5 specific questions to answer.
    2. "likelySourceTypes": Array of strings (e.g., "academic", "news", "reference").
    3. "stopConditions": Array of strings (e.g., "Found 3 sources confirming Q3 revenue").`;
    
    logLlmDiagnostics("buildQueryPlan-PreSearch", planPrompt);
    
    // Call LLM here and parse the JSON to inject into your existing query generation logic.
    // This prevents immediate blind fan-out.
  }

  return {
    queriesByRole: queriesByRole as Record<WorkerRole, ReadonlyArray<string>>,
    usedAI,
    topicKeywords,
    dynamicSpecs,
  };
} // <--- THIS BRACE CLOSES THE FUNCTION

export async function summariseFindings(
  sources: ReadonlyArray<CrawledSource>,
  topic: string,
  useAI: boolean,
  status: StatusFn,
): Promise<ReadonlyArray<AgentMessage>> {
  if (!useAI || sources.length === 0) return [];
  const sourceSummaries = sources
    .slice(0, 20)
    .map(
      (s, i) =>
        `[${i + 1}] ${s.workerLabel}: ${s.title} - ${s.text.slice(0, FINDINGS_SUMMARY_SOURCE_CHARS)}`,
    )
    .join("\n\n");
  const prompt = `You are a research coordinator. A team of research agents collected these sources on "${topic}":
${sourceSummaries}
Summarise:
The 3-5 most important findings discovered so far (one line each)
3-5 specific questions or angles that were NOT covered and need follow-up
Output format:
FINDINGS:
finding 1
finding 2
...
FOLLOW_UP:
question 1
question 2
...`;
  const raw = await callLoadedModel(
    prompt,
    AI_FINDINGS_SUMMARY_MAX_TOKENS,
    AI_FINDINGS_SUMMARY_TEMPERATURE,
  );
  if (!raw) return [];
  const findings: string[] = [];
  const followUps: string[] = [];
  let section: "findings" | "followup" | null = null;
  for (const line of raw.split("\n")) {
    const trimmed = line.trim();
    if (/^FINDINGS:/i.test(trimmed)) {
      section = "findings";
      continue;
    }
    if (/^FOLLOW.?UP:/i.test(trimmed)) {
      section = "followup";
      continue;
    }
    const item = trimmed.replace(/^[-*•]\s*/, "").trim();
    if (item.length < 5) continue;
    if (section === "findings") findings.push(item);
    if (section === "followup") followUps.push(item);
  }
  if (findings.length === 0 && followUps.length === 0) return [];
  status(
    `AI summarised ${findings.length} key findings, ${followUps.length} follow-up suggestions`,
  );
  return [
    {
      fromWorker: "round-coordinator",
      keyFindings: findings,
      suggestedFollowUps: followUps,
    },
  ];
}

const GAP_ROLE_MAP: Readonly<
  Record<
    string,
    {
      role: WorkerRole;
      followLinks: boolean;
      tiers?: ReadonlyArray<import("../types").SourceTier>;
    }
  >
> = {
  overview: { role: "breadth", followLinks: false },
  mechanism: { role: "technical", followLinks: true },
  history: { role: "breadth", followLinks: false },
  current: { role: "recency", followLinks: false },
  applications: { role: "breadth", followLinks: false },
  challenges: { role: "critical", followLinks: false },
  comparison: { role: "comparative", followLinks: false },
  evidence: {
    role: "academic",
    followLinks: true,
    tiers: ["academic", "government", "reference"],
  },
  expert: { role: "primary", followLinks: true, tiers: ["academic", "news"] },
  future: { role: "recency", followLinks: false },
  controversy: { role: "critical", followLinks: false },
  economics: { role: "statistical", followLinks: false },
};

export async function buildAdaptiveGapFill(
  topic: string,
  coveredIds: ReadonlyArray<string>,
  priorMessages: ReadonlyArray<AgentMessage>,
  useAI: boolean,
  status: StatusFn,
  profile: DepthProfile,
): Promise<ReadonlyArray<AdaptiveGapPlan>> {
  const gaps = detectGaps(coveredIds);
  if (gaps.length === 0) {
    status("All research dimensions covered - no gap queries needed");
    return [];
  }
  status(`Gaps: ${gaps.map((g) => g.label).join(", ")}`);
  const byRole = new Map<
    WorkerRole,
    {
      dimIds: string[];
      dimLabels: string[];
      queries: string[];
      followLinks: boolean;
      tiers?: ReadonlyArray<import("../types").SourceTier>;
    }
  >();
  const shortTopic = topic;
  for (const gap of gaps) {
    const mapping = GAP_ROLE_MAP[gap.id] ?? {
      role: "breadth" as WorkerRole,
      followLinks: false,
    };
    const existing = byRole.get(mapping.role) ?? {
      dimIds: [],
      dimLabels: [],
      queries: [],
      followLinks: mapping.followLinks,
      tiers: mapping.tiers,
    };
    existing.dimIds.push(gap.id);
    existing.dimLabels.push(gap.label);
    existing.queries.push(...gap.queries(shortTopic));
    byRole.set(mapping.role, existing);
  }
  if (useAI) {
    const followUpContext = priorMessages
      .flatMap((m) => m.suggestedFollowUps)
      .slice(0, 6);
    const entries = Array.from(byRole.entries());
    const aiResults = await Promise.allSettled(
      entries.map(async ([role, group]) => {
        const queryCount = Math.min(
          group.dimLabels.length * 3,
          profile.maxGapFillQueries,
        );
        const prompt = `You are a research assistant. A research session on "${topic}" is missing these angles:
${group.dimLabels.join(", ")}
${followUpContext.length > 0 ? `Previous round suggested exploring:\n${followUpContext.join("\n")}\n` : ""}
Generate ${queryCount} specific search queries to fill these gaps.
The queries should be best suited for a ${role} research agent.
Make queries diverse - cover different angles and phrasings.
Return ONLY the queries, one per line.
Queries:`;
        return { role, raw: await callLoadedModel(prompt) };
      }),
    );
    for (const result of aiResults) {
      if (result.status !== "fulfilled" || !result.value.raw) continue;
      const { role, raw } = result.value;
      const parsed = parseLines(raw, profile.maxGapFillQueries);
      if (parsed.length >= 2) {
        const group = byRole.get(role);
        if (group) group.queries = [...parsed];
      }
    }
  }
  const plans: AdaptiveGapPlan[] = [];
  for (const [role, group] of byRole) {
    // Determine if we have API keys available (passed via profile or config context)
    // Note: Since planner doesn't have direct cfg access, we rely on the fact that 
    // extraEngines in the profile will be used by the orchestrator for gap tasks.
    // We just need to ensure the gap task signals that it's a gap fill.
    plans.push({
      role,
      label: `Gap-fill: ${group.dimLabels.slice(0, 3).join(", ")}`,
      queries: group.queries.slice(0, profile.maxGapFillQueries),
      followLinks: group.followLinks,
      preferredTiers: group.tiers,
    });
  }
  status(`${plans.length} adaptive gap-fill worker(s) planned`);
  return plans;
}

const STOP_WORDS = new Set([
  "the",
  "a",
  "an",
  "is",
  "in",
  "of",
  "and",
  "or",
  "for",
  "to",
  "how",
  "what",
  "why",
  "when",
  "does",
  "with",
  "from",
  "that",
  "could",
  "which",
  "about",
  "their",
  "this",
  "these",
  "those",
  "would",
  "should",
  "current",
  "hypothetical",
  "scenarios",
  "lead",
]);

function extractKeywords(topic: string): ReadonlyArray<string> {
  return topic
    .toLowerCase()
    .replace(/[^a-z0-9\s]/g, "")
    .split(/\s+/)
    .filter((w) => w.length > 2 && !STOP_WORDS.has(w))
    .slice(0, 8);
}

function shortenTopic(topic: string): string {
  const colonIdx = topic.indexOf(":");
  const dashIdx = topic.indexOf(" - ");
  const sepIdx = colonIdx > 3 ? colonIdx : dashIdx > 3 ? dashIdx : -1;
  let core: string;
  if (sepIdx > 3 && sepIdx < topic.length * 0.6) {
    core = topic.slice(0, sepIdx).trim();
  } else {
    core = topic;
  }
  const words = core
    .replace(/[,;()]/g, " ")
    .split(/\s+/)
    .filter((w) => w.length > 1 && !STOP_WORDS.has(w.toLowerCase()));
  let result = "";
  let count = 0;
  for (const w of words) {
    if (count >= 6 || result.length + w.length > 58) break;
    result += (result ? " " : "") + w;
    count++;
  }
  if (result.length < 5) {
    result = topic.split(/\s+/).slice(0, 5).join(" ");
  }
  return result;
}