import { AIError, type GenerateOptions, type Message } from "./index.ts";
import type { AIProvider, AIResult } from "./providers.ts";

export interface HttpAIProviderOptions {
  apiKey?: string;
  model: string;
  baseUrl?: string;
  fetch?: typeof fetch;
  headers?: HeadersInit;
}
function key(configured: string | undefined, name: string): string {
  const value = configured ?? process.env[name];
  if (!value) throw new AIError(`Missing ${name}.`, 0, "authentication_error");
  return value;
}
function promptMessages(prompt: string | Message[], options?: GenerateOptions) {
  const messages =
    options?.messages ??
    (typeof prompt === "string" ? [{ role: "user" as const, content: prompt }] : prompt);
  return options?.system ? [{ role: "system", content: options.system }, ...messages] : messages;
}

export function openAIProvider(options: HttpAIProviderOptions): AIProvider {
  const send = options.fetch ?? fetch;
  const base = (options.baseUrl ?? "https://api.openai.com/v1").replace(/\/$/, "");
  return {
    name: "openai",
    capabilities: { streaming: false, structuredOutput: true, tools: true, usage: true },
    async generate(prompt, call = {}) {
      const response = await send(`${base}/chat/completions`, {
        method: "POST",
        headers: {
          authorization: `Bearer ${key(options.apiKey, "OPENAI_API_KEY")}`,
          "content-type": "application/json",
          ...Object.fromEntries(new Headers(options.headers)),
        },
        body: JSON.stringify({
          model: call.model ?? options.model,
          messages: promptMessages(prompt, call),
          ...(call.maxTokens ? { max_completion_tokens: call.maxTokens } : {}),
        }),
        signal: call.signal,
      });
      if (!response.ok)
        throw new AIError(`OpenAI returned ${response.status}.`, response.status, "provider_error");
      const body = (await response.json()) as any;
      return {
        value: String(body.choices?.[0]?.message?.content ?? ""),
        provider: "openai",
        model: body.model,
        finishReason: body.choices?.[0]?.finish_reason,
        usage: {
          inputTokens: body.usage?.prompt_tokens,
          outputTokens: body.usage?.completion_tokens,
          totalTokens: body.usage?.total_tokens,
        },
        toolCalls: body.choices?.[0]?.message?.tool_calls?.map((tool: any) => ({
          id: String(tool.id),
          name: String(tool.function?.name),
          arguments: JSON.parse(tool.function?.arguments ?? "{}"),
        })),
        raw: body,
      };
    },
  };
}

export function googleAIProvider(options: HttpAIProviderOptions): AIProvider {
  const send = options.fetch ?? fetch;
  const base = (options.baseUrl ?? "https://generativelanguage.googleapis.com/v1beta").replace(
    /\/$/,
    "",
  );
  return {
    name: "google",
    capabilities: { streaming: false, structuredOutput: true, tools: true, usage: true },
    async generate(prompt, call = {}) {
      const response = await send(
        `${base}/models/${encodeURIComponent(call.model ?? options.model)}:generateContent?key=${encodeURIComponent(key(options.apiKey, "GOOGLE_AI_API_KEY"))}`,
        {
          method: "POST",
          headers: {
            "content-type": "application/json",
            ...Object.fromEntries(new Headers(options.headers)),
          },
          body: JSON.stringify({
            contents: promptMessages(prompt, call)
              .filter((message) => message.role !== "system")
              .map((message) => ({
                role: message.role === "assistant" ? "model" : "user",
                parts: [{ text: message.content }],
              })),
            ...(call.system ? { systemInstruction: { parts: [{ text: call.system }] } } : {}),
          }),
          signal: call.signal,
        },
      );
      if (!response.ok)
        throw new AIError(
          `Google AI returned ${response.status}.`,
          response.status,
          "provider_error",
        );
      const body = (await response.json()) as any;
      return {
        value: String(
          body.candidates?.[0]?.content?.parts?.map((part: any) => part.text ?? "").join("") ?? "",
        ),
        provider: "google",
        model: call.model ?? options.model,
        finishReason: body.candidates?.[0]?.finishReason,
        usage: {
          inputTokens: body.usageMetadata?.promptTokenCount,
          outputTokens: body.usageMetadata?.candidatesTokenCount,
          totalTokens: body.usageMetadata?.totalTokenCount,
        },
        raw: body,
      };
    },
  };
}

