import type { CallToolResult } from '@modelcontextprotocol/sdk/types.js'
import { describe, expect, it, vi } from 'vitest'
import { type Backend, BackendDisposedError } from '../../../src/backend.js'
import type { ManagedConnection, Manager } from '../../../src/manager.js'
import { withManaged } from '../../../src/mcp/managed.js'

type TestConfig = { name: string; type: string }
type TestBackend = Backend & { ping: () => Promise<string> }
type TestManaged = ManagedConnection<TestConfig, TestBackend>

const config: TestConfig = { name: 'c', type: 'postgres' }

const managed = (backend: Partial<TestBackend> = {}): TestManaged => ({
    backend: { dispose: vi.fn(async () => {}), ping: vi.fn(async () => 'ok'), ...backend },
    config,
    source: 'store'
})

type FakeManager = Manager<TestConfig, TestBackend> & {
    get: ReturnType<typeof vi.fn>
    invalidate: ReturnType<typeof vi.fn>
    disposeAll: ReturnType<typeof vi.fn>
}

const fakeManager = (): FakeManager => ({
    get: vi.fn(),
    invalidate: vi.fn(async () => {}),
    disposeAll: vi.fn(async () => {})
})

const formatError = (cfg: TestConfig, error: unknown): string =>
    '[' + cfg.type + '] ' + (error instanceof Error ? error.message : String(error))

const text = (result: CallToolResult): string => {
    const first = result.content[0]
    return first && first.type === 'text' ? first.text : ''
}

describe('withManaged', () => {
    it('returns ok with the fn result on success', async () => {
        const manager = fakeManager()
        manager.get.mockResolvedValue(managed())
        const fn = vi.fn(async () => ({ value: 42 }))

        const result = await withManaged(manager, 'c', fn, { formatError })

        expect(result.isError).toBeFalsy()
        expect(JSON.parse(text(result))).toEqual({ value: 42 })
        expect(manager.get).toHaveBeenCalledTimes(1)
    })

    it('fails when manager.get throws, without invoking fn', async () => {
        const manager = fakeManager()
        manager.get.mockRejectedValue(new Error('no such connection'))
        const fn = vi.fn()

        const result = await withManaged(manager, 'c', fn, { formatError })

        expect(result.isError).toBe(true)
        expect(text(result)).toBe('no such connection')
        expect(fn).not.toHaveBeenCalled()
    })

    it('retries once when the backend was disposed under the caller', async () => {
        const manager = fakeManager()
        manager.get.mockResolvedValue(managed())
        const fn = vi
            .fn()
            .mockRejectedValueOnce(new BackendDisposedError())
            .mockResolvedValueOnce({ value: 7 })

        const result = await withManaged(manager, 'c', fn, { formatError })

        expect(result.isError).toBeFalsy()
        expect(JSON.parse(text(result))).toEqual({ value: 7 })
        expect(manager.get).toHaveBeenCalledTimes(2)
        expect(fn).toHaveBeenCalledTimes(2)
    })

    it('gives up with a reconfiguring message after a second stale error', async () => {
        const manager = fakeManager()
        manager.get.mockResolvedValue(managed())
        const fn = vi.fn().mockRejectedValue(new BackendDisposedError())

        const result = await withManaged(manager, 'c', fn, { formatError })

        expect(result.isError).toBe(true)
        expect(text(result)).toContain('being reconfigured')
        expect(manager.get).toHaveBeenCalledTimes(2)
        expect(fn).toHaveBeenCalledTimes(2)
    })

    it('formats non-stale errors via formatError, no retry', async () => {
        const manager = fakeManager()
        manager.get.mockResolvedValue(managed())
        const fn = vi.fn().mockRejectedValue(new Error('boom'))

        const result = await withManaged(manager, 'c', fn, { formatError })

        expect(result.isError).toBe(true)
        expect(text(result)).toBe('[postgres] boom')
        expect(manager.get).toHaveBeenCalledTimes(1)
        expect(fn).toHaveBeenCalledTimes(1)
    })

    it('honors a custom isStaleError predicate', async () => {
        const manager = fakeManager()
        manager.get.mockResolvedValue(managed())
        const fn = vi
            .fn()
            .mockRejectedValueOnce(Object.assign(new Error('stale'), { code: 'STALE' }))
            .mockResolvedValueOnce({ value: 1 })

        const result = await withManaged(manager, 'c', fn, {
            formatError,
            isStaleError: (error) =>
                error instanceof Error && (error as { code?: string }).code === 'STALE'
        })

        expect(result.isError).toBeFalsy()
        expect(JSON.parse(text(result))).toEqual({ value: 1 })
        expect(fn).toHaveBeenCalledTimes(2)
    })
})