import { describe, it, expect, vi } from 'vitest'
import { Agent } from '../agent'
import type { AgentDefinition, AgentConfig, LLMAdapterFactory } from '../types'
import type { Contract } from '../../../base/types'
import { z } from 'zod'
import { createMockExecutionContext, createMockTool } from '../../../test-utils/mocks'
import { MockAdapter } from '../../../llm/adapters/mock-adapter'

/**
 * Create a mock LLMAdapterFactory that returns a MockAdapter instance.
 * Used to satisfy the Agent constructor's second parameter.
 */
function createMockAdapterFactory(): LLMAdapterFactory {
  const adapter = new MockAdapter()
  return () => adapter
}

describe('Agent Execution', () => {

  // Helper: Create minimal agent definition with mock provider
  const createAgentDefinition = (overrides?: Partial<AgentDefinition>): AgentDefinition => {
    const config: AgentConfig = {
      type: 'agent',
      kind: 'utility',
      resourceId: 'test-agent',
      name: 'Test Agent',
      description: 'Agent for testing',
      environment: 'dev',
      systemPrompt: 'You are a test agent'
    }

    const contract: Contract = {
      inputSchema: z.object({ query: z.string() }),
      outputSchema: z.object({ response: z.string(), processedAt: z.string() })
    }

    return {
      config,
      contract,
      tools: [],
      modelConfig: {
        model: 'mock',
        provider: 'mock',
        apiKey: 'test-key'
      },
      ...overrides
    }
  }


  describe('Constructor', () => {
    it('creates agent instance from definition', () => {
      const definition = createAgentDefinition()
      const agent = new Agent(definition, createMockAdapterFactory())

      expect(agent.config).toEqual(definition.config)
      expect(agent.contract).toEqual(definition.contract)
      expect(agent.getTools()).toEqual(definition.tools)
      expect(agent.modelConfig).toEqual(definition.modelConfig)
    })

    it('stores tools from definition', () => {
      const tools = [createMockTool('toolA'), createMockTool('toolB')]
      const definition = createAgentDefinition({ tools })
      const agent = new Agent(definition, createMockAdapterFactory())

      const registeredTools = agent.getTools()
      expect(registeredTools).toHaveLength(2)
      expect(registeredTools[0].name).toBe('toolA')
      expect(registeredTools[1].name).toBe('toolB')
    })
  })

  describe('Phase 1: Initialization', () => {
    it('validates input against contract schema', async () => {
      const definition = createAgentDefinition()
      const agent = new Agent(definition, createMockAdapterFactory())
      const context = createMockExecutionContext()

      const input = { query: 'test query' }
      const output = await agent.execute(input, context)

      // If input validation failed, execution would throw
      expect(output).toBeDefined()
      expect(output).toHaveProperty('response')
    })

    it('throws error on invalid input', async () => {
      const definition = createAgentDefinition()
      const agent = new Agent(definition, createMockAdapterFactory())
      const context = createMockExecutionContext()

      const invalidInput = { wrong: 'field' }

      await expect(agent.execute(invalidInput as unknown, context)).rejects.toThrow()
    })

    it('creates memory manager with empty memory by default', async () => {
      const definition = createAgentDefinition()
      const agent = new Agent(definition, createMockAdapterFactory())
      const context = createMockExecutionContext()

      const input = { query: 'test' }
      await agent.execute(input, context)

      expect(agent.memoryManager).toBeDefined()
      const snapshot = agent.getMemorySnapshot()
      expect(snapshot).toBeDefined()
    })

    it('preloads memory if preloadMemory callback provided', async () => {
      const preloadMemory = vi.fn(async () => ({
        sessionMemory: {
          context: {
            type: 'context' as const,
            content: 'Preloaded context',
            timestamp: Date.now()
          }
        },
        history: []
      }))

      const definition = createAgentDefinition({ preloadMemory })
      const agent = new Agent(definition, createMockAdapterFactory())
      const context = createMockExecutionContext()

      const input = { query: 'test' }
      await agent.execute(input, context)

      expect(preloadMemory).toHaveBeenCalledWith(context)

      const snapshot = agent.getMemorySnapshot()
      expect(snapshot?.sessionMemory['context']).toBeDefined()
      expect(snapshot?.sessionMemory['context'].content).toBe('Preloaded context')
    })

    it('adds input entry to memory history', async () => {
      const definition = createAgentDefinition()
      const agent = new Agent(definition, createMockAdapterFactory())
      const context = createMockExecutionContext()

      const input = { query: 'test query' }
      await agent.execute(input, context)

      const snapshot = agent.getMemorySnapshot()
      const inputEntry = snapshot?.history.find((entry) => entry.type === 'input')

      expect(inputEntry).toBeDefined()
      expect(inputEntry?.content).toContain('test query')
    })
  })

  describe('Phase 2: Iteration Loop (Mock Adapter)', () => {
    it('completes execution with mock adapter', async () => {
      const definition = createAgentDefinition()
      const agent = new Agent(definition, createMockAdapterFactory())
      const context = createMockExecutionContext()

      const input = { query: 'test' }
      const output = await agent.execute(input, context)

      // Mock adapter returns immediately with complete action
      expect(output).toBeDefined()
      expect(output.response).toBe('Mock response')
    })

    it('respects maxIterations constraint', async () => {
      const definition = createAgentDefinition({
        config: {
          type: 'agent',
          kind: 'utility',
          resourceId: 'test-agent',
          name: 'Test Agent',
          description: 'Agent for testing',
          environment: 'dev',
          systemPrompt: 'Test',
          constraints: { maxIterations: 1 }
        }
      })

      const agent = new Agent(definition, createMockAdapterFactory())
      const context = createMockExecutionContext()

      const input = { query: 'test' }

      // Mock adapter completes immediately, so this won't hit the limit
      const output = await agent.execute(input, context)
      expect(output).toBeDefined()
    })

    it('stores reasoning to memory history', async () => {
      const definition = createAgentDefinition()
      const agent = new Agent(definition, createMockAdapterFactory())
      const context = createMockExecutionContext()

      const input = { query: 'test' }
      await agent.execute(input, context)

      const snapshot = agent.getMemorySnapshot()
      const reasoningEntry = snapshot?.history.find((entry) => entry.type === 'reasoning')

      expect(reasoningEntry).toBeDefined()
      expect(reasoningEntry?.content).toContain('Mock reasoning')
    })
  })

  describe('Phase 3: Completion', () => {
    it('generates final output from execution history', async () => {
      const definition = createAgentDefinition()
      const agent = new Agent(definition, createMockAdapterFactory())
      const context = createMockExecutionContext()

      const input = { query: 'test' }
      const output = await agent.execute(input, context)

      expect(output).toBeDefined()
      expect(output.response).toBe('Mock response')
      expect(output.processedAt).toBeDefined()
    })

    it('validates output against contract schema', async () => {
      const definition = createAgentDefinition()
      const agent = new Agent(definition, createMockAdapterFactory())
      const context = createMockExecutionContext()

      const input = { query: 'test' }
      const output = await agent.execute(input, context)

      // If validation failed, execute would throw
      expect(output).toBeDefined()
      expect(output).toHaveProperty('response')
      expect(output).toHaveProperty('processedAt')
    })

    it('captures memory snapshot after completion', async () => {
      const definition = createAgentDefinition()
      const agent = new Agent(definition, createMockAdapterFactory())
      const context = createMockExecutionContext()

      const input = { query: 'test' }
      await agent.execute(input, context)

      const snapshot = agent.getMemorySnapshot()

      expect(snapshot).toBeDefined()
      expect(snapshot?.history).toBeDefined()
      expect(snapshot?.sessionMemory).toBeDefined()
    })
  })

  describe('Error Handling', () => {
    it('throws error from initialization phase on invalid input', async () => {
      const definition = createAgentDefinition()
      const agent = new Agent(definition, createMockAdapterFactory())
      const context = createMockExecutionContext()

      // Invalid input will cause initialization to fail
      const invalidInput = { wrong: 'field' }

      await expect(agent.execute(invalidInput as unknown, context)).rejects.toThrow()
    })

    it('validates output schema strictly', async () => {
      // Create definition with strict output schema
      const strictDefinition = createAgentDefinition({
        contract: {
          inputSchema: z.object({ query: z.string() }),
          outputSchema: z.object({
            response: z.string(),
            count: z.number() // Mock adapter won't provide this
          })
        }
      })

      const agent = new Agent(strictDefinition, createMockAdapterFactory())
      const context = createMockExecutionContext()

      const input = { query: 'test' }

      // Should fail because mock adapter doesn't match strict schema
      await expect(agent.execute(input, context)).rejects.toThrow()
    })

    it('handles preloadMemory callback errors gracefully', async () => {
      const preloadMemory = vi.fn(async () => {
        throw new Error('Failed to load memory from database')
      })

      const definition = createAgentDefinition({ preloadMemory })
      const agent = new Agent(definition, createMockAdapterFactory())
      const context = createMockExecutionContext()

      const input = { query: 'test' }

      await expect(agent.execute(input, context)).rejects.toThrow('Failed to load memory from database')
      expect(preloadMemory).toHaveBeenCalledWith(context)
    })

    it('handles null input gracefully', async () => {
      const definition = createAgentDefinition()
      const agent = new Agent(definition, createMockAdapterFactory())
      const context = createMockExecutionContext()

      await expect(agent.execute(null as unknown, context)).rejects.toThrow()
    })

    it('handles undefined input gracefully', async () => {
      const definition = createAgentDefinition()
      const agent = new Agent(definition, createMockAdapterFactory())
      const context = createMockExecutionContext()

      await expect(agent.execute(undefined as unknown, context)).rejects.toThrow()
    })

    it('handles empty object when required fields missing', async () => {
      const definition = createAgentDefinition({
        contract: {
          inputSchema: z.object({
            query: z.string(),
            requiredField: z.string()
          }),
          outputSchema: z.object({ response: z.string() })
        }
      })

      const agent = new Agent(definition, createMockAdapterFactory())
      const context = createMockExecutionContext()

      const invalidInput = { query: 'test' } // missing requiredField

      await expect(agent.execute(invalidInput as unknown, context)).rejects.toThrow()
    })

    it('handles input with extra unexpected fields', async () => {
      const definition = createAgentDefinition({
        contract: {
          inputSchema: z.object({ query: z.string() }).strict(),
          outputSchema: z.object({ response: z.string(), processedAt: z.string() })
        }
      })

      const agent = new Agent(definition, createMockAdapterFactory())
      const context = createMockExecutionContext()

      const inputWithExtras = {
        query: 'test',
        unexpected: 'field',
        another: 'value'
      }

      // Strict schema should reject extra fields
      await expect(agent.execute(inputWithExtras as unknown, context)).rejects.toThrow()
    })

    it('validates complex nested input schemas', async () => {
      const definition = createAgentDefinition({
        contract: {
          inputSchema: z.object({
            query: z.string(),
            options: z.object({
              temperature: z.number().min(0).max(1),
              maxTokens: z.number().positive()
            })
          }),
          outputSchema: z.object({ response: z.string(), processedAt: z.string() })
        }
      })

      const agent = new Agent(definition, createMockAdapterFactory())
      const context = createMockExecutionContext()

      // Invalid nested field (temperature > 1)
      const invalidInput = {
        query: 'test',
        options: {
          temperature: 1.5, // exceeds max
          maxTokens: 100
        }
      }

      await expect(agent.execute(invalidInput as unknown, context)).rejects.toThrow()
    })

    it('handles output validation failures with detailed error', async () => {
      const definition = createAgentDefinition({
        contract: {
          inputSchema: z.object({ query: z.string() }),
          outputSchema: z.object({
            response: z.string().min(10), // Mock returns shorter string
            processedAt: z.string(),
            requiredField: z.string() // Mock won't provide this
          })
        }
      })

      const agent = new Agent(definition, createMockAdapterFactory())
      const context = createMockExecutionContext()

      const input = { query: 'test' }

      await expect(agent.execute(input, context)).rejects.toThrow()
    })

    it('logs errors to context logger', async () => {
      const definition = createAgentDefinition()
      const agent = new Agent(definition, createMockAdapterFactory())
      const context = createMockExecutionContext()

      const invalidInput = { wrong: 'field' }

      try {
        await agent.execute(invalidInput as unknown, context)
      } catch {
        // Expected to throw
      }

      // Logger should have recorded the error (implementation detail, but useful to verify)
      expect(context.logger.info).toHaveBeenCalled()
    })

    it('handles maxIterations edge case (zero iterations)', async () => {
      const definition = createAgentDefinition({
        config: {
          type: 'agent',
          kind: 'utility',
          resourceId: 'test-agent',
          name: 'Test Agent',
          description: 'Agent for testing',
          environment: 'dev',
          systemPrompt: 'Test',
          constraints: { maxIterations: 0 } // Edge case: no iterations allowed
        }
      })

      const agent = new Agent(definition, createMockAdapterFactory())
      const context = createMockExecutionContext()

      const input = { query: 'test' }

      // Mock adapter completes immediately regardless of maxIterations
      // Test that zero iterations doesn't cause crash
      const output = await agent.execute(input, context)
      expect(output).toBeDefined()
    })

    it('handles negative maxIterations', async () => {
      const definition = createAgentDefinition({
        config: {
          type: 'agent',
          kind: 'utility',
          resourceId: 'test-agent',
          name: 'Test Agent',
          description: 'Agent for testing',
          environment: 'dev',
          systemPrompt: 'Test',
          constraints: { maxIterations: -1 } // Invalid constraint
        }
      })

      const agent = new Agent(definition, createMockAdapterFactory())
      const context = createMockExecutionContext()

      const input = { query: 'test' }

      // Should handle invalid constraint gracefully
      await expect(agent.execute(input, context)).rejects.toThrow()
    })
  })

  describe('Memory Management', () => {
    it('maintains memory across iteration and completion phases', async () => {
      const definition = createAgentDefinition()
      const agent = new Agent(definition, createMockAdapterFactory())
      const context = createMockExecutionContext()

      const input = { query: 'test' }
      await agent.execute(input, context)

      const snapshot = agent.getMemorySnapshot()

      // Should have input entry, reasoning entry
      expect(snapshot?.history.length).toBeGreaterThanOrEqual(2)

      const inputEntry = snapshot?.history.find((e) => e.type === 'input')
      const reasoningEntry = snapshot?.history.find((e) => e.type === 'reasoning')

      expect(inputEntry).toBeDefined()
      expect(reasoningEntry).toBeDefined()
    })

    it('preloaded session memory is preserved', async () => {
      const preloadMemory = async () => ({
        sessionMemory: {
          task: {
            type: 'context' as const,
            content: 'Build feature X',
            timestamp: Date.now()
          },
          guidelines: {
            type: 'context' as const,
            content: 'Follow best practices',
            timestamp: Date.now()
          }
        },
        history: []
      })

      const definition = createAgentDefinition({ preloadMemory })
      const agent = new Agent(definition, createMockAdapterFactory())
      const context = createMockExecutionContext()

      const input = { query: 'test' }
      await agent.execute(input, context)

      const snapshot = agent.getMemorySnapshot()

      expect(snapshot?.sessionMemory['task']).toBeDefined()
      expect(snapshot?.sessionMemory['task'].content).toBe('Build feature X')
      expect(snapshot?.sessionMemory['guidelines']).toBeDefined()
      expect(snapshot?.sessionMemory['guidelines'].content).toBe('Follow best practices')
    })
  })

  describe('Execution Context', () => {
    it('passes execution context to all phases', async () => {
      const definition = createAgentDefinition()
      const agent = new Agent(definition, createMockAdapterFactory())
      const context = createMockExecutionContext()

      const input = { query: 'test' }
      await agent.execute(input, context)

      // Execution completed successfully, context was passed through
      expect(context.executionId).toBe('test-exec-123')
      expect(context.organizationId).toBe('test-org-123')
      expect(context.userId).toBe('test-user-123')
    })

    it('uses scoped logger for execution tracking', async () => {
      const definition = createAgentDefinition()
      const agent = new Agent(definition, createMockAdapterFactory())
      const context = createMockExecutionContext()

      const input = { query: 'test' }
      await agent.execute(input, context)

      // Logger should have been called for lifecycle events
      expect(context.logger.info).toHaveBeenCalled()
    })
  })

  describe('Tool Registry', () => {
    it('initializes registry with base tools', () => {
      const tools = [createMockTool('tool1'), createMockTool('tool2')]
      const definition = createAgentDefinition({ tools })
      const agent = new Agent(definition, createMockAdapterFactory())

      const registeredTools = agent.getTools()
      expect(registeredTools).toHaveLength(2)
      expect(registeredTools.find(t => t.name === 'tool1')).toBeDefined()
      expect(registeredTools.find(t => t.name === 'tool2')).toBeDefined()
    })

    it('registers new tools dynamically', () => {
      const tool1 = createMockTool('tool1')
      const definition = createAgentDefinition({ tools: [tool1] })
      const agent = new Agent(definition, createMockAdapterFactory())

      const newTools = [createMockTool('tool2'), createMockTool('tool3')]
      agent.registerTools(newTools)

      const registeredTools = agent.getTools()
      expect(registeredTools).toHaveLength(3)
      expect(registeredTools.find(t => t.name === 'tool1')).toBeDefined()
      expect(registeredTools.find(t => t.name === 'tool2')).toBeDefined()
      expect(registeredTools.find(t => t.name === 'tool3')).toBeDefined()
    })

    it('gracefully skips duplicate tool registration', () => {
      const tool1 = createMockTool('tool1')
      const definition = createAgentDefinition({ tools: [tool1] })
      const agent = new Agent(definition, createMockAdapterFactory())

      const duplicateTool = createMockTool('tool1')
      // Should not throw - gracefully skips duplicate
      agent.registerTools([duplicateTool])

      // Tool count should remain 1 (duplicate skipped)
      const registeredTools = agent.getTools()
      expect(registeredTools).toHaveLength(1)
    })

    it('getTools returns array of all registered tools', () => {
      const tools = [createMockTool('tool1'), createMockTool('tool2'), createMockTool('tool3')]
      const definition = createAgentDefinition({ tools })
      const agent = new Agent(definition, createMockAdapterFactory())

      const registeredTools = agent.getTools()
      expect(registeredTools).toHaveLength(3)
    })

    it('allows registering multiple tools at once', () => {
      const tool1 = createMockTool('tool1')
      const definition = createAgentDefinition({ tools: [tool1] })
      const agent = new Agent(definition, createMockAdapterFactory())

      const newTools = [
        createMockTool('tool2'),
        createMockTool('tool3'),
        createMockTool('tool4')
      ]
      agent.registerTools(newTools)

      const registeredTools = agent.getTools()
      expect(registeredTools).toHaveLength(4)
    })
  })
})