/** OpenAI-compatible local servers including Ollama, llama.cpp and vLLM. */
export function localAIProvider(
  options: Omit<HttpAIProviderOptions, "apiKey"> & { apiKey?: string },
): AIProvider {
  const provider = openAIProvider({
    ...options,
    apiKey: options.apiKey ?? "local",
    baseUrl: options.baseUrl ?? "http://localhost:11434/v1",
  });
  return {
    ...provider,
    name: "local",
    generate: async (prompt, call) => ({
      ...(await provider.generate(prompt, call)),
      provider: "local",
    }),
  };
}

export interface EmbeddingProvider {
  name: string;
  embed(
    values: string[],
    options?: { model?: string; signal?: AbortSignal },
  ): Promise<{ vectors: number[][]; usage?: { tokens?: number } }>;
}
export function openAIEmbeddings(options: HttpAIProviderOptions): EmbeddingProvider {
  return {
    name: "openai",
    async embed(values, call = {}) {
      const response = await (options.fetch ?? fetch)(
        `${(options.baseUrl ?? "https://api.openai.com/v1").replace(/\/$/, "")}/embeddings`,
        {
          method: "POST",
          headers: {
            authorization: `Bearer ${key(options.apiKey, "OPENAI_API_KEY")}`,
            "content-type": "application/json",
          },
          body: JSON.stringify({ model: call.model ?? options.model, input: values }),
          signal: call.signal,
        },
      );
      if (!response.ok)
        throw new AIError(`Embedding provider returned ${response.status}.`, response.status);
      const body = (await response.json()) as any;
      return {
        vectors: body.data.map((item: any) => item.embedding as number[]),
        usage: { tokens: body.usage?.total_tokens },
      };
    },
  };
}

export interface VectorRecord<T = Record<string, unknown>> {
  id: string;
  vector: number[];
  text: string;
  metadata: T;
}
export interface VectorMatch<T = Record<string, unknown>> extends VectorRecord<T> {
  score: number;
}
export interface VectorStore<T = Record<string, unknown>> {
  upsert(records: VectorRecord<T>[]): Promise<void>;
  query(
    vector: number[],
    limit?: number,
    filter?: (metadata: T) => boolean,
  ): Promise<VectorMatch<T>[]>;
  delete(ids: string[]): Promise<void>;
}
export function memoryVectorStore<T = Record<string, unknown>>(): VectorStore<T> {
  const records = new Map<string, VectorRecord<T>>();
  const cosine = (a: number[], b: number[]) => {
    if (a.length !== b.length || !a.length) return 0;
    let dot = 0,
      aa = 0,
      bb = 0;
    for (let i = 0; i < a.length; i++) {
      dot += a[i]! * b[i]!;
      aa += a[i]! ** 2;
      bb += b[i]! ** 2;
    }
    return aa && bb ? dot / Math.sqrt(aa * bb) : 0;
  };
  return {
    async upsert(values) {
      for (const value of values) records.set(value.id, { ...value, vector: [...value.vector] });
    },
    async query(vector, limit = 5, filter) {
      return [...records.values()]
        .filter((record) => !filter || filter(record.metadata))
        .map((record) => ({ ...record, score: cosine(vector, record.vector) }))
        .sort((a, b) => b.score - a.score)
        .slice(0, Math.max(0, limit));
    },
    async delete(ids) {
      for (const id of ids) records.delete(id);
    },
  };
}

