import { describe, expect, it } from 'vitest'
import { Agent } from '../agent.js'
import { AgentAsTool } from '../agent-as-tool.js'
import { MockMessageModel } from '../../__fixtures__/mock-message-model.js'
import { collectGenerator } from '../../__fixtures__/model-test-helpers.js'
import { createMockContext } from '../../__fixtures__/tool-helpers.js'
import { ToolValidationError } from '../../errors.js'
import { Tool, ToolStreamEvent } from '../../tools/tool.js'
import { ToolResultBlock } from '../../types/messages.js'
import { SessionManager } from '../../session/session-manager.js'
import type { SnapshotStorage } from '../../session/storage.js'

describe('AgentAsTool', () => {
  describe('properties', () => {
    it('uses agent name as default tool name', () => {
      const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hi' })
      const agent = new Agent({ model, name: 'researcher' })
      const tool = new AgentAsTool({ agent })

      expect(tool.name).toBe('researcher')
    })

    it('allows overriding the tool name', () => {
      const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hi' })
      const agent = new Agent({ model, name: 'researcher' })
      const tool = new AgentAsTool({ agent, name: 'research-tool' })

      expect(tool.name).toBe('research-tool')
    })

    it('uses agent description as default tool description', () => {
      const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hi' })
      const agent = new Agent({ model, name: 'researcher', description: 'Finds information' })
      const tool = new AgentAsTool({ agent })

      expect(tool.description).toBe('Finds information')
    })

    it('falls back to generic description when agent has no description', () => {
      const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hi' })
      const agent = new Agent({ model, name: 'researcher' })
      const tool = new AgentAsTool({ agent })

      expect(tool.description).toBe('Use the researcher agent by providing a natural language input')
    })

    it('allows overriding the tool description', () => {
      const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hi' })
      const agent = new Agent({ model, name: 'researcher' })
      const tool = new AgentAsTool({ agent, description: 'Custom description' })

      expect(tool.description).toBe('Custom description')
    })

    it('exposes the wrapped agent via getter', () => {
      const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hi' })
      const agent = new Agent({ model, name: 'researcher' })
      const tool = new AgentAsTool({ agent })

      expect(tool.agent).toBe(agent)
    })

    it('has correct toolSpec shape', () => {
      const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hi' })
      const agent = new Agent({ model, name: 'researcher', description: 'Finds info' })
      const tool = new AgentAsTool({ agent })

      expect(tool.toolSpec).toEqual({
        name: 'researcher',
        description: 'Finds info',
        inputSchema: {
          type: 'object',
          properties: {
            input: {
              type: 'string',
              description: 'The natural language input to send to the agent.',
            },
          },
          required: ['input'],
        },
      })
    })
  })

  describe('name validation', () => {
    it('throws when registered with agent name containing spaces', () => {
      const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hi' })
      const subAgent = new Agent({ model, name: 'Strands Agent' })

      expect(() => new Agent({ model, tools: [subAgent] })).toThrow(ToolValidationError)
    })

    it('throws when registered with explicit name containing invalid characters', () => {
      const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hi' })
      const subAgent = new Agent({ model, name: 'researcher' })

      expect(() => new Agent({ model, tools: [subAgent.asTool({ name: 'has spaces' })] })).toThrow(ToolValidationError)
    })

    it('accepts valid name with hyphens and underscores', () => {
      const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hi' })
      const agent = new Agent({ model, name: 'my_research-agent' })
      const tool = new AgentAsTool({ agent })

      expect(tool.name).toBe('my_research-agent')
    })
  })

  describe('stream', () => {
    it('invokes the wrapped agent and returns text result', async () => {
      const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Agent response' })
      const agent = new Agent({ model, name: 'test-agent', printer: false })
      const tool = new AgentAsTool({ agent })

      const context = createMockContext({
        name: 'test-agent',
        toolUseId: 'tool-1',
        input: { input: 'Hello agent' },
      })

      const { result } = await collectGenerator(tool.stream(context))

      expect(result.toolUseId).toBe('tool-1')
      expect(result.status).toBe('success')
      expect(result.content).toHaveLength(1)
      expect(result.content[0]).toEqual(
        expect.objectContaining({
          type: 'textBlock',
          text: 'Agent response',
        })
      )
    })

    it('yields ToolStreamEvents wrapping sub-agent events', async () => {
      const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hello' })
      const agent = new Agent({ model, name: 'test-agent', printer: false })
      const tool = new AgentAsTool({ agent })

      const context = createMockContext({
        name: 'test-agent',
        toolUseId: 'tool-1',
        input: { input: 'Hi' },
      })

      const { items } = await collectGenerator(tool.stream(context))

      expect(items.length).toBeGreaterThan(0)
      for (const item of items) {
        expect(item).toBeInstanceOf(ToolStreamEvent)
      }
    })

    it('unwraps toolStreamUpdateEvent by yielding inner ToolStreamEvent directly', async () => {
      // Create a tool that yields ToolStreamEvents during execution.
      // When the sub-agent runs this tool, the agent loop wraps each yielded
      // ToolStreamEvent in a ToolStreamUpdateEvent. The AgentAsTool should
      // unwrap these back to bare ToolStreamEvents instead of double-wrapping.
      const streamingTool = {
        name: 'streaming-tool',
        description: 'A tool that yields stream events',
        toolSpec: {
          name: 'streaming-tool',
          description: 'A tool that yields stream events',
          inputSchema: { type: 'object' as const, properties: {} },
        },
        async *stream(context: any) {
          yield new ToolStreamEvent({ data: 'progress-1' })
          yield new ToolStreamEvent({ data: 'progress-2' })
          return new ToolResultBlock({
            toolUseId: context.toolUse.toolUseId,
            status: 'success' as const,
            content: [],
          })
        },
      } as Tool

      // Turn 1: model requests tool use → triggers the streaming tool
      // Turn 2: model responds with text after tool result
      const model = new MockMessageModel()
        .addTurn({
          type: 'toolUseBlock',
          name: 'streaming-tool',
          toolUseId: 'sub-tool-1',
          input: {},
        })
        .addTurn({ type: 'textBlock', text: 'Final response' })

      const agent = new Agent({ model, name: 'test-agent', tools: [streamingTool], printer: false })
      const tool = new AgentAsTool({ agent })

      const context = createMockContext({
        name: 'test-agent',
        toolUseId: 'outer-tool-1',
        input: { input: 'Do something' },
      })

      const { items } = await collectGenerator(tool.stream(context))

      // All yielded items should be ToolStreamEvent instances
      for (const item of items) {
        expect(item).toBeInstanceOf(ToolStreamEvent)
      }

      // Find the unwrapped events from the streaming tool.
      // If unwrapping works correctly, data is the original string.
      // If double-wrapped, data would be a ToolStreamUpdateEvent object.
      const progressEvents = items.filter((item) => item.data === 'progress-1' || item.data === 'progress-2')

      expect(progressEvents).toHaveLength(2)
      expect(progressEvents[0]!.data).toBe('progress-1')
      expect(progressEvents[1]!.data).toBe('progress-2')
    })

    it('returns error result on agent failure', async () => {
      const model = new MockMessageModel().addTurn(new Error('Model failed'))
      const agent = new Agent({ model, name: 'test-agent', printer: false })
      const tool = new AgentAsTool({ agent })

      const context = createMockContext({
        name: 'test-agent',
        toolUseId: 'tool-1',
        input: { input: 'Hello' },
      })

      const { result } = await collectGenerator(tool.stream(context))

      expect(result.toolUseId).toBe('tool-1')
      expect(result.status).toBe('error')
    })

    it('returns error result when agent is already busy', async () => {
      const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Slow response' })
      const agent = new Agent({ model, name: 'test-agent', printer: false })
      const tool = new AgentAsTool({ agent })

      const context1 = createMockContext({
        name: 'test-agent',
        toolUseId: 'tool-1',
        input: { input: 'First call' },
      })
      const context2 = createMockContext({
        name: 'test-agent',
        toolUseId: 'tool-2',
        input: { input: 'Second call' },
      })

      // Start first call but don't fully consume it
      const gen1 = tool.stream(context1)
      await gen1.next()

      // Second call should get an error immediately
      const { result } = await collectGenerator(tool.stream(context2))

      expect(result.toolUseId).toBe('tool-2')
      expect(result.status).toBe('error')
      expect(result.content[0]).toEqual(
        expect.objectContaining({
          type: 'textBlock',
          text: expect.stringContaining('already processing'),
        })
      )

      // Clean up first generator
      await collectGenerator(gen1)
    })
  })

  describe('preserveContext', () => {
    it('resets agent state between invocations when false (default)', async () => {
      const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Response' })
      const agent = new Agent({ model, name: 'test-agent', printer: false })
      const tool = new AgentAsTool({ agent })

      const context1 = createMockContext({
        name: 'test-agent',
        toolUseId: 'tool-1',
        input: { input: 'Hello' },
      })
      const context2 = createMockContext({
        name: 'test-agent',
        toolUseId: 'tool-2',
        input: { input: 'Hello again' },
      })

      await collectGenerator(tool.stream(context1))
      const messagesAfterFirst = agent.messages.length

      await collectGenerator(tool.stream(context2))
      const messagesAfterSecond = agent.messages.length

      // State is reset so both produce the same message count
      expect(messagesAfterSecond).toBe(messagesAfterFirst)
    })

    it('preserves agent state across invocations when true', async () => {
      const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Response' })
      const agent = new Agent({ model, name: 'test-agent', printer: false })
      const tool = new AgentAsTool({ agent, preserveContext: true })

      const context1 = createMockContext({
        name: 'test-agent',
        toolUseId: 'tool-1',
        input: { input: 'Hello' },
      })
      const context2 = createMockContext({
        name: 'test-agent',
        toolUseId: 'tool-2',
        input: { input: 'Hello again' },
      })

      await collectGenerator(tool.stream(context1))
      const messagesAfterFirst = agent.messages.length

      await collectGenerator(tool.stream(context2))
      const messagesAfterSecond = agent.messages.length

      // Messages should accumulate across invocations
      expect(messagesAfterSecond).toBeGreaterThan(messagesAfterFirst)
    })

    it('snapshots at construction time, not first invocation', async () => {
      const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Response' })
      const agent = new Agent({ model, name: 'test-agent', printer: false })
      const tool = new AgentAsTool({ agent })
      const messagesAtConstruction = agent.messages.length

      // Modify agent state after tool creation
      await agent.invoke('Direct invocation')
      expect(agent.messages.length).toBeGreaterThan(messagesAtConstruction)

      // First tool call restores to construction-time state, then runs
      const context1 = createMockContext({
        name: 'test-agent',
        toolUseId: 'tool-1',
        input: { input: 'Hello' },
      })
      await collectGenerator(tool.stream(context1))
      const messagesAfterFirstTool = agent.messages.length

      // Second tool call should produce the same count — both reset to construction baseline
      const context2 = createMockContext({
        name: 'test-agent',
        toolUseId: 'tool-2',
        input: { input: 'Hello again' },
      })
      await collectGenerator(tool.stream(context2))

      expect(agent.messages.length).toBe(messagesAfterFirstTool)
    })
  })

  describe('sessionManager validation', () => {
    const mockStorage: SnapshotStorage = {
      saveSnapshot: async () => {},
      loadSnapshot: async () => null,
      listSnapshotIds: async () => [],
      deleteSession: async () => {},
      loadManifest: async () => ({ schemaVersion: '1.0', updatedAt: '12:00:00' }),
      saveManifest: async () => {},
    }

    it('throws when preserveContext is false and agent has a sessionManager', () => {
      const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hi' })
      const sessionManager = new SessionManager({ storage: { snapshot: mockStorage } })
      const agent = new Agent({ model, name: 'test-agent', sessionManager })

      expect(() => new AgentAsTool({ agent })).toThrow(/SessionManager.*conflicts with preserveContext=false/)
    })

    it('throws when preserveContext is explicitly false and agent has a sessionManager', () => {
      const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hi' })
      const sessionManager = new SessionManager({ storage: { snapshot: mockStorage } })
      const agent = new Agent({ model, name: 'test-agent', sessionManager })

      expect(() => new AgentAsTool({ agent, preserveContext: false })).toThrow(
        /SessionManager.*conflicts with preserveContext=false/
      )
    })

    it('allows preserveContext true with sessionManager', () => {
      const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hi' })
      const sessionManager = new SessionManager({ storage: { snapshot: mockStorage } })
      const agent = new Agent({ model, name: 'test-agent', sessionManager })

      expect(() => new AgentAsTool({ agent, preserveContext: true })).not.toThrow()
    })
  })

  describe('Agent.asTool', () => {
    it('returns an AgentAsTool instance', () => {
      const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hi' })
      const agent = new Agent({ model, name: 'researcher' })

      const tool = agent.asTool()

      expect(tool).toBeInstanceOf(AgentAsTool)
      expect(tool.name).toBe('researcher')
    })

    it('passes options through to AgentAsTool', () => {
      const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hi' })
      const agent = new Agent({ model, name: 'researcher' })

      const tool = agent.asTool({ name: 'custom-name', description: 'Custom desc' })

      expect(tool.name).toBe('custom-name')
      expect(tool.description).toBe('Custom desc')
    })
  })

  describe('Agent in ToolList', () => {
    it('auto-wraps Agent as AgentAsTool when passed in tools array', () => {
      const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hi' })
      const subAgent = new Agent({ model, name: 'sub-agent', description: 'A sub agent' })
      const parentAgent = new Agent({ model, tools: [subAgent] })

      const registeredTool = parentAgent.toolRegistry.get('sub-agent')
      expect(registeredTool).toBeInstanceOf(AgentAsTool)
      expect(registeredTool!.name).toBe('sub-agent')
    })

    it('auto-wraps Agent in nested tools arrays', () => {
      const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hi' })
      const subAgent = new Agent({ model, name: 'nested-agent' })
      const parentAgent = new Agent({ model, tools: [[subAgent]] })

      const registeredTool = parentAgent.toolRegistry.get('nested-agent')
      expect(registeredTool).toBeInstanceOf(AgentAsTool)
    })
  })
})
