import { Injectable, Logger } from '@nestjs/common';
import { ConfigService } from '@nestjs/config';

/**
 * Client for the AI Gateway, whose API mirrors Dify's workflow API.
 * The gateway is currently a passthrough — once the real gateway lands, only
 * env vars change.
 *
 * Two ops we need:
 *   - run     : POST /v1/workflows/run — JSON in, JSON out (blocking)
 *   - stream  : POST /v1/workflows/run with response_mode=streaming, returns SSE
 *
 * Auth: Bearer <AI_GATEWAY_API_KEY>.
 */
export interface GatewayRunInput {
  workflowId: string;
  inputs: Record<string, unknown>;
  user: string;            // tenant scope identifier (tenantId:userId)
  responseMode?: 'blocking' | 'streaming';
  conversationId?: string;
}

export interface GatewayRunResult {
  outputs: Record<string, unknown>;
  text: string;
  metadata: {
    promptTokens: number;
    completionTokens: number;
    totalTokens: number;
    costCents: number;
    raw?: unknown;
  };
}

export type GatewayStreamChunk =
  | { type: 'token'; data: string }
  | { type: 'done'; data: { outputs: Record<string, unknown>; text: string; usage: GatewayRunResult['metadata'] } }
  | { type: 'error'; data: { message: string } };

@Injectable()
export class AiGatewayClient {
  private readonly logger = new Logger(AiGatewayClient.name);

  constructor(private readonly config: ConfigService) {}

  /**
   * Blocking call. Throws on HTTP/network/gateway error so callers can catch
   * and record the failure on AgentInvocation.
   */
  async run(input: GatewayRunInput): Promise<GatewayRunResult> {
    const { url, key } = this.endpoint();
    const res = await fetch(`${url}/v1/workflows/run`, {
      method: 'POST',
      headers: this.headers(key),
      body: JSON.stringify({
        workflow_id: input.workflowId,
        inputs: input.inputs,
        user: input.user,
        response_mode: 'blocking',
        conversation_id: input.conversationId,
      }),
    });
    if (!res.ok) {
      const text = await res.text().catch(() => '');
      throw new Error(`AI Gateway ${res.status}: ${text || res.statusText}`);
    }
    const json = (await res.json()) as Record<string, unknown>;
    return this.normalize(json);
  }

  /**
   * Streaming call — yields parsed chunks. The gateway returns Dify-shaped
   * SSE: `data: { event, ... }` lines. We translate to a small union type.
   */
  async *stream(input: GatewayRunInput): AsyncGenerator<GatewayStreamChunk, void, void> {
    const { url, key } = this.endpoint();
    const res = await fetch(`${url}/v1/workflows/run`, {
      method: 'POST',
      headers: this.headers(key),
      body: JSON.stringify({
        workflow_id: input.workflowId,
        inputs: input.inputs,
        user: input.user,
        response_mode: 'streaming',
        conversation_id: input.conversationId,
      }),
    });
    if (!res.ok || !res.body) {
      const text = res.body ? await res.text().catch(() => '') : '';
      yield { type: 'error', data: { message: `AI Gateway ${res.status}: ${text || res.statusText}` } };
      return;
    }

    const decoder = new TextDecoder();
    const reader = res.body.getReader();
    let buffer = '';
    let collectedText = '';
    let lastUsage: GatewayRunResult['metadata'] = {
      promptTokens: 0,
      completionTokens: 0,
      totalTokens: 0,
      costCents: 0,
    };
    let lastOutputs: Record<string, unknown> = {};

    try {
      while (true) {
        const { value, done } = await reader.read();
        if (done) break;
        buffer += decoder.decode(value, { stream: true });
        let idx;
        while ((idx = buffer.indexOf('\n')) !== -1) {
          const line = buffer.slice(0, idx).trim();
          buffer = buffer.slice(idx + 1);
          if (!line.startsWith('data:')) continue;
          const payload = line.slice(5).trim();
          if (!payload || payload === '[DONE]') continue;
          let event: Record<string, unknown>;
          try {
            event = JSON.parse(payload) as Record<string, unknown>;
          } catch {
            continue;
          }
          const evType = event['event'];
          if (evType === 'text_chunk' || evType === 'message') {
            const data = (event['data'] as Record<string, unknown> | undefined) ?? event;
            const text = (data['text'] as string | undefined) ?? (event['answer'] as string | undefined) ?? '';
            if (text) {
              collectedText += text;
              yield { type: 'token', data: text };
            }
          } else if (evType === 'workflow_finished' || evType === 'message_end') {
            const data = (event['data'] as Record<string, unknown> | undefined) ?? event;
            lastOutputs = (data['outputs'] as Record<string, unknown> | undefined) ?? {};
            const usage = (data['metadata'] as Record<string, unknown> | undefined)?.['usage'] as
              | Record<string, unknown>
              | undefined;
            if (usage) {
              lastUsage = {
                promptTokens: Number(usage['prompt_tokens'] ?? 0),
                completionTokens: Number(usage['completion_tokens'] ?? 0),
                totalTokens: Number(usage['total_tokens'] ?? 0),
                costCents: Number(usage['cost_cents'] ?? 0),
                raw: usage,
              };
            }
          } else if (evType === 'error') {
            const data = (event['data'] as Record<string, unknown> | undefined) ?? event;
            yield {
              type: 'error',
              data: { message: (data['message'] as string | undefined) ?? 'gateway error' },
            };
            return;
          }
        }
      }
    } finally {
      reader.releaseLock();
    }

    yield {
      type: 'done',
      data: { outputs: lastOutputs, text: collectedText, usage: lastUsage },
    };
  }

  private endpoint(): { url: string; key: string } {
    const url = this.config.getOrThrow<string>('ai.gatewayUrl');
    const key = this.config.getOrThrow<string>('ai.gatewayApiKey');
    return { url: url.replace(/\/+$/, ''), key };
  }

  private headers(key: string): Record<string, string> {
    return {
      Authorization: `Bearer ${key}`,
      'Content-Type': 'application/json',
      Accept: 'application/json',
    };
  }

  private normalize(json: Record<string, unknown>): GatewayRunResult {
    const data = (json['data'] as Record<string, unknown> | undefined) ?? json;
    const outputs = (data['outputs'] as Record<string, unknown> | undefined) ?? {};
    const text =
      (outputs['text'] as string | undefined) ??
      (data['answer'] as string | undefined) ??
      '';
    const usage = (data['metadata'] as Record<string, unknown> | undefined)?.['usage'] as
      | Record<string, unknown>
      | undefined;
    return {
      outputs,
      text,
      metadata: {
        promptTokens: Number(usage?.['prompt_tokens'] ?? 0),
        completionTokens: Number(usage?.['completion_tokens'] ?? 0),
        totalTokens: Number(usage?.['total_tokens'] ?? 0),
        costCents: Number(usage?.['cost_cents'] ?? 0),
        raw: usage,
      },
    };
  }
}
