import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'
import { SpanlensClient } from '../client.js'
import { observeOpenAI, observeAnthropic, observeGemini, observeOllama } from '../observe.js'

describe('observeOpenAI / observeAnthropic / observeGemini', () => {
  let fetchMock: ReturnType<typeof vi.fn>

  beforeEach(() => {
    fetchMock = vi.fn().mockResolvedValue(
      new Response(JSON.stringify({ success: true }), { status: 200 }),
    )
    vi.stubGlobal('fetch', fetchMock)
  })

  afterEach(() => {
    vi.unstubAllGlobals()
  })

  it('observeOpenAI injects tracing headers into callback', async () => {
    const client = new SpanlensClient({ apiKey: 'k', baseUrl: 'http://x' })
    const trace = client.startTrace({ name: 't' })

    let receivedHeaders: Record<string, string> | null = null
    await observeOpenAI(trace, 'call-gpt4', async (headers) => {
      receivedHeaders = headers
      return {
        model: 'gpt-4o',
        usage: { prompt_tokens: 50, completion_tokens: 100, total_tokens: 150 },
      }
    })

    expect(receivedHeaders).not.toBeNull()
    expect(receivedHeaders!['x-trace-id']).toBe(trace.traceId)
    expect(receivedHeaders!['x-span-id']).toMatch(/^[0-9a-f-]{36}$/)
  })

  it('observeOpenAI forwards promptVersion option as x-spanlens-prompt-version header', async () => {
    const client = new SpanlensClient({ apiKey: 'k', baseUrl: 'http://x' })
    const trace = client.startTrace({ name: 't' })

    let receivedHeaders: Record<string, string> | null = null
    await observeOpenAI(
      trace,
      { name: 'call-v3', promptVersion: 'chatbot-system@3' },
      async (headers) => {
        receivedHeaders = headers
        return {
          model: 'gpt-4o-mini',
          usage: { prompt_tokens: 10, completion_tokens: 20, total_tokens: 30 },
        }
      },
    )

    expect(receivedHeaders!['x-spanlens-prompt-version']).toBe('chatbot-system@3')
    // traceparent-style headers should still be present
    expect(receivedHeaders!['x-trace-id']).toBe(trace.traceId)
  })

  it('observeAnthropic forwards promptVersion option', async () => {
    const client = new SpanlensClient({ apiKey: 'k', baseUrl: 'http://x' })
    const trace = client.startTrace({ name: 't' })

    let receivedHeaders: Record<string, string> | null = null
    await observeAnthropic(
      trace,
      { name: 'call', promptVersion: 'greeter@latest' },
      async (headers) => {
        receivedHeaders = headers
        return {
          model: 'claude-3-5-sonnet-20241022',
          usage: { input_tokens: 5, output_tokens: 7 },
        }
      },
    )

    expect(receivedHeaders!['x-spanlens-prompt-version']).toBe('greeter@latest')
  })

  it('observeGemini forwards promptVersion option', async () => {
    const client = new SpanlensClient({ apiKey: 'k', baseUrl: 'http://x' })
    const trace = client.startTrace({ name: 't' })

    let receivedHeaders: Record<string, string> | null = null
    await observeGemini(
      trace,
      { name: 'call', promptVersion: 'uuid-like-id-12345' },
      async (headers) => {
        receivedHeaders = headers
        return {
          response: {
            usageMetadata: { promptTokenCount: 3, candidatesTokenCount: 4, totalTokenCount: 7 },
          },
        }
      },
    )

    expect(receivedHeaders!['x-spanlens-prompt-version']).toBe('uuid-like-id-12345')
  })

  it('omits prompt version header when option is not provided', async () => {
    const client = new SpanlensClient({ apiKey: 'k', baseUrl: 'http://x' })
    const trace = client.startTrace({ name: 't' })

    let receivedHeaders: Record<string, string> | null = null
    await observeOpenAI(trace, 'call', async (headers) => {
      receivedHeaders = headers
      return {
        model: 'gpt-4o',
        usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 },
      }
    })

    expect(receivedHeaders!['x-spanlens-prompt-version']).toBeUndefined()
  })

  it('forwards logBody option as x-spanlens-log-body header', async () => {
    const client = new SpanlensClient({ apiKey: 'k', baseUrl: 'http://x' })
    const trace = client.startTrace({ name: 't' })

    let receivedHeaders: Record<string, string> | null = null
    await observeOpenAI(
      trace,
      { name: 'call', logBody: 'meta' },
      async (headers) => {
        receivedHeaders = headers
        return {
          model: 'gpt-4o',
          usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 },
        }
      },
    )

    expect(receivedHeaders!['x-spanlens-log-body']).toBe('meta')
  })

  it('omits logBody header when option is not provided (server defaults to full)', async () => {
    const client = new SpanlensClient({ apiKey: 'k', baseUrl: 'http://x' })
    const trace = client.startTrace({ name: 't' })

    let receivedHeaders: Record<string, string> | null = null
    await observeOpenAI(trace, 'call', async (headers) => {
      receivedHeaders = headers
      return {
        model: 'gpt-4o',
        usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 },
      }
    })

    expect(receivedHeaders!['x-spanlens-log-body']).toBeUndefined()
  })

  it('observeOpenAI auto-parses usage into span.end', async () => {
    const client = new SpanlensClient({ apiKey: 'k', baseUrl: 'http://x' })
    const trace = client.startTrace({ name: 't' })

    await observeOpenAI(trace, 'call', async () => ({
      model: 'gpt-4o-mini',
      usage: { prompt_tokens: 10, completion_tokens: 20, total_tokens: 30 },
    }))

    const patchCall = fetchMock.mock.calls.find(
      ([, init]) => (init as RequestInit).method === 'PATCH',
    )
    expect(patchCall).toBeDefined()
    const body = JSON.parse((patchCall![1] as RequestInit).body as string) as Record<string, unknown>
    expect(body.total_tokens).toBe(30)
    expect(body.prompt_tokens).toBe(10)
    expect(body.completion_tokens).toBe(20)
    expect(body.status).toBe('completed')
    expect((body.metadata as Record<string, unknown>).model).toBe('gpt-4o-mini')
  })

  it('observeOpenAI creates span with spanType=llm', async () => {
    const client = new SpanlensClient({ apiKey: 'k', baseUrl: 'http://x' })
    const trace = client.startTrace({ name: 't' })

    await observeOpenAI(trace, 'call', async () => ({ usage: { total_tokens: 0 } }))
    await new Promise((r) => setTimeout(r, 10))

    const postCall = fetchMock.mock.calls.find(
      ([url, init]) =>
        typeof url === 'string' &&
        url.includes('/spans') &&
        (init as RequestInit).method === 'POST',
    )
    expect(postCall).toBeDefined()
    const body = JSON.parse((postCall![1] as RequestInit).body as string) as Record<string, unknown>
    expect(body.span_type).toBe('llm')
    expect(body.name).toBe('call')
  })

  it('observeAnthropic parses input_tokens/output_tokens', async () => {
    const client = new SpanlensClient({ apiKey: 'k', baseUrl: 'http://x' })
    const trace = client.startTrace({ name: 't' })

    await observeAnthropic(trace, 'msg', async () => ({
      model: 'claude-haiku-4-5',
      usage: { input_tokens: 15, output_tokens: 45 },
    }))

    const patchCall = fetchMock.mock.calls.find(
      ([, init]) => (init as RequestInit).method === 'PATCH',
    )
    const body = JSON.parse((patchCall![1] as RequestInit).body as string) as Record<string, unknown>
    expect(body.prompt_tokens).toBe(15)
    expect(body.completion_tokens).toBe(45)
    expect(body.total_tokens).toBe(60)
  })

  it('observeGemini parses usageMetadata from real GenerateContentResult shape', async () => {
    const client = new SpanlensClient({ apiKey: 'k', baseUrl: 'http://x' })
    const trace = client.startTrace({ name: 't' })

    // Real Gemini SDK: generateContent() returns { response: GenerateContentResponse }
    await observeGemini(trace, 'gen', async () => ({
      response: {
        modelVersion: 'gemini-2.0-flash',
        candidates: [{ content: { parts: [{ text: 'hello' }], role: 'model' } }],
        usageMetadata: {
          promptTokenCount: 5,
          candidatesTokenCount: 15,
          totalTokenCount: 20,
        },
      },
    }))

    const patchCall = fetchMock.mock.calls.find(
      ([, init]) => (init as RequestInit).method === 'PATCH',
    )
    const body = JSON.parse((patchCall![1] as RequestInit).body as string) as Record<string, unknown>
    expect(body.total_tokens).toBe(20)
    expect(body.prompt_tokens).toBe(5)
    expect(body.completion_tokens).toBe(15)
    expect((body.metadata as Record<string, unknown>)?.model).toBe('gemini-2.0-flash')
  })

  it('observeOpenAI captures full response as output in span.end', async () => {
    const client = new SpanlensClient({ apiKey: 'k', baseUrl: 'http://x' })
    const trace = client.startTrace({ name: 't' })

    const mockResponse = {
      id: 'chatcmpl-abc',
      model: 'gpt-4o-mini',
      choices: [{ message: { role: 'assistant', content: 'Hello!' }, finish_reason: 'stop' }],
      usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 },
    }

    await observeOpenAI(trace, 'call', async () => mockResponse)

    const patchCall = fetchMock.mock.calls.find(
      ([, init]) => (init as RequestInit).method === 'PATCH',
    )
    expect(patchCall).toBeDefined()
    const body = JSON.parse((patchCall![1] as RequestInit).body as string) as Record<string, unknown>
    expect(body.output).toEqual(mockResponse)
  })

  it('observeOpenAI omits output for stream-like responses', async () => {
    const client = new SpanlensClient({ apiKey: 'k', baseUrl: 'http://x' })
    const trace = client.startTrace({ name: 't' })

    const streamLike = { usage: { total_tokens: 5 }, [Symbol.asyncIterator]: () => ({}) }

    await observeOpenAI(trace, 'stream-call', async () => streamLike as unknown as typeof streamLike)

    const patchCall = fetchMock.mock.calls.find(
      ([, init]) => (init as RequestInit).method === 'PATCH',
    )
    const body = JSON.parse((patchCall![1] as RequestInit).body as string) as Record<string, unknown>
    expect(body.output).toBeUndefined()
  })

  it('marks span as error and rethrows when callback fails', async () => {
    const client = new SpanlensClient({ apiKey: 'k', baseUrl: 'http://x' })
    const trace = client.startTrace({ name: 't' })

    await expect(
      observeOpenAI(trace, 'boom', async () => {
        throw new Error('api failed')
      }),
    ).rejects.toThrow('api failed')

    const errPatch = fetchMock.mock.calls.find(([, init]) => {
      const ri = init as RequestInit | undefined
      if (!ri || ri.method !== 'PATCH') return false
      const body = JSON.parse(ri.body as string) as Record<string, unknown>
      return body.status === 'error' && body.error_message === 'api failed'
    })
    expect(errPatch).toBeDefined()
  })

  // ── Provider tag (Ollama + override) ───────────────────────────────────────

  it('observeOpenAI stamps metadata.provider="openai" by default', async () => {
    const client = new SpanlensClient({ apiKey: 'k', baseUrl: 'http://x' })
    const trace = client.startTrace({ name: 't' })

    await observeOpenAI(trace, 'call', async () => ({
      model: 'gpt-4o-mini',
      usage: { prompt_tokens: 1, completion_tokens: 2, total_tokens: 3 },
    }))

    const patchCall = fetchMock.mock.calls.find(
      ([, init]) => (init as RequestInit).method === 'PATCH',
    )
    const body = JSON.parse((patchCall![1] as RequestInit).body as string) as Record<string, unknown>
    expect((body.metadata as Record<string, unknown>).provider).toBe('openai')
    // Model still flows through alongside the new provider tag.
    expect((body.metadata as Record<string, unknown>).model).toBe('gpt-4o-mini')
  })

  it('observeOllama parses OpenAI-compatible response and tags provider="ollama"', async () => {
    const client = new SpanlensClient({ apiKey: 'k', baseUrl: 'http://x' })
    const trace = client.startTrace({ name: 't' })

    // Ollama's /v1 endpoint returns an OpenAI-shaped payload.
    await observeOllama(trace, 'chat', async () => ({
      model: 'llama3.2',
      usage: { prompt_tokens: 12, completion_tokens: 34, total_tokens: 46 },
    }))

    const patchCall = fetchMock.mock.calls.find(
      ([, init]) => (init as RequestInit).method === 'PATCH',
    )
    expect(patchCall).toBeDefined()
    const body = JSON.parse((patchCall![1] as RequestInit).body as string) as Record<string, unknown>
    expect(body.prompt_tokens).toBe(12)
    expect(body.completion_tokens).toBe(34)
    expect(body.total_tokens).toBe(46)
    expect((body.metadata as Record<string, unknown>).provider).toBe('ollama')
    expect((body.metadata as Record<string, unknown>).model).toBe('llama3.2')
  })

  it('observeOpenAI provider override wins over default tag', async () => {
    const client = new SpanlensClient({ apiKey: 'k', baseUrl: 'http://x' })
    const trace = client.startTrace({ name: 't' })

    // User points OpenAI SDK at vLLM (OpenAI-compatible) and wants the
    // dashboard to label it as 'vllm' not 'openai'.
    await observeOpenAI(
      trace,
      { name: 'vllm-call', provider: 'vllm' },
      async () => ({
        model: 'meta-llama/Llama-3-8B',
        usage: { prompt_tokens: 1, completion_tokens: 2, total_tokens: 3 },
      }),
    )

    const patchCall = fetchMock.mock.calls.find(
      ([, init]) => (init as RequestInit).method === 'PATCH',
    )
    const body = JSON.parse((patchCall![1] as RequestInit).body as string) as Record<string, unknown>
    expect((body.metadata as Record<string, unknown>).provider).toBe('vllm')
  })

  it('works when parent is a SpanHandle (nested LLM call)', async () => {
    const client = new SpanlensClient({ apiKey: 'k', baseUrl: 'http://x' })
    const trace = client.startTrace({ name: 't' })
    const outer = trace.span({ name: 'agent_loop' })

    let receivedHeaders: Record<string, string> | null = null
    await observeOpenAI(outer, 'inner-call', async (headers) => {
      receivedHeaders = headers
      return { usage: { total_tokens: 0 } }
    })

    expect(receivedHeaders).not.toBeNull()
    // The inner span should be a child of `outer`
    await new Promise((r) => setTimeout(r, 10))
    const spanPosts = fetchMock.mock.calls.filter(
      ([url, init]) =>
        typeof url === 'string' &&
        url.includes('/spans') &&
        (init as RequestInit).method === 'POST',
    )
    // Outer span post + inner span post = 2
    expect(spanPosts.length).toBeGreaterThanOrEqual(2)
    const innerPost = spanPosts.find((call) => {
      const body = JSON.parse((call[1] as RequestInit).body as string) as Record<string, unknown>
      return body.name === 'inner-call'
    })
    expect(innerPost).toBeDefined()
    const innerBody = JSON.parse((innerPost![1] as RequestInit).body as string) as Record<string, unknown>
    expect(innerBody.parent_span_id).toBe(outer.spanId)
  })
})