export function createRagPipeline<T>(options: {
  embeddings: EmbeddingProvider;
  store: VectorStore<T>;
  generate: (prompt: string, options?: GenerateOptions) => Promise<AIResult<string>>;
  maxContextCharacters?: number;
}) {
  return {
    async index(documents: Array<{ id: string; text: string; metadata: T }>) {
      const embedded = await options.embeddings.embed(documents.map((document) => document.text));
      await options.store.upsert(
        documents.map((document, index) => ({ ...document, vector: embedded.vectors[index]! })),
      );
      return documents.length;
    },
    async ask(
      question: string,
      call: GenerateOptions & { limit?: number; filter?: (metadata: T) => boolean } = {},
    ) {
      const embedded = await options.embeddings.embed([question], { signal: call.signal });
      const matches = await options.store.query(embedded.vectors[0]!, call.limit, call.filter);
      const context = matches
        .map((match, index) => `[${index + 1}] ${match.text}`)
        .join("\n\n")
        .slice(0, options.maxContextCharacters ?? 12_000);
      const result = await options.generate(
        `Answer using only the supplied context. Cite sources as [n].\n\nContext:\n${context}\n\nQuestion: ${question}`,
        call,
      );
      return { ...result, sources: matches };
    },
  };
}

export interface ConversationStore {
  load(id: string): Promise<Message[]>;
  append(id: string, messages: Message[]): Promise<void>;
  clear(id: string): Promise<void>;
}
export function memoryConversationStore(maxMessages = 100): ConversationStore {
  const conversations = new Map<string, Message[]>();
  return {
    async load(id) {
      return [...(conversations.get(id) ?? [])];
    },
    async append(id, messages) {
      conversations.set(id, [...(conversations.get(id) ?? []), ...messages].slice(-maxMessages));
    },
    async clear(id) {
      conversations.delete(id);
    },
  };
}
export function promptTemplate(template: string) {
  return (variables: Record<string, string | number>) =>
    template.replace(/\{\{\s*([\w.-]+)\s*\}\}/g, (_match, name: string) => {
      if (!(name in variables)) throw new Error(`WRN-AI-PROMPT-VARIABLE:${name}`);
      return String(variables[name]);
    });
}
export type AIGuardrail = (input: {
  prompt: string | Message[];
  output?: string;
}) => void | Promise<void>;
export const maxPromptLength =
  (maximum: number): AIGuardrail =>
  ({ prompt }) => {
    const length =
      typeof prompt === "string"
        ? prompt.length
        : prompt.reduce((sum, message) => sum + message.content.length, 0);
    if (length > maximum)
      throw new AIError("Prompt exceeds configured guardrail.", 400, "guardrail");
  };
export function guardedProvider(provider: AIProvider, guardrails: AIGuardrail[]): AIProvider {
  return {
    ...provider,
    async generate(prompt, options) {
      for (const guardrail of guardrails) await guardrail({ prompt });
      const result = await provider.generate(prompt, options);
      for (const guardrail of guardrails) await guardrail({ prompt, output: result.value });
      return result;
    },
  };
}
export function aiRateLimiter(options: { limit: number; windowMs: number; now?: () => number }) {
  const buckets = new Map<string, { count: number; reset: number }>();
  const now = options.now ?? Date.now;
  return (key: string) => {
    const time = now();
    const bucket = buckets.get(key);
    if (!bucket || bucket.reset <= time) {
      buckets.set(key, { count: 1, reset: time + options.windowMs });
      return { allowed: true, remaining: options.limit - 1 };
    }
    if (bucket.count >= options.limit)
      return { allowed: false, remaining: 0, retryAfterMs: bucket.reset - time };
    bucket.count++;
    return { allowed: true, remaining: options.limit - bucket.count };
  };
}
export async function evaluateAI(
  cases: Array<{
    name: string;
    prompt: string;
    expected?: string;
    score?: (output: string) => number | Promise<number>;
  }>,
  generate: (prompt: string) => Promise<string>,
) {
  const results = [];
  for (const item of cases) {
    const started = performance.now();
    const output = await generate(item.prompt);
    const score = item.score
      ? await item.score(output)
      : item.expected === undefined
        ? 1
        : output.includes(item.expected)
          ? 1
          : 0;
    results.push({ name: item.name, output, score, durationMs: performance.now() - started });
  }
  return {
    results,
    score: results.length
      ? results.reduce((sum, result) => sum + result.score, 0) / results.length
      : 0,
  };
}