#!/usr/bin/env python3
"""
Test runner for Graphiti MCP integration tests.
Provides various test execution modes and reporting options.
"""

import argparse
import asyncio
import json
import os
import subprocess
import sys
import time
from pathlib import Path
from typing import Dict, List, Optional

import pytest


class TestRunner:
    """Orchestrate test execution with various configurations."""

    def __init__(self, args):
        self.args = args
        self.test_dir = Path(__file__).parent
        self.results = {}

    def check_prerequisites(self) -> Dict[str, bool]:
        """Check if required services and dependencies are available."""
        checks = {}

        # Check for OpenAI API key if not using mocks
        if not self.args.mock_llm:
            checks['openai_api_key'] = bool(os.environ.get('OPENAI_API_KEY'))
        else:
            checks['openai_api_key'] = True

        # Check database availability based on backend
        if self.args.database == 'neo4j':
            checks['neo4j'] = self._check_neo4j()
        elif self.args.database == 'falkordb':
            checks['falkordb'] = self._check_falkordb()
        elif self.args.database == 'kuzu':
            checks['kuzu'] = True  # KuzuDB is embedded

        # Check Python dependencies
        checks['mcp'] = self._check_python_package('mcp')
        checks['pytest'] = self._check_python_package('pytest')
        checks['pytest-asyncio'] = self._check_python_package('pytest-asyncio')

        return checks

    def _check_neo4j(self) -> bool:
        """Check if Neo4j is available."""
        try:
            import neo4j
            # Try to connect
            uri = os.environ.get('NEO4J_URI', 'bolt://localhost:7687')
            user = os.environ.get('NEO4J_USER', 'neo4j')
            password = os.environ.get('NEO4J_PASSWORD', 'graphiti')

            driver = neo4j.GraphDatabase.driver(uri, auth=(user, password))
            with driver.session() as session:
                session.run("RETURN 1")
            driver.close()
            return True
        except:
            return False

    def _check_falkordb(self) -> bool:
        """Check if FalkorDB is available."""
        try:
            import redis
            uri = os.environ.get('FALKORDB_URI', 'redis://localhost:6379')
            r = redis.from_url(uri)
            r.ping()
            return True
        except:
            return False

    def _check_python_package(self, package: str) -> bool:
        """Check if a Python package is installed."""
        try:
            __import__(package.replace('-', '_'))
            return True
        except ImportError:
            return False

    def run_test_suite(self, suite: str) -> int:
        """Run a specific test suite."""
        pytest_args = ['-v', '--tb=short']

        # Add database marker
        if self.args.database:
            pytest_args.extend(['-m', f'not requires_{db}'
                              for db in ['neo4j', 'falkordb', 'kuzu']
                              if db != self.args.database])

        # Add suite-specific arguments
        if suite == 'unit':
            pytest_args.extend(['-m', 'unit', 'test_*.py'])
        elif suite == 'integration':
            pytest_args.extend(['-m', 'integration or not unit', 'test_*.py'])
        elif suite == 'comprehensive':
            pytest_args.append('test_comprehensive_integration.py')
        elif suite == 'async':
            pytest_args.append('test_async_operations.py')
        elif suite == 'stress':
            pytest_args.extend(['-m', 'slow', 'test_stress_load.py'])
        elif suite == 'smoke':
            # Quick smoke test - just basic operations
            pytest_args.extend([
                'test_comprehensive_integration.py::TestCoreOperations::test_server_initialization',
                'test_comprehensive_integration.py::TestCoreOperations::test_add_text_memory'
            ])
        elif suite == 'all':
            pytest_args.append('.')
        else:
            pytest_args.append(suite)

        # Add coverage if requested
        if self.args.coverage:
            pytest_args.extend(['--cov=../src', '--cov-report=html'])

        # Add parallel execution if requested
        if self.args.parallel:
            pytest_args.extend(['-n', str(self.args.parallel)])

        # Add verbosity
        if self.args.verbose:
            pytest_args.append('-vv')

        # Add markers to skip
        if self.args.skip_slow:
            pytest_args.extend(['-m', 'not slow'])

        # Add timeout override
        if self.args.timeout:
            pytest_args.extend(['--timeout', str(self.args.timeout)])

        # Add environment variables
        env = os.environ.copy()
        if self.args.mock_llm:
            env['USE_MOCK_LLM'] = 'true'
        if self.args.database:
            env['DATABASE_PROVIDER'] = self.args.database

        # Run tests
        print(f"Running {suite} tests with pytest args: {' '.join(pytest_args)}")
        return pytest.main(pytest_args)

    def run_performance_benchmark(self):
        """Run performance benchmarking suite."""
        print("Running performance benchmarks...")

        # Import test modules
        from test_comprehensive_integration import TestPerformance
        from test_async_operations import TestAsyncPerformance

        # Run performance tests
        result = pytest.main([
            '-v',
            'test_comprehensive_integration.py::TestPerformance',
            'test_async_operations.py::TestAsyncPerformance',
            '--benchmark-only' if self.args.benchmark_only else '',
        ])

        return result

    def generate_report(self):
        """Generate test execution report."""
        report = []
        report.append("\n" + "=" * 60)
        report.append("GRAPHITI MCP TEST EXECUTION REPORT")
        report.append("=" * 60)

        # Prerequisites check
        checks = self.check_prerequisites()
        report.append("\nPrerequisites:")
        for check, passed in checks.items():
            status = "✅" if passed else "❌"
            report.append(f"  {status} {check}")

        # Test configuration
        report.append(f"\nConfiguration:")
        report.append(f"  Database: {self.args.database}")
        report.append(f"  Mock LLM: {self.args.mock_llm}")
        report.append(f"  Parallel: {self.args.parallel or 'No'}")
        report.append(f"  Timeout: {self.args.timeout}s")

        # Results summary (if available)
        if self.results:
            report.append(f"\nResults:")
            for suite, result in self.results.items():
                status = "✅ Passed" if result == 0 else f"❌ Failed ({result})"
                report.append(f"  {suite}: {status}")

        report.append("=" * 60)
        return "\n".join(report)


def main():
    """Main entry point for test runner."""
    parser = argparse.ArgumentParser(
        description='Run Graphiti MCP integration tests',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Test Suites:
  unit          - Run unit tests only
  integration   - Run integration tests
  comprehensive - Run comprehensive integration test suite
  async         - Run async operation tests
  stress        - Run stress and load tests
  smoke         - Run quick smoke tests
  all           - Run all tests

Examples:
  python run_tests.py smoke                    # Quick smoke test
  python run_tests.py integration --parallel 4 # Run integration tests in parallel
  python run_tests.py stress --database neo4j  # Run stress tests with Neo4j
  python run_tests.py all --coverage          # Run all tests with coverage
        """
    )

    parser.add_argument(
        'suite',
        choices=['unit', 'integration', 'comprehensive', 'async', 'stress', 'smoke', 'all'],
        help='Test suite to run'
    )

    parser.add_argument(
        '--database',
        choices=['neo4j', 'falkordb', 'kuzu'],
        default='kuzu',
        help='Database backend to test (default: kuzu)'
    )

    parser.add_argument(
        '--mock-llm',
        action='store_true',
        help='Use mock LLM for faster testing'
    )

    parser.add_argument(
        '--parallel',
        type=int,
        metavar='N',
        help='Run tests in parallel with N workers'
    )

    parser.add_argument(
        '--coverage',
        action='store_true',
        help='Generate coverage report'
    )

    parser.add_argument(
        '--verbose',
        action='store_true',
        help='Verbose output'
    )

    parser.add_argument(
        '--skip-slow',
        action='store_true',
        help='Skip slow tests'
    )

    parser.add_argument(
        '--timeout',
        type=int,
        default=300,
        help='Test timeout in seconds (default: 300)'
    )

    parser.add_argument(
        '--benchmark-only',
        action='store_true',
        help='Run only benchmark tests'
    )

    parser.add_argument(
        '--check-only',
        action='store_true',
        help='Only check prerequisites without running tests'
    )

    args = parser.parse_args()

    # Create test runner
    runner = TestRunner(args)

    # Check prerequisites
    if args.check_only:
        print(runner.generate_report())
        sys.exit(0)

    # Check if prerequisites are met
    checks = runner.check_prerequisites()
    if not all(checks.values()):
        print("⚠️  Some prerequisites are not met:")
        for check, passed in checks.items():
            if not passed:
                print(f"  ❌ {check}")

        if not args.mock_llm and not checks.get('openai_api_key'):
            print("\nHint: Use --mock-llm to run tests without OpenAI API key")

        response = input("\nContinue anyway? (y/N): ")
        if response.lower() != 'y':
            sys.exit(1)

    # Run tests
    print(f"\n🚀 Starting test execution: {args.suite}")
    start_time = time.time()

    if args.benchmark_only:
        result = runner.run_performance_benchmark()
    else:
        result = runner.run_test_suite(args.suite)

    duration = time.time() - start_time

    # Store results
    runner.results[args.suite] = result

    # Generate and print report
    print(runner.generate_report())
    print(f"\n⏱️  Test execution completed in {duration:.2f} seconds")

    # Exit with test result code
    sys.exit(result)


if __name__ == "__main__":
    main()