import type { SpendGuard } from '../spend-guard';
import {
  createFrameworkGuard,
  modelFromUnknown,
  objectRecord,
  outputTokensFromParams,
  preflightFrameworkCall,
  recordFrameworkReceipt,
  redactDataPlane,
  settleFrameworkCall,
  stringValue,
  type FrameworkAdapterOptions,
  type FrameworkPreflight,
} from './common';

/**
 * LangChain.js callback handler that runs `guard.decide()` before every LLM /
 * chat-model call and settles real token usage after it completes.
 *
 * Hook point: LangChain's `BaseCallbackHandler` lifecycle methods
 * `handleChatModelStart` / `handleLLMStart` (preflight, may throw
 * AgentGuardBlockedError) and `handleLLMEnd` (usage settlement). We do NOT
 * extend the real BaseCallbackHandler (that would force a peer-dep import at
 * module load); LangChain duck-types handlers by these method names, and we set
 * `awaitHandlers = true` so a thrown block aborts the call.
 *
 * This adapter reuses the shared framework core (createFrameworkGuard +
 * preflightFrameworkCall + settleFrameworkCall). It never proxies prompts,
 * completions, or provider keys — only redacted request shape leaves for the
 * signed receipt.
 *
 * @example
 *   import { ChatOpenAI } from '@langchain/openai';
 *   import { createLangChainHandler } from '@agentguard-run/spend/frameworks/langchain';
 *
 *   const handler = createLangChainHandler({ policy, scope });
 *   const llm = new ChatOpenAI({ callbacks: [handler] });
 *   await llm.invoke('hello'); // throws AgentGuardBlockedError if a cap blocks
 */
export interface LangChainAdapterOptions extends FrameworkAdapterOptions {
  framework?: 'langchain';
}

interface LangChainSerialized {
  id?: string[];
  kwargs?: Record<string, unknown>;
  [key: string]: unknown;
}

export class AgentGuardLangChainHandler {
  readonly name = 'agentguard_spend';
  /** Ensures LangChain awaits the handler so a thrown block actually aborts. */
  readonly awaitHandlers = true;
  readonly guard: SpendGuard;
  private readonly opts: LangChainAdapterOptions;
  private readonly pending = new Map<string, FrameworkPreflight>();

  constructor(opts: LangChainAdapterOptions) {
    this.opts = { ...opts, framework: 'langchain' };
    this.guard = createFrameworkGuard(this.opts);
  }

  async handleChatModelStart(
    llm: LangChainSerialized,
    messages: unknown[][],
    runId: string,
    _parentRunId?: string,
    extraParams?: Record<string, unknown>,
    _tags?: string[],
    metadata?: Record<string, unknown>,
  ): Promise<void> {
    await this.start(runId, llm, messages, extraParams, metadata);
  }

  async handleLLMStart(
    llm: LangChainSerialized,
    prompts: unknown[],
    runId: string,
    _parentRunId?: string,
    extraParams?: Record<string, unknown>,
    _tags?: string[],
    metadata?: Record<string, unknown>,
  ): Promise<void> {
    await this.start(runId, llm, prompts, extraParams, metadata);
  }

  async handleLLMEnd(output: unknown, runId: string): Promise<void> {
    const preflight = this.pending.get(runId);
    if (!preflight) return;
    this.pending.delete(runId);
    await settleFrameworkCall(preflight, normalizeUsage(output));
    await recordFrameworkReceipt(this.guard, 'langchain', {
      event: 'llm_end',
      decisionId: preflight.decision.decisionId,
      model: preflight.decision.modelResolved,
    });
  }

  async handleLLMError(_error: unknown, runId: string): Promise<void> {
    this.pending.delete(runId);
  }

  private async start(
    runId: string,
    llm: LangChainSerialized,
    payload: unknown,
    extraParams?: Record<string, unknown>,
    metadata?: Record<string, unknown>,
  ): Promise<void> {
    const invocation = objectRecord(extraParams?.invocation_params) ?? {};
    const model = modelFromUnknown(
      invocation.model ?? invocation.model_name ?? lastId(llm) ?? objectRecord(llm?.kwargs)?.model,
      this.opts.defaultModel ?? 'unknown',
    );
    const preflight = await preflightFrameworkCall(this.guard, this.opts, {
      framework: 'langchain',
      model,
      params: payload,
      outputTokens: outputTokensFromParams(invocation, this.opts.defaultOutputTokens ?? 1024),
      metadata: redactDataPlane(metadata ?? {}) as Record<string, unknown>,
      workflowId: stringValue(metadata?.workflowId) ?? stringValue(metadata?.thread_id),
      requestShape: {
        toolCount: Array.isArray(invocation.tools) ? invocation.tools.length : 0,
        maxTokens: invocation.max_tokens ?? invocation.maxTokens,
        stream: invocation.stream === true,
      },
    });
    this.pending.set(runId, preflight);
  }
}

export function createLangChainHandler(opts: LangChainAdapterOptions): AgentGuardLangChainHandler {
  return new AgentGuardLangChainHandler(opts);
}

export const agentguardLangChainHandler = createLangChainHandler;

function lastId(llm: LangChainSerialized | undefined): string | undefined {
  const id = llm?.id;
  return Array.isArray(id) && id.length ? id[id.length - 1] : undefined;
}

/**
 * Map a LangChain LLMResult into the `{ usage: { inputTokens, outputTokens } }`
 * shape that settleFrameworkCall expects. LangChain reports token usage on
 * `llmOutput.tokenUsage` (OpenAI) or `llmOutput.usage`.
 */
function normalizeUsage(output: unknown): { usage: Record<string, unknown> } | null {
  const record = objectRecord(output);
  const llmOutput = objectRecord(record?.llmOutput);
  const tokenUsage = objectRecord(llmOutput?.tokenUsage) ?? objectRecord(llmOutput?.usage);
  if (!tokenUsage) return null;
  return {
    usage: {
      inputTokens: tokenUsage.promptTokens ?? tokenUsage.prompt_tokens ?? tokenUsage.input_tokens ?? tokenUsage.inputTokens,
      outputTokens: tokenUsage.completionTokens ?? tokenUsage.completion_tokens ?? tokenUsage.output_tokens ?? tokenUsage.outputTokens,
    },
  };
}