import { ensureCustomAgentsTable } from "./custom-agents.js";
import { createId } from "./id.js";
import { prisma } from "./prisma.js";
import { resolveSystemAgentInstructions } from "./system-agents.js";
import { runJsonChatCompletion } from "./tenant-ai.js";

export type OperationalAgentCandidateScope =
  | "builtin"
  | "system_agent"
  | "custom_agent"
  | "form_agent";

export type OperationalAgentCandidate = {
  agentKey: string;
  agentName: string;
  agentScope: OperationalAgentCandidateScope;
  description: string | null;
  promptText?: string | null;
  operationTypes?: string[];
  tags?: string[];
  metadata?: Record<string, unknown>;
};

type OperationalRoutingUsage = {
  inputTokens: number;
  outputTokens: number;
  totalTokens: number;
};

export type OperationalRoutingDecision = {
  selectedAgentKey: string;
  selectedAgentName: string;
  selectedAgentScope: OperationalAgentCandidateScope;
  rationale: string;
  confidence: number;
  selectedBy: "router" | "fallback";
  usage: OperationalRoutingUsage | null;
  model: string | null;
};

let ensureOperationalAgentRoutingTablePromise: Promise<void> | null = null;

function normalizeAgentText(value: string | null | undefined) {
  return String(value ?? "").replace(/\s+/g, " ").trim();
}

function buildOperationalFallbackDecision(input: {
  candidate: OperationalAgentCandidate;
  rationale: string;
}) {
  return {
    selectedAgentKey: input.candidate.agentKey,
    selectedAgentName: input.candidate.agentName,
    selectedAgentScope: input.candidate.agentScope,
    rationale: normalizeAgentText(input.rationale) || "Fallback routing selected this agent.",
    confidence: 0.51,
    selectedBy: "fallback" as const,
    usage: null,
    model: null,
  } satisfies OperationalRoutingDecision;
}

async function resolveLawFirmName(lawFirmId: string) {
  const [row] = await prisma.$queryRaw<Array<{ name: string | null }>>`
    SELECT name
    FROM law_firms
    WHERE id = ${lawFirmId}
    LIMIT 1
  `;

  return normalizeAgentText(row?.name) || "workspace";
}

export async function ensureOperationalAgentRoutingTable() {
  if (!ensureOperationalAgentRoutingTablePromise) {
    ensureOperationalAgentRoutingTablePromise = (async () => {
      await prisma.$executeRawUnsafe(`
        CREATE TABLE IF NOT EXISTS operational_agent_routing_decisions (
          id CHAR(36) NOT NULL PRIMARY KEY,
          law_firm_id CHAR(36) NOT NULL,
          ai_run_id CHAR(36) NULL,
          source_entity_type VARCHAR(50) NOT NULL,
          source_entity_id CHAR(36) NULL,
          operation_type VARCHAR(50) NOT NULL,
          stage_code VARCHAR(80) NOT NULL,
          selected_agent_scope VARCHAR(50) NOT NULL,
          selected_agent_key VARCHAR(120) NOT NULL,
          selected_agent_name VARCHAR(180) NOT NULL,
          decision_json JSON NOT NULL,
          created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
          KEY idx_operational_agent_routing_law_firm (law_firm_id, created_at),
          KEY idx_operational_agent_routing_ai_run (ai_run_id, created_at),
          KEY idx_operational_agent_routing_source (source_entity_type, source_entity_id, created_at),
          KEY idx_operational_agent_routing_stage (operation_type, stage_code, created_at),
          CONSTRAINT fk_operational_agent_routing_law_firm
            FOREIGN KEY (law_firm_id) REFERENCES law_firms (id),
          CONSTRAINT fk_operational_agent_routing_ai_run
            FOREIGN KEY (ai_run_id) REFERENCES ai_runs (id)
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
      `);
    })().catch((error) => {
      ensureOperationalAgentRoutingTablePromise = null;
      throw error;
    });
  }

  await ensureOperationalAgentRoutingTablePromise;
}

export async function listOperationalCustomAgentCandidates(input: { lawFirmId: string }) {
  await ensureCustomAgentsTable();

  const rows = await prisma.$queryRaw<
    Array<{
      id: string;
      agent_name: string;
      description: string | null;
      system_prompt: string;
      team_name: string;
    }>
  >`
    SELECT
      ca.id,
      ca.agent_name,
      ca.description,
      ca.system_prompt,
      t.name AS team_name
    FROM custom_agents ca
    INNER JOIN teams t ON t.id = ca.team_id
    WHERE ca.law_firm_id = ${input.lawFirmId}
      AND ca.archived_at IS NULL
    ORDER BY ca.updated_at DESC, ca.agent_name ASC
    LIMIT 50
  `;

  return rows.map((row) => ({
    agentKey: `custom:${row.id}`,
    agentName: normalizeAgentText(row.agent_name) || "Custom agent",
    agentScope: "custom_agent" as const,
    description:
      normalizeAgentText(row.description) ||
      `Custom operational agent from team ${normalizeAgentText(row.team_name) || "workspace"}.`,
    promptText: String(row.system_prompt ?? "").trim() || null,
    operationTypes: ["document_extraction", "form_fill"],
    tags: ["custom", "workspace", "tenant_defined"],
    metadata: {
      customAgentId: row.id,
      teamName: normalizeAgentText(row.team_name) || null,
    },
  })) satisfies OperationalAgentCandidate[];
}

export async function routeOperationalAgentSelection(input: {
  lawFirmId: string;
  operationType: "document_extraction" | "form_fill" | "process_generation";
  stageCode: string;
  sourceEntityType: string;
  sourceEntityId?: string | null;
  availableAgents: OperationalAgentCandidate[];
  fallbackAgentKey?: string | null;
  context: Record<string, unknown>;
  decisionRules?: string[];
}) {
  const normalizedCandidates = Array.from(
    new Map(
      input.availableAgents
        .filter((candidate) => normalizeAgentText(candidate.agentKey))
        .map((candidate) => [candidate.agentKey, candidate]),
    ).values(),
  );

  if (!normalizedCandidates.length) {
    throw new Error("Operational router requires at least one available agent");
  }

  const fallbackCandidate =
    normalizedCandidates.find(
      (candidate) => candidate.agentKey === normalizeAgentText(input.fallbackAgentKey),
    ) ?? normalizedCandidates[0];

  if (normalizedCandidates.length === 1) {
    return buildOperationalFallbackDecision({
      candidate: fallbackCandidate,
      rationale: "Only one compatible agent was available for this operational step.",
    });
  }

  const lawFirmName = await resolveLawFirmName(input.lawFirmId);

  try {
    const router = await resolveSystemAgentInstructions({
      lawFirmId: input.lawFirmId,
      lawFirmName,
      agentCode: "document_generation_orchestrator",
    });

    const completion = await runJsonChatCompletion({
      lawFirmId: input.lawFirmId,
      systemPrompt: router.instructions,
      userPrompt: JSON.stringify(
        {
          operation: {
            operationType: input.operationType,
            stageCode: normalizeAgentText(input.stageCode),
            sourceEntityType: normalizeAgentText(input.sourceEntityType),
            sourceEntityId: normalizeAgentText(input.sourceEntityId) || null,
          },
          context: input.context,
          availableAgents: normalizedCandidates.map((candidate) => ({
            agentKey: candidate.agentKey,
            agentName: candidate.agentName,
            agentScope: candidate.agentScope,
            description: candidate.description,
            operationTypes: candidate.operationTypes ?? [],
            tags: candidate.tags ?? [],
            metadata: candidate.metadata ?? null,
          })),
          decisionRules:
            input.decisionRules?.filter((rule) => normalizeAgentText(rule)).map((rule) => rule.trim()) ??
            [],
          outputContract: {
            selectedAgentKey: "string",
            rationale: "string",
            confidence: "number between 0 and 1",
          },
        },
        null,
        2,
      ),
      maxCompletionTokens: 900,
    });

    const selectedAgentKey = normalizeAgentText(
      String(completion.json.selectedAgentKey ?? fallbackCandidate.agentKey),
    );
    const selectedCandidate =
      normalizedCandidates.find((candidate) => candidate.agentKey === selectedAgentKey) ??
      fallbackCandidate;
    const confidence = Math.max(
      0,
      Math.min(1, Number(completion.json.confidence ?? 0.5) || 0.5),
    );
    const rationale =
      normalizeAgentText(String(completion.json.rationale ?? "")) ||
      "Operational router selected the agent with the best contextual fit.";

    return {
      selectedAgentKey: selectedCandidate.agentKey,
      selectedAgentName: selectedCandidate.agentName,
      selectedAgentScope: selectedCandidate.agentScope,
      rationale:
        selectedCandidate.agentKey === selectedAgentKey ?
          rationale
        : `${rationale} The requested agent was unavailable, so a compatible fallback was used.`,
      confidence:
        selectedCandidate.agentKey === selectedAgentKey ? confidence : Math.min(confidence, 0.45),
      selectedBy: "router" as const,
      usage: completion.usage,
      model: completion.model,
    } satisfies OperationalRoutingDecision;
  } catch (error) {
    return buildOperationalFallbackDecision({
      candidate: fallbackCandidate,
      rationale: `Fallback routing was used for this operational step: ${error instanceof Error ? error.message : "unexpected router failure"}.`,
    });
  }
}

export async function recordOperationalAgentRoutingDecision(input: {
  lawFirmId: string;
  aiRunId?: string | null;
  sourceEntityType: string;
  sourceEntityId?: string | null;
  operationType: string;
  stageCode: string;
  decision: OperationalRoutingDecision;
  availableAgents: OperationalAgentCandidate[];
  contextSummary?: Record<string, unknown> | null;
}) {
  await ensureOperationalAgentRoutingTable();

  await prisma.$executeRaw`
    INSERT INTO operational_agent_routing_decisions (
      id,
      law_firm_id,
      ai_run_id,
      source_entity_type,
      source_entity_id,
      operation_type,
      stage_code,
      selected_agent_scope,
      selected_agent_key,
      selected_agent_name,
      decision_json,
      created_at
    ) VALUES (
      ${createId()},
      ${input.lawFirmId},
      ${input.aiRunId ?? null},
      ${normalizeAgentText(input.sourceEntityType) || "unknown"},
      ${normalizeAgentText(input.sourceEntityId) || null},
      ${normalizeAgentText(input.operationType) || "unknown"},
      ${normalizeAgentText(input.stageCode) || "unknown"},
      ${input.decision.selectedAgentScope},
      ${input.decision.selectedAgentKey},
      ${input.decision.selectedAgentName},
      ${JSON.stringify({
        selectedAgentKey: input.decision.selectedAgentKey,
        selectedAgentName: input.decision.selectedAgentName,
        selectedAgentScope: input.decision.selectedAgentScope,
        rationale: input.decision.rationale,
        confidence: input.decision.confidence,
        selectedBy: input.decision.selectedBy,
        routerModel: input.decision.model,
        availableAgents: input.availableAgents.map((candidate) => ({
          agentKey: candidate.agentKey,
          agentName: candidate.agentName,
          agentScope: candidate.agentScope,
          description: candidate.description,
          operationTypes: candidate.operationTypes ?? [],
          tags: candidate.tags ?? [],
          metadata: candidate.metadata ?? null,
        })),
        contextSummary: input.contextSummary ?? null,
      })},
      CURRENT_TIMESTAMP
    )
  `;
}