/**
 * AgentAudit TypeScript SDK — v3. Single-file, zero-dependency.
 *
 * Two ways to use it:
 *
 * (A) Auto-instrumentation — RECOMMENDED. Wrap your existing SDK client and
 *     every LLM / tool call inside the trace context is captured automatically.
 *
 *       import OpenAI from "openai";
 *       import { createAgentAudit } from "@agentaudit/sdk";
 *
 *       const aa     = createAgentAudit({ apiKey: process.env.AGENTAUDIT_KEY! });
 *       const openai = aa.wrapOpenAI(new OpenAI());
 *       const search = aa.tool("web_search", async (q: string) => fetchResults(q));
 *
 *       await aa.withTrace(
 *         { agent_name: "billing-copilot", external_user_id: user.id },
 *         async () => {
 *           const hits = await search("acme invoice");
 *           return openai.chat.completions.create({
 *             model: "gpt-5.5",
 *             messages: [{ role: "user", content: hits[0].snippet }],
 *           });
 *         },
 *       );
 *
 *     Supported wrappers: wrapOpenAI, wrapAnthropic, wrapFetch (generic HTTP),
 *     wrapLangChain (callback handler), tool()/traced() for arbitrary functions.
 *
 * (B) Manual — full control. See `aa.trace(...).llm() / .tool() / .end()`.
 *
 * Runtimes: Node 18+, Deno, Bun, Cloudflare Workers, Vercel Edge. Keys stay
 * server-side. Uses AsyncLocalStorage where available; falls back to a
 * synchronous "current trace" pointer elsewhere (document your concurrency).
 */


export type Scope = "telemetry" | "evaluations" | "audit" | "governance";

export type AgentAuditConfig = {
  apiKey: string;
  baseUrl?: string;
  timeoutMs?: number;
  fetch?: typeof fetch;
};

export type Message = {
  role: "system" | "user" | "assistant" | "tool" | "developer";
  content?: string | unknown[];
  name?: string;
  tool_call_id?: string;
  tool_calls?: unknown[];
};

export type SpanKind = "agent" | "llm" | "tool" | "retrieval" | "guardrail" | "other";

export type SpanInit = {
  name: string;
  kind?: SpanKind;
  parent_span_id?: string;
  model?: string;
  provider?: string;
  input?: unknown;
  metadata?: Record<string, unknown>;
};

export type Usage = {
  prompt_tokens?: number;
  completion_tokens?: number;
  total_tokens?: number;
  cost_usd?: number;
};

export type TraceInit = {
  trace_id?: string;
  agent_id?: string;
  agent_name?: string;
  session_id?: string;
  external_user_id?: string;
  metadata?: Record<string, unknown>;
};

export type LegacyTelemetryEvent = {
  trace_id: string;
  outcome: "success" | "failure" | "escalated" | "blocked" | string;
  latency_ms?: number;
  tool_count?: number;
  detail?: string;
  agent_id?: string;
  agent_name?: string;
  model?: string;
  provider?: string;
  prompt_tokens?: number;
  completion_tokens?: number;
  total_tokens?: number;
  cost_usd?: number;
  session_id?: string;
  external_user_id?: string;
};

export type EvaluationReport = {
  name: string;
  agent_id?: string;
  agent_name?: string;
  harness_id?: string;
  category?: string;
  status?: "scheduled" | "running" | "passing" | "failing";
  pass_rate?: number;
  notes?: string;
  severity_threshold?: number;
};

export type AuditEntry = {
  actor: string;
  action: string;
  subject?: string;
  rationale?: string;
  framework?: string;
  trace_id?: string;
  agent_id?: string;
  retention_until?: string;
  metadata?: Record<string, unknown>;
};

export type PromotionCheckInput = { agent_id: string; harness_id?: string; min_pass_rate?: number; max_age_hours?: number };
export type PromotionCheckResult = { allowed: boolean; reason: string; message?: string; threshold?: number; evaluation?: any };

export class AgentAuditError extends Error {
  constructor(public status: number, message: string, public body?: unknown) { super(`[AgentAudit ${status}] ${message}`); }
}

const DEFAULT_BASE = "https://agentauditcentral.co.uk";

function uid(prefix: string) { return `${prefix}_${Date.now().toString(36)}${Math.random().toString(36).slice(2, 10)}`; }

export function createAgentAudit(config: AgentAuditConfig) {
  if (!config.apiKey || !config.apiKey.startsWith("aa_live_")) {
    throw new Error("AgentAudit: apiKey must be an aa_live_… ingestion key");
  }
  const baseUrl = (config.baseUrl ?? DEFAULT_BASE).replace(/\/$/, "");
  const timeoutMs = config.timeoutMs ?? 15_000;
  const doFetch = config.fetch ?? globalThis.fetch;

  async function request<T>(path: string, body: unknown): Promise<T> {
    const ctrl = new AbortController();
    const t = setTimeout(() => ctrl.abort(), timeoutMs);
    try {
      const res = await doFetch(`${baseUrl}${path}`, {
        method: "POST",
        headers: { "content-type": "application/json", authorization: `Bearer ${config.apiKey}` },
        body: JSON.stringify(body),
        signal: ctrl.signal,
      });
      const txt = await res.text();
      const json = txt ? safeJson(txt) : {};
      if (!res.ok) {
        const msg = (json as { error?: string })?.error || res.statusText;
        throw new AgentAuditError(res.status, msg, json);
      }
      return json as T;
    } finally { clearTimeout(t); }
  }

  function trace(init: TraceInit = {}) {
    const trace_id = init.trace_id ?? uid("trace");
    const started = new Date();
    const spans: any[] = [];

    function span(spec: SpanInit) {
      const span_id = uid("span");
      const startedAt = new Date();
      const row: any = {
        span_id,
        parent_span_id: spec.parent_span_id,
        kind: spec.kind ?? "other",
        name: spec.name,
        status: "success",
        model: spec.model,
        provider: spec.provider,
        input: spec.input,
        metadata: spec.metadata,
        started_at: startedAt.toISOString(),
      };
      const api = {
        span_id,
        setInput(v: unknown)      { row.input = v; return api; },
        setOutput(v: unknown)     { row.output = v; return api; },
        setMessages(m: Message[]) { row.messages = m; return api; },
        setModel(m: string, provider?: string) { row.model = m; if (provider) row.provider = provider; return api; },
        setUsage(u: Usage)        { Object.assign(row, u); return api; },
        setArguments(v: unknown)  { row.tool_arguments = v; return api; },
        setResult(v: unknown)     { row.tool_result = v; return api; },
        setError(e: unknown)      { row.error = e instanceof Error ? e.message : String(e); row.status = "error"; return api; },
        setStatus(s: string)      { row.status = s; return api; },
        setMetadata(v: Record<string, unknown>) { row.metadata = { ...(row.metadata ?? {}), ...v }; return api; },
        end() {
          const endedAt = new Date();
          row.ended_at = endedAt.toISOString();
          row.latency_ms = endedAt.getTime() - startedAt.getTime();
          spans.push(row);
          return api;
        },
      };
      return api;
    }

    return {
      trace_id,
      llm:       (s: Omit<SpanInit, "kind">) => span({ ...s, kind: "llm" }),
      tool:      (s: Omit<SpanInit, "kind"> & { arguments?: unknown }) => {
        const sp = span({ ...s, kind: "tool" });
        if (s.arguments !== undefined) sp.setArguments(s.arguments);
        return sp;
      },
      retrieval: (s: Omit<SpanInit, "kind">) => span({ ...s, kind: "retrieval" }),
      guardrail: (s: Omit<SpanInit, "kind">) => span({ ...s, kind: "guardrail" }),
      span,
      /** Submit the full trace with all spans in one call. */
      async end(final: { outcome?: string; detail?: string; error?: string; latency_ms?: number; metadata?: Record<string, unknown> } = {}) {
        const ended = new Date();
        return request<{ ok: true; id: string; trace_id: string; spans_ingested: number; tokens: number | null; cost_usd: number | null }>(
          "/api/public/traces",
          {
            trace_id,
            agent_id: init.agent_id,
            agent_name: init.agent_name,
            session_id: init.session_id,
            external_user_id: init.external_user_id,
            outcome: final.outcome ?? "success",
            detail: final.detail,
            error: final.error,
            latency_ms: final.latency_ms ?? ended.getTime() - started.getTime(),
            started_at: started.toISOString(),
            ended_at: ended.toISOString(),
            metadata: { ...(init.metadata ?? {}), ...(final.metadata ?? {}) },
            spans,
          },
        );
      },
    };
  }

  // ─────────────────────────────────────────────────────────────────────────
  // Auto-instrumentation layer
  // ─────────────────────────────────────────────────────────────────────────

  type ActiveTrace = ReturnType<typeof trace>;
  type Ctx = { t: ActiveTrace; parent_span_id?: string };

  // Prefer AsyncLocalStorage where available; fall back to a mutable pointer.
  let alsPromise: Promise<any> | null = null;
  async function getStore(): Promise<any> {
    if (alsPromise) return alsPromise;
    alsPromise = (async () => {
      try {
        const mod = await import(/* @vite-ignore */ "node:async_hooks");
        return new mod.AsyncLocalStorage<Ctx>();
      } catch { return null; }
    })();
    return alsPromise;
  }
  let fallbackCtx: Ctx | undefined;
  function currentSync(): Ctx | undefined { return fallbackCtx; }
  async function runInCtx<T>(ctx: Ctx, fn: () => Promise<T>): Promise<T> {
    const store = await getStore();
    if (store) return store.run(ctx, fn);
    const prev = fallbackCtx; fallbackCtx = ctx;
    try { return await fn(); } finally { fallbackCtx = prev; }
  }
  async function current(): Promise<Ctx | undefined> {
    const store = await getStore();
    return store ? store.getStore() : fallbackCtx;
  }

  // Simple per-1K-token price table (USD). Override with setPricing().
  const pricing: Record<string, { in: number; out: number }> = {
    "gpt-5.5":            { in: 0.0025, out: 0.010 },
    "gpt-4o":             { in: 0.0025, out: 0.010 },
    "gpt-4o-mini":        { in: 0.00015, out: 0.0006 },
    "o4-mini":            { in: 0.0011, out: 0.0044 },
    "claude-sonnet-4":    { in: 0.003,  out: 0.015 },
    "claude-3-5-sonnet":  { in: 0.003,  out: 0.015 },
    "claude-3-5-haiku":   { in: 0.0008, out: 0.004 },
    "gemini-2.5-pro":     { in: 0.00125, out: 0.010 },
  };
  function estimateCost(model: string | undefined, pt = 0, ct = 0): number | undefined {
    if (!model) return undefined;
    const key = Object.keys(pricing).find((k) => model.toLowerCase().includes(k));
    if (!key) return undefined;
    const p = pricing[key];
    return +(((pt / 1000) * p.in) + ((ct / 1000) * p.out)).toFixed(6);
  }
  function setPricing(model: string, price: { in: number; out: number }) { pricing[model.toLowerCase()] = price; }

  /**
   * Run `fn` inside a trace context. Any wrapped client used inside `fn`
   * (wrapOpenAI, wrapAnthropic, tool(), traced(), fetch wrapper) auto-attaches
   * its spans to this trace. The trace is submitted on completion (or on
   * error, marked failure). Returns whatever `fn` returns.
   */
  async function withTrace<T>(init: TraceInit & { outcome?: string }, fn: () => Promise<T>): Promise<T> {
    const t = trace(init);
    return runInCtx({ t }, async () => {
      try {
        const result = await fn();
        await t.end({ outcome: init.outcome ?? "success" });
        return result;
      } catch (e) {
        await t.end({ outcome: "failure", error: e instanceof Error ? e.message : String(e) });
        throw e;
      }
    });
  }

  /** Wrap any async function as a `tool` span. Auto-captures args/result/error. */
  function traced<A extends unknown[], R>(name: string, fn: (...args: A) => Promise<R>) {
    return (async (...args: A): Promise<R> => {
      const ctx = await current();
      if (!ctx) return fn(...args); // outside a trace — no-op
      const sp = ctx.t.tool({ name, arguments: args.length === 1 ? args[0] : args });
      try {
        const out = await fn(...args);
        sp.setResult(out).end();
        return out;
      } catch (e) { sp.setError(e).end(); throw e; }
    }) as (...args: A) => Promise<R>;
  }
  /** Alias — reads better for domain functions ("tool" implies LLM-callable). */
  const toolFn = traced;

  /** Manual span open (useful for retrieval / guardrail). Call `.end()` on the result. */
  async function span(kind: SpanKind, name: string, metadata?: Record<string, unknown>) {
    const ctx = await current();
    if (!ctx) return null;
    const sp = ctx.t.span({ kind, name, metadata });
    return sp;
  }

  // OpenAI wrapper — supports chat.completions.create and responses.create.
  function wrapOpenAI<T extends object>(client: T): T {
    return new Proxy(client, {
      get(target, prop, receiver) {
        const value = Reflect.get(target, prop, receiver);
        if (prop === "chat") return wrapPath(value, "chat.completions.create");
        if (prop === "responses") return wrapPath(value, "responses.create");
        if (prop === "completions") return wrapPath(value, "completions.create");
        return value;
      },
    });
  }
  function wrapPath(obj: any, dottedMethod: string): any {
    const [head, ...rest] = dottedMethod.split(".");
    if (rest.length === 0) {
      const orig = obj[head]?.bind(obj);
      if (typeof orig !== "function") return obj;
      return new Proxy(obj, {
        get(t, p) { return p === head ? instrumentChatCall(orig, dottedMethod) : (t as any)[p]; },
      });
    }
    const inner = obj[head];
    if (!inner) return obj;
    return new Proxy(obj, {
      get(t, p) { return p === head ? wrapPath(inner, rest.join(".")) : (t as any)[p]; },
    });
  }
  function instrumentChatCall(orig: (...args: any[]) => any, name: string) {
    return async function instrumented(this: any, params: any, ...rest: any[]) {
      const ctx = await current();
      if (!ctx) return orig.call(this, params, ...rest);
      const model = params?.model;
      const sp = ctx.t.llm({ name, model, provider: "openai" });
      sp.setMessages(params?.messages ?? params?.input ?? []).setInput(redact(params));
      try {
        const res = await orig.call(this, params, ...rest);
        // Streaming? tap the async iterator without breaking user consumption
        if (res && typeof res[Symbol.asyncIterator] === "function") return tapOpenAIStream(res, sp, model);
        const usage = res?.usage ?? {};
        const pt = usage.prompt_tokens ?? usage.input_tokens;
        const ct = usage.completion_tokens ?? usage.output_tokens;
        sp.setOutput(extractOpenAIText(res))
          .setUsage({ prompt_tokens: pt, completion_tokens: ct, total_tokens: usage.total_tokens, cost_usd: estimateCost(model, pt, ct) })
          .end();
        return res;
      } catch (e) { sp.setError(e).end(); throw e; }
    };
  }
  function extractOpenAIText(res: any): unknown {
    return res?.choices?.[0]?.message ?? res?.output_text ?? res?.output ?? res;
  }
  async function* tapOpenAIStream(iter: any, sp: any, model: string | undefined) {
    let text = "", pt = 0, ct = 0;
    try {
      for await (const chunk of iter) {
        const delta = chunk?.choices?.[0]?.delta?.content ?? chunk?.delta ?? "";
        if (typeof delta === "string") text += delta;
        if (chunk?.usage) { pt = chunk.usage.prompt_tokens ?? pt; ct = chunk.usage.completion_tokens ?? ct; }
        yield chunk;
      }
      sp.setOutput(text).setUsage({ prompt_tokens: pt, completion_tokens: ct, cost_usd: estimateCost(model, pt, ct) }).end();
    } catch (e) { sp.setError(e).end(); throw e; }
  }

  // Anthropic wrapper — messages.create.
  function wrapAnthropic<T extends object>(client: T): T {
    return new Proxy(client, {
      get(target, prop, receiver) {
        if (prop === "messages") {
          const messages = Reflect.get(target, prop, receiver);
          return new Proxy(messages, {
            get(t: any, p) {
              if (p !== "create") return t[p];
              const orig = t.create.bind(t);
              return async function (this: any, params: any) {
                const ctx = await current();
                if (!ctx) return orig(params);
                const sp = ctx.t.llm({ name: "messages.create", model: params?.model, provider: "anthropic" });
                sp.setMessages(params?.messages ?? []).setInput(redact(params));
                try {
                  const res = await orig(params);
                  const u = res?.usage ?? {};
                  const pt = u.input_tokens, ct = u.output_tokens;
                  sp.setOutput(res?.content ?? res)
                    .setUsage({ prompt_tokens: pt, completion_tokens: ct, cost_usd: estimateCost(params?.model, pt, ct) })
                    .end();
                  return res;
                } catch (e) { sp.setError(e).end(); throw e; }
              };
            },
          });
        }
        return Reflect.get(target, prop, receiver);
      },
    });
  }

  /**
   * Generic fetch wrapper. Recognises OpenAI/Anthropic/Gemini/OpenRouter
   * endpoints from the URL and records prompts, model, usage. Useful if you
   * call providers by raw HTTP instead of an SDK.
   */
  function wrapFetch(baseFetch: typeof fetch = globalThis.fetch): typeof fetch {
    return (async (input: RequestInfo | URL, init?: RequestInit) => {
      const url = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url;
      const method = (init?.method ?? "GET").toUpperCase();
      const provider = detectProvider(url);
      const ctx = provider && method === "POST" ? await current() : undefined;
      if (!ctx) return baseFetch(input as any, init);
      const bodyStr = typeof init?.body === "string" ? init.body : undefined;
      const body = bodyStr ? safeJson(bodyStr) as any : undefined;
      const model = body?.model;
      const sp = ctx.t.llm({ name: `${provider}:http`, model, provider });
      if (body?.messages) sp.setMessages(body.messages);
      sp.setInput(redact(body));
      try {
        const res = await baseFetch(input as any, init);
        const clone = res.clone();
        try {
          const json: any = await clone.json();
          const u = json?.usage ?? {};
          const pt = u.prompt_tokens ?? u.input_tokens;
          const ct = u.completion_tokens ?? u.output_tokens;
          sp.setOutput(extractOpenAIText(json)).setUsage({ prompt_tokens: pt, completion_tokens: ct, cost_usd: estimateCost(model, pt, ct) }).end();
        } catch { sp.end(); }
        return res;
      } catch (e) { sp.setError(e).end(); throw e; }
    }) as typeof fetch;
  }
  function detectProvider(url: string): string | null {
    const u = url.toLowerCase();
    if (u.includes("api.openai.com")) return "openai";
    if (u.includes("api.anthropic.com")) return "anthropic";
    if (u.includes("generativelanguage.googleapis.com")) return "google";
    if (u.includes("openrouter.ai")) return "openrouter";
    if (u.includes("api.mistral.ai")) return "mistral";
    if (u.includes("api.groq.com")) return "groq";
    return null;
  }

  /**
   * LangChain / LangGraph callback handler. Pass into `.invoke(input, { callbacks: [aa.langchainCallbackHandler()] })`
   * (or into `RunnableConfig`). Captures LLM and tool runs.
   */
  function langchainCallbackHandler() {
    const spans = new Map<string, any>();
    const get = () => current();
    return {
      handleLLMStart: async (llm: any, prompts: string[], runId: string, _parent: any, extra: any) => {
        const ctx = await get(); if (!ctx) return;
        const model = extra?.invocation_params?.model ?? llm?.id?.[llm.id.length - 1];
        const sp = ctx.t.llm({ name: "llm", model, provider: extra?.invocation_params?._type ?? "langchain" });
        sp.setInput({ prompts });
        spans.set(runId, sp);
      },
      handleLLMEnd: async (out: any, runId: string) => {
        const sp = spans.get(runId); if (!sp) return;
        const u = out?.llmOutput?.tokenUsage ?? out?.llmOutput?.usage ?? {};
        sp.setOutput(out?.generations).setUsage({ prompt_tokens: u.promptTokens ?? u.prompt_tokens, completion_tokens: u.completionTokens ?? u.completion_tokens }).end();
      },
      handleLLMError: async (err: Error, runId: string) => { spans.get(runId)?.setError(err).end(); },
      handleToolStart: async (tool: any, input: string, runId: string) => {
        const ctx = await get(); if (!ctx) return;
        const sp = ctx.t.tool({ name: tool?.name ?? "tool", arguments: input });
        spans.set(runId, sp);
      },
      handleToolEnd: async (output: string, runId: string) => { spans.get(runId)?.setResult(output).end(); },
      handleToolError: async (err: Error, runId: string) => { spans.get(runId)?.setError(err).end(); },
    };
  }

  return {
    /** Rich trace API — manual. */
    trace,

    // Auto-instrumentation
    withTrace,
    wrapOpenAI,
    wrapAnthropic,
    wrapFetch,
    langchainCallbackHandler,
    tool: toolFn,
    traced,
    span,
    setPricing,

    /** Legacy run-summary telemetry (no per-step data). Prefer trace(). */
    telemetry: {
      async emit(event: LegacyTelemetryEvent) { return request("/api/public/telemetry", event); },
      async emitBatch(events: LegacyTelemetryEvent[]) { return request("/api/public/telemetry", { events }); },
    },

    evaluations: {
      async report(r: EvaluationReport): Promise<{ ok: true; id: string }> { return request("/api/public/evaluations", r); },
    },

    audit: {
      async record(entry: AuditEntry): Promise<{ ok: true; id: string; row_hash: string; prev_hash: string | null }> {
        return request("/api/public/audit", entry);
      },
    },

    governance: {
      async promotionCheck(input: PromotionCheckInput): Promise<PromotionCheckResult> {
        return request("/api/public/governance/promotion-check", input);
      },
    },
  };
}

/** Best-effort redaction of obviously-sensitive keys before capture. */
function redact(v: any): any {
  if (!v || typeof v !== "object") return v;
  const out: any = Array.isArray(v) ? [] : {};
  for (const k of Object.keys(v)) {
    if (/api[-_]?key|authorization|password|secret|token/i.test(k)) out[k] = "[redacted]";
    else out[k] = v[k];
  }
  return out;
}


function safeJson(txt: string): unknown { try { return JSON.parse(txt); } catch { return { raw: txt }; } }

export type AgentAuditClient = ReturnType<typeof createAgentAudit>;