import type { AgentTask } from '@pig/db';
import { z } from 'zod';
import { zodToJsonSchema } from 'zod-to-json-schema';

export interface AgentTool {
  name: string;
  description: string;
  inputSchema: z.ZodTypeAny;
  execute(input: unknown, signal?: AbortSignal): Promise<unknown>;
}

interface ToolDefinition<TSchema extends z.ZodTypeAny> {
  name: string;
  description: string;
  inputSchema: TSchema;
  execute(input: z.infer<TSchema>, signal?: AbortSignal): Promise<unknown>;
}

/** Keep tool construction typed while exposing no ambient coding-agent tools. */
export function defineTool<TSchema extends z.ZodTypeAny>(
  definition: ToolDefinition<TSchema>,
): AgentTool {
  return {
    name: definition.name,
    description: definition.description,
    inputSchema: definition.inputSchema,
    execute: async (input, signal) => definition.execute(definition.inputSchema.parse(input), signal),
  };
}

export interface AgentProviderRequest {
  task: AgentTask;
  tools: AgentTool[];
  signal?: AbortSignal;
}

export interface AgentProviderResult {
  summary: string;
  inputTokens: number | null;
  outputTokens: number | null;
  result: Record<string, unknown>;
}

export interface AgentProvider {
  readonly model: string;
  run(request: AgentProviderRequest): Promise<AgentProviderResult>;
}

export interface PrimeOpenAIProviderOptions {
  apiKey: string;
  baseUrl?: string;
  model?: string;
  maxTokens?: number;
  /** Total attempts per model call, including the first. */
  maxAttempts?: number;
  /** Deadline for one attempt, headers and body together. */
  timeoutMs?: number;
  maxBackoffMs?: number;
  onRetry?: InferenceRetryPolicy['onRetry'];
  fetchImpl?: typeof fetch;
}

const toolCallSchema = z.object({
  id: z.string(),
  type: z.literal('function').optional(),
  function: z.object({
    name: z.string(),
    arguments: z.string(),
  }),
});

const completionSchema = z.object({
  choices: z
    .array(
      z.object({
        message: z.object({
          content: z.string().nullable().optional(),
          tool_calls: z.array(toolCallSchema).optional(),
        }),
      }),
    )
    .min(1),
  usage: z
    .object({
      prompt_tokens: z.number().int().nonnegative().optional(),
      completion_tokens: z.number().int().nonnegative().optional(),
    })
    .optional(),
});

type ToolCall = z.infer<typeof toolCallSchema>;
type ChatMessage =
  | { role: 'system' | 'user'; content: string }
  | { role: 'assistant'; content: string | null; tool_calls?: ToolCall[] }
  | { role: 'tool'; tool_call_id: string; name: string; content: string };

export class PrimeOpenAIProvider implements AgentProvider {
  readonly model: string;
  private readonly baseUrl: string;
  private readonly maxTokens: number;
  private readonly retry: InferenceRetryPolicy;
  private readonly fetchImpl: typeof fetch;

  constructor(private readonly options: PrimeOpenAIProviderOptions) {
    this.model = options.model ?? 'nvidia/nemotron-3-nano-30b-a3b';
    this.baseUrl = (options.baseUrl ?? 'https://api.pinference.ai/api/v1').replace(/\/$/, '');
    this.maxTokens = options.maxTokens ?? 1_024;
    // Nobody is waiting on a queued task, so it can afford the fuller budget:
    // five attempts, and a deadline that covers the whole non-streamed body.
    this.retry = {
      maxAttempts: options.maxAttempts ?? 5,
      timeoutMs: options.timeoutMs ?? 60_000,
      maxBackoffMs: options.maxBackoffMs ?? 30_000,
      onRetry: options.onRetry,
    };
    this.fetchImpl = options.fetchImpl ?? fetch;
  }

  async run(request: AgentProviderRequest): Promise<AgentProviderResult> {
    const toolsByName = new Map(request.tools.map((tool) => [tool.name, tool]));
    const messages: ChatMessage[] = [
      { role: 'system', content: systemPrompt },
      { role: 'user', content: taskPrompt(request.task) },
    ];
    let inputTokens = 0;
    let outputTokens = 0;
    let toolCallCount = 0;

    // `budget` counts model calls, not tools. A final answer after a tool is a
    // separate call and must fit inside the budget the queue row authorised.
    for (let turn = 0; turn < Math.max(1, request.task.budget); turn += 1) {
      // The schema check sits outside the retry on purpose: a truncated body is
      // worth another attempt, but a response the schema rejects will be
      // rejected identically five times over and each one costs credit.
      const payload = await withInferenceRetries(this.retry, request.signal, async (attemptSignal) => {
        const response = await this.fetchImpl(`${this.baseUrl}/chat/completions`, {
          method: 'POST',
          headers: {
            authorization: `Bearer ${this.options.apiKey}`,
            'content-type': 'application/json',
            accept: 'application/json',
          },
          body: JSON.stringify({
            model: this.model,
            messages,
            tools: request.tools.map((tool) => ({
              type: 'function',
              function: {
                name: tool.name,
                description: tool.description,
                parameters: zodToJsonSchema(tool.inputSchema, {
                  $refStrategy: 'none',
                  target: 'openAi',
                }),
              },
            })),
            tool_choice: 'auto',
            parallel_tool_calls: false,
            temperature: 0,
            max_tokens: this.maxTokens,
            // Nemotron otherwise spends a tight response budget thinking aloud
            // and can truncate before emitting the tool call or extraction.
            reasoning_effort: 'none',
          }),
          signal: attemptSignal,
        });

        if (!response.ok) throw await inferenceErrorFor(response);
        return (await response.json()) as unknown;
      });

      const completion = completionSchema.parse(payload);
      inputTokens += completion.usage?.prompt_tokens ?? 0;
      outputTokens += completion.usage?.completion_tokens ?? 0;
      const message = completion.choices[0]!.message;
      const toolCalls = message.tool_calls ?? [];
      messages.push({
        role: 'assistant',
        content: message.content ?? null,
        ...(toolCalls.length > 0 ? { tool_calls: toolCalls } : {}),
      });

      if (toolCalls.length === 0) {
        const summary = message.content?.trim();
        if (!summary) throw new Error('Piggy returned neither text nor a tool call.');
        return {
          summary,
          inputTokens: inputTokens || null,
          outputTokens: outputTokens || null,
          result: { messages, toolCallCount },
        };
      }

      for (const toolCall of toolCalls) {
        toolCallCount += 1;
        const tool = toolsByName.get(toolCall.function.name);
        let content: string;

        if (!tool) {
          // A hallucinated coding tool is an error result, never an ambient
          // capability lookup. Only the explicit PIG registry can execute.
          content = JSON.stringify({ error: `Tool ${toolCall.function.name} is not available.` });
        } else {
          try {
            const args = JSON.parse(toolCall.function.arguments) as unknown;
            content = JSON.stringify({ ok: true, result: await tool.execute(args, request.signal) });
          } catch (error) {
            content = JSON.stringify({
              ok: false,
              error: error instanceof Error ? error.message : String(error),
            });
          }
        }

        messages.push({
          role: 'tool',
          tool_call_id: toolCall.id,
          name: toolCall.function.name,
          content,
        });
      }
    }

    throw new Error(`Piggy exhausted its ${request.task.budget} model-call budget.`);
  }
}

const systemPrompt = `You are Piggy, PIG's internal CRM worker.
Use only the tools provided by the PIG application. You have no shell, filesystem, browser, or hidden tools.
Never invent facts, source URLs, affiliations, or email addresses. A claim is not stored unless pig_record_fact succeeds.
Every derived claim requires a source URL and a short evidence excerpt. If the task supplies insufficient evidence, say so and stop.
Be concise. In the final response, state what you stored, what you could not establish, and why.`;

function taskPrompt(task: AgentTask): string {
  return JSON.stringify(
    {
      kind: task.kind,
      subject: task.subject,
      reason: task.reason,
      payload: task.payload,
    },
    null,
    2,
  );
}

/**
 * Timeout and retry for both inference paths — the queued worker here and the
 * interactive chat in `chat.ts`.
 *
 * Neither had either. A hung upstream hung the chat until the browser gave up,
 * and because `PiggyWorker` renews its lease at half the lease interval for as
 * long as the model call is outstanding, one hung socket pinned a queued task
 * for the life of the process. `packages/prime/src/client.ts` already solved
 * this shape for the compute API — exponential backoff with full jitter,
 * `Retry-After` honoured when the server offers one, 429 and 5xx retried and
 * every other 4xx never — so this follows it rather than inventing a second
 * policy for the same upstream operator.
 *
 * The deadline is per attempt and covers exactly what the attempt awaits. The
 * worker awaits the whole JSON body inside it. The chat awaits only the
 * response headers, because a flat deadline over a streamed answer would kill
 * a legitimately long one; its stream is guarded by an idle timeout instead.
 */
export interface InferenceRetryPolicy {
  /** Total attempts, including the first. */
  maxAttempts: number;
  /** Deadline for a single attempt. */
  timeoutMs: number;
  /** Ceiling on the backoff between attempts. */
  maxBackoffMs: number;
  onRetry?: (info: { attempt: number; delayMs: number; reason: string }) => void;
}

export class PiggyInferenceError extends Error {
  constructor(
    message: string,
    /** Absent when the attempt never got a response at all. */
    readonly status?: number,
    /** What the server asked us to wait, when it said. */
    readonly retryAfterMs?: number,
  ) {
    super(message);
    this.name = 'PiggyInferenceError';
  }

  /** A 4xx that is not 429 will fail identically however often it is retried. */
  get isRetryable(): boolean {
    return this.status === undefined || this.status === 429 || this.status >= 500;
  }
}

/** Drains a failed response and turns it into the error the policy classifies on. */
export async function inferenceErrorFor(response: Response): Promise<PiggyInferenceError> {
  const body = (await response.text().catch(() => '')).slice(0, 500);
  return new PiggyInferenceError(
    `Piggy inference ${response.status}: ${body || response.statusText}`,
    response.status,
    parseRetryAfter(response.headers.get('retry-after')) ?? undefined,
  );
}

export async function withInferenceRetries<T>(
  policy: InferenceRetryPolicy,
  signal: AbortSignal | undefined,
  attempt: (attemptSignal: AbortSignal) => Promise<T>,
): Promise<T> {
  let lastError: unknown;

  for (let n = 1; n <= policy.maxAttempts; n += 1) {
    const deadline = new AbortController();
    const timer = setTimeout(
      () =>
        deadline.abort(
          new PiggyInferenceError(`Piggy inference did not respond within ${policy.timeoutMs}ms.`),
        ),
      policy.timeoutMs,
    );
    let delayMs: number | undefined;

    try {
      return await attempt(anySignal(signal, deadline.signal));
    } catch (error) {
      // The caller hung up — the browser navigated away, or the worker lost its
      // lease. Retrying would spend credit on an answer nobody will read.
      if (signal?.aborted) throw signal.reason ?? error;
      const retryable = !(error instanceof PiggyInferenceError) || error.isRetryable;
      if (!retryable || n === policy.maxAttempts) throw error;
      lastError = error;
      delayMs =
        (error instanceof PiggyInferenceError ? error.retryAfterMs : undefined) ??
        backoffMs(n, policy.maxBackoffMs);
      policy.onRetry?.({
        attempt: n,
        delayMs,
        reason: error instanceof Error ? error.message : 'network error',
      });
    } finally {
      clearTimeout(timer);
    }

    // Backing off outside the try keeps the attempt's deadline from outliving
    // the attempt it was guarding and aborting the next one on arrival.
    await sleep(delayMs ?? 0, signal);
  }

  throw lastError ?? new Error('Piggy inference request failed.');
}

/**
 * `AbortSignal.any([undefined])` throws, and the caller's signal is optional on
 * every path into inference, so the list is filtered rather than assumed dense.
 */
export function anySignal(...signals: (AbortSignal | undefined)[]): AbortSignal {
  return AbortSignal.any(signals.filter((signal): signal is AbortSignal => signal !== undefined));
}

/**
 * Exponential backoff with full jitter. Jitter matters more than the curve:
 * without it the worker and every open chat that hit the same rate limit retry
 * in lockstep and reproduce the limit that caused it.
 */
function backoffMs(attempt: number, ceilingMs: number): number {
  return Math.round(Math.random() * Math.min(ceilingMs, 1_000 * 2 ** (attempt - 1)));
}

function parseRetryAfter(header: string | null): number | null {
  if (!header) return null;
  const seconds = Number(header);
  if (Number.isFinite(seconds)) return Math.max(0, seconds * 1_000);
  const date = Date.parse(header);
  if (Number.isFinite(date)) return Math.max(0, date - Date.now());
  return null;
}

/** Sleeps, but wakes immediately if the caller gives up mid-backoff. */
export function sleep(ms: number, signal?: AbortSignal): Promise<void> {
  return new Promise((resolve, reject) => {
    if (signal?.aborted) {
      reject(signal.reason);
      return;
    }
    let timer: ReturnType<typeof setTimeout> | undefined;
    const onAbort = () => {
      clearTimeout(timer);
      reject(signal?.reason);
    };
    timer = setTimeout(() => {
      signal?.removeEventListener('abort', onAbort);
      resolve();
    }, ms);
    signal?.addEventListener('abort', onAbort, { once: true });
  });
}
