import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
import type { CallToolResult } from '@modelcontextprotocol/sdk/types.js'
import { beforeEach, describe, expect, it, vi } from 'vitest'
import * as z from 'zod'
import { type Backend, BackendDisposedError } from '../../../src/backend.js'
import type { ResolvedConnection } from '../../../src/config/baseSchema.js'
import type { Registry } from '../../../src/config/registry.js'
import { RegistryError } from '../../../src/config/registry.js'
import type { ManagedConnection, Manager } from '../../../src/manager.js'
import {
    type ConnectionToolsDeps,
    registerConnectionTools
} from '../../../src/mcp/connectionTools.js'

// Tiny inline config + schema: name/type plus a readonly default so the schema's
// input type (readonly optional) differs from its output, exercising the
// z.ZodType<TConfig, ZodTypeDef, unknown> input-variance contract.
const schema = z
    .object({
        name: z.string(),
        type: z.enum(['postgres', 'mysql']),
        host: z.string(),
        readonly: z.boolean().default(false)
    })
    .strict()

type TestConfig = z.infer<typeof schema>
type TestBackend = Backend & { ping: () => Promise<string> }
type TestManaged = ManagedConnection<TestConfig, TestBackend>

const resolved = (config: TestConfig): ResolvedConnection<TestConfig> => ({
    config,
    source: 'store',
    hash: 'h'
})

const publicView = (r: ResolvedConnection<TestConfig>): unknown => ({
    name: r.config.name,
    type: r.config.type,
    host: r.config.host,
    readonly: r.config.readonly,
    source: r.source
})

type FakeRegistry = Registry<TestConfig> & {
    list: ReturnType<typeof vi.fn>
    resolve: ReturnType<typeof vi.fn>
    add: ReturnType<typeof vi.fn>
    update: ReturnType<typeof vi.fn>
    remove: ReturnType<typeof vi.fn>
}

const fakeRegistry = (): FakeRegistry => ({
    list: vi.fn(() => []),
    resolve: vi.fn(),
    add: vi.fn(),
    update: vi.fn(),
    remove: vi.fn()
})

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 managed = (ping: () => Promise<string>): TestManaged => ({
    backend: { dispose: vi.fn(async () => {}), ping },
    config: { name: 'c', type: 'postgres', host: 'h', readonly: false },
    source: 'store'
})

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

const descriptions = {
    list: 'list desc',
    add: 'add desc',
    update: 'update desc',
    remove: 'remove desc',
    test: 'test desc'
}

type Handler = (args: Record<string, unknown>) => Promise<CallToolResult>

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

