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;
  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 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;
    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) {
      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: request.signal,
      });

      if (!response.ok) {
        const body = await response.text().catch(() => '');
        throw new Error(
          `Piggy inference ${response.status}: ${body.slice(0, 500) || response.statusText}`,
        );
      }

      const completion = completionSchema.parse(await response.json());
      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,
  );
}