describe('registerConnectionTools', () => {
    let registry: FakeRegistry
    let manager: FakeManager
    let handlers: Map<string, Handler>
    let configs: Map<string, { description?: string }>

    const invoke = (name: string, args: Record<string, unknown> = {}): Promise<CallToolResult> => {
        const handler = handlers.get(name)
        if (!handler) {
            throw new Error('tool not registered: ' + name)
        }
        return handler(args)
    }

    const register = (overrides: Partial<ConnectionToolsDeps<TestConfig, TestBackend>> = {}) => {
        const server = new McpServer({ name: 'test', version: '0.0.0' })
        handlers = new Map()
        configs = new Map()
        vi.spyOn(server, 'registerTool').mockImplementation(
            (name: string, config: { description?: string }, handler: unknown) => {
                configs.set(name, config)
                handlers.set(name, handler as Handler)
                return {
                    enable: vi.fn(),
                    disable: vi.fn(),
                    update: vi.fn(),
                    remove: vi.fn()
                } as never
            }
        )
        registerConnectionTools(server, {
            registry,
            manager,
            fullSchema: schema,
            patchSchema: z.object({ host: z.string().optional() }).strict(),
            publicView,
            ping: async (m) => ({ version: await m.backend.ping() }),
            formatError,
            descriptions,
            ...overrides
        })
    }

    beforeEach(() => {
        registry = fakeRegistry()
        manager = fakeManager()
    })

    it('registers all five connection tools with their descriptions', () => {
        register()
        expect([...handlers.keys()].sort()).toEqual([
            'add_connection',
            'list_connections',
            'remove_connection',
            'test_connection',
            'update_connection'
        ])
        expect(configs.get('list_connections')?.description).toBe('list desc')
        expect(configs.get('add_connection')?.description).toBe('add desc')
        expect(configs.get('update_connection')?.description).toBe('update desc')
        expect(configs.get('remove_connection')?.description).toBe('remove desc')
        expect(configs.get('test_connection')?.description).toBe('test desc')
    })

    it('list_connections maps each entry through publicView', async () => {
        registry.list.mockReturnValue([
            resolved({ name: 'a', type: 'postgres', host: 'ha', readonly: false }),
            resolved({ name: 'b', type: 'mysql', host: 'hb', readonly: true })
        ])
        register()

        const result = await invoke('list_connections')

        expect(result.isError).toBeFalsy()
        expect(JSON.parse(text(result))).toEqual([
            { name: 'a', type: 'postgres', host: 'ha', readonly: false, source: 'store' },
            { name: 'b', type: 'mysql', host: 'hb', readonly: true, source: 'store' }
        ])
    })

    it('add_connection parses, persists and reports the public view', async () => {
        registry.add.mockImplementation((config: TestConfig) => resolved(config))
        register()

        const result = await invoke('add_connection', {
            name: 'fresh',
            type: 'mysql',
            host: 'db'
        })

        expect(result.isError).toBeFalsy()
        expect(registry.add).toHaveBeenCalledWith({
            name: 'fresh',
            type: 'mysql',
            host: 'db',
            readonly: false
        })
        expect(JSON.parse(text(result))).toEqual({
            added: { name: 'fresh', type: 'mysql', host: 'db', readonly: false, source: 'store' }
        })
    })

    it('add_connection reports a duplicate error from the registry', async () => {
        registry.add.mockImplementation(() => {
            throw new RegistryError("connection 'fresh' already exists (source: store)")
        })
        register()

        const result = await invoke('add_connection', {
            name: 'fresh',
            type: 'mysql',
            host: 'db'
        })

        expect(result.isError).toBe(true)
        expect(text(result)).toContain('already exists')
    })

    it('update_connection patches, invalidates the cache and returns the public view', async () => {
        registry.update.mockReturnValue(
            resolved({ name: 'c', type: 'postgres', host: 'new', readonly: false })
        )
        register()

        const result = await invoke('update_connection', {
            name: 'c',
            patch: { host: 'new' }
        })

        expect(result.isError).toBeFalsy()
        expect(registry.update).toHaveBeenCalledWith('c', { host: 'new' })
        expect(manager.invalidate).toHaveBeenCalledWith('c')
        expect(JSON.parse(text(result)).updated.host).toBe('new')
    })

    it('update_connection surfaces registry errors and does not invalidate', async () => {
        registry.update.mockImplementation(() => {
            throw new RegistryError('read-only at runtime')
        })
        register()

        const result = await invoke('update_connection', { name: 'c', patch: {} })

        expect(result.isError).toBe(true)
        expect(text(result)).toContain('read-only')
        expect(manager.invalidate).not.toHaveBeenCalled()
    })

    it('remove_connection deletes and invalidates the cache', async () => {
        register()

        const result = await invoke('remove_connection', { name: 'c' })

        expect(result.isError).toBeFalsy()
        expect(registry.remove).toHaveBeenCalledWith('c')
        expect(manager.invalidate).toHaveBeenCalledWith('c')
        expect(JSON.parse(text(result))).toEqual({ removed: 'c' })
    })

    it('remove_connection reports registry errors', async () => {
        registry.remove.mockImplementation(() => {
            throw new RegistryError('not found')
        })
        register()

        const result = await invoke('remove_connection', { name: 'ghost' })

        expect(result.isError).toBe(true)
        expect(text(result)).toContain('not found')
    })

    it('test_connection reports ping info plus latency', async () => {
        manager.get.mockResolvedValue(managed(vi.fn(async () => '17.2')))
        register()

        const result = await invoke('test_connection', { name: 'c' })

        expect(result.isError).toBeFalsy()
        const payload = JSON.parse(text(result)) as Record<string, unknown>
        expect(payload.ok).toBe(true)
        expect(payload.version).toBe('17.2')
        expect(typeof payload.latencyMs).toBe('number')
    })

    it('test_connection retries once when the backend was disposed under it', async () => {
        const ping = vi
            .fn<() => Promise<string>>()
            .mockRejectedValueOnce(new BackendDisposedError())
            .mockResolvedValueOnce('17.2')
        manager.get.mockResolvedValue(managed(ping))
        register()

        const result = await invoke('test_connection', { name: 'c' })

        expect(result.isError).toBeFalsy()
        expect((JSON.parse(text(result)) as Record<string, unknown>).version).toBe('17.2')
        expect(manager.get).toHaveBeenCalledTimes(2)
    })

    it('test_connection formats backend failures via formatError', async () => {
        manager.get.mockResolvedValue(
            managed(
                vi.fn(async () => {
                    throw new Error('connect ECONNREFUSED')
                })
            )
        )
        register()

        const result = await invoke('test_connection', { name: 'c' })

        expect(result.isError).toBe(true)
        expect(text(result)).toBe('[postgres] connect ECONNREFUSED')
    })

    it('test_connection formats the error after a second disposal (no further retry)', async () => {
        // First disposal triggers one rebuild; the second is on attempt 1, where the
        // retry guard (attempt === 0) no longer holds, so formatError is reported.
        manager.get.mockResolvedValue(
            managed(
                vi.fn(async () => {
                    throw new BackendDisposedError()
                })
            )
        )
        register()

        const result = await invoke('test_connection', { name: 'c' })

        expect(result.isError).toBe(true)
        expect(text(result)).toBe(
            '[postgres] connection was reconfigured or removed: retry the call'
        )
        expect(manager.get).toHaveBeenCalledTimes(2)
    })

    it('test_connection reports manager.get failures plainly', async () => {
        manager.get.mockRejectedValue(new Error('ssh auth failed'))
        register()

        const result = await invoke('test_connection', { name: 'c' })

        expect(result.isError).toBe(true)
        expect(text(result)).toBe('ssh auth failed')
    })
})