#!/usr/bin/env python3
"""
Agent Orchestrator - Comprehensive System Test

This script tests all major components of the orchestration system.
Run with: python scripts/run_system_test.py
"""

import asyncio
import sys
import os
import tempfile
from pathlib import Path
from datetime import datetime

# Add src to path
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))

# Test results tracking
results = []
passed = 0
failed = 0


def test_result(name: str, success: bool, message: str = ""):
    """Record a test result."""
    global passed, failed
    status = "✓ PASS" if success else "✗ FAIL"
    results.append((name, success, message))
    if success:
        passed += 1
        print(f"  {status}: {name}")
    else:
        failed += 1
        print(f"  {status}: {name} - {message}")


def section(title: str):
    """Print a section header."""
    print()
    print("=" * 60)
    print(f"  {title}")
    print("=" * 60)


async def main():
    """Run all system tests."""
    print()
    print("╔══════════════════════════════════════════════════════════╗")
    print("║       Agent Orchestrator - System Test Suite             ║")
    print("║                     Version 1.8                          ║")
    print("╚══════════════════════════════════════════════════════════╝")
    print()
    print(f"  Start Time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"  Python: {sys.version.split()[0]}")

    # =========================================================================
    # Phase 1: Core Imports
    # =========================================================================
    section("Phase 1: Core Module Imports")

    # Config
    try:
        from agent_orchestrator.config import Config, get_config
        test_result("Import config module", True)
    except Exception as e:
        test_result("Import config module", False, str(e))

    # Persistence
    try:
        from agent_orchestrator.persistence import OrchestratorDB
        test_result("Import persistence module", True)
    except Exception as e:
        test_result("Import persistence module", False, str(e))

    # Journal
    try:
        from agent_orchestrator.journal import (
            ProjectJournal, StatusPacket, TaskArtifacts,
        )
        test_result("Import journal module", True)
    except Exception as e:
        test_result("Import journal module", False, str(e))

    # Secrets
    try:
        from agent_orchestrator.secrets import (
            SecretRedactor, SecretFileGuard,
        )
        test_result("Import secrets module", True)
    except Exception as e:
        test_result("Import secrets module", False, str(e))

    # Adapters
    try:
        from agent_orchestrator.adapters import (
            BaseAdapter, LLMAdapter, CLIAgentAdapter,
            AgentResponse, UsageStats,
        )
        test_result("Import adapters module", True)
    except Exception as e:
        test_result("Import adapters module", False, str(e))

    # =========================================================================
    # Phase 2: Memory Layer
    # =========================================================================
    section("Phase 2: Memory Layer Imports")

    try:
        from agent_orchestrator.memory import (
            OperationalMemory, KnowledgeMemory, WorkingMemory,
            MemoryWriteGate,
            MemoryItem, MemoryPatch, MemoryTier,
        )
        test_result("Import memory module", True)
    except Exception as e:
        test_result("Import memory module", False, str(e))

    # =========================================================================
    # Phase 3: Control Loop & Risk
    # =========================================================================
    section("Phase 3: Control & Risk Imports")

    try:
        from agent_orchestrator.control import (
            AgentControlLoop, AgentHealthCheck, ControlAction,
        )
        test_result("Import control module", True)
    except Exception as e:
        test_result("Import control module", False, str(e))

    try:
        from agent_orchestrator.risk import (
            RiskPolicy, AutonomyGate, RiskDecision,
        )
        test_result("Import risk module", True)
    except Exception as e:
        test_result("Import risk module", False, str(e))

    # =========================================================================
    # Phase 4: Interrupt Handlers
    # =========================================================================
    section("Phase 4: Interrupt Handler Imports")

    try:
        from agent_orchestrator.interrupt import (
            CLIInterruptHandler, AsyncInterruptHandler,
            ApprovalDecision, ApprovalResponse,
        )
        test_result("Import interrupt module", True)
    except Exception as e:
        test_result("Import interrupt module", False, str(e))

    # =========================================================================
    # Phase 5: Budget Controls
    # =========================================================================
    section("Phase 5: Budget Control Imports")

    try:
        from agent_orchestrator.budget import (
            AgentBudget, BudgetEnforcer, DailyUsage,
            MCPServerBudget, ToolBudget, MCPUsageTracker,
            MCPRegistry, AccessLevel, RegistryDecision,
        )
        test_result("Import budget module", True)
    except Exception as e:
        test_result("Import budget module", False, str(e))

    # =========================================================================
    # Phase 6: Merge Gate
    # =========================================================================
    section("Phase 6: Merge Gate Imports")

    try:
        from agent_orchestrator.merge import (
            MergeReadiness, ReadinessReport, CheckResult,
            MergeGate, MergeRequest, MergeResult,
        )
        test_result("Import merge module", True)
    except Exception as e:
        test_result("Import merge module", False, str(e))

    # =========================================================================
    # Phase 7: Observability
    # =========================================================================
    section("Phase 7: Observability Imports")

    try:
        from agent_orchestrator.observability import (
            TokenAuditClient, UsageEntry, UsageSummary,
            AIObserverClient, MetricPoint,
            AlertManager, Alert, AlertRule, AlertSeverity,
            AuditTrail, AuditEvent, EventType, EventSeverity,
        )
        test_result("Import observability module", True)
    except Exception as e:
        test_result("Import observability module", False, str(e))

    # =========================================================================
    # Phase 8: Reliability
    # =========================================================================
    section("Phase 8: Reliability Imports")

    try:
        from agent_orchestrator.reliability import (
            RateLimiter, RateLimitConfig, RateLimiterRegistry,
            RetryConfig, RetryManager, CircuitBreaker,
            FallbackModel, FallbackConfig, retry,
            GracefulShutdown, ShutdownConfig, ShutdownState,
        )
        test_result("Import reliability module", True)
    except Exception as e:
        test_result("Import reliability module", False, str(e))

    # =========================================================================
    # Functional Tests
    # =========================================================================
    section("Functional Tests: Configuration")

    try:
        config = Config()
        test_result("Create Config instance", True)
        test_result("Config has api_keys", hasattr(config, 'api_keys'))
        test_result("Config has database", hasattr(config, 'database'))
        test_result("Config has budgets", hasattr(config, 'budgets'))
    except Exception as e:
        test_result("Create Config instance", False, str(e))

    section("Functional Tests: Database")

    try:
        from agent_orchestrator.persistence import OrchestratorDB
        with tempfile.TemporaryDirectory() as tmpdir:
            db_path = Path(tmpdir) / "test.db"
            db = OrchestratorDB(db_path)
            test_result("Create OrchestratorDB", True)
            test_result("Database file exists", db_path.exists())
    except Exception as e:
        test_result("Database operations", False, str(e))

    section("Functional Tests: Secret Redactor")

    try:
        from agent_orchestrator.secrets import SecretRedactor

        # Test API key redaction
        test_text = "My key is sk-ant-api03-1234567890abcdef"
        redacted = SecretRedactor.redact(test_text)
        has_key = "sk-ant-api03" in redacted
        test_result("Redact API key", not has_key,
                    "Key still visible" if has_key else "")

        # Test password redaction
        test_text2 = 'password = "secret123"'
        redacted2 = SecretRedactor.redact(test_text2)
        test_result("Redact password", "secret123" not in redacted2)

    except Exception as e:
        test_result("Secret redaction", False, str(e))

    section("Functional Tests: Risk Gate")

    try:
        from agent_orchestrator.risk import RiskPolicy
        policy = RiskPolicy()

        # Test file classification
        high_risk = policy.classify_file(".env")
        test_result("Classify .env as high risk", high_risk.name == "HIGH")

        medium_risk = policy.classify_file("src/main.py")
        test_result("Classify source file", medium_risk.name in ["LOW", "MEDIUM"])

        # Test command classification
        critical = policy.classify_command("git push --force origin main")
        test_result("Classify force push as critical", critical.name == "CRITICAL")

        low_risk = policy.classify_command("ls -la")
        test_result("Classify ls as low risk", low_risk.name == "LOW")

    except Exception as e:
        test_result("Risk classification", False, str(e))

    section("Functional Tests: Rate Limiter")

    try:
        from agent_orchestrator.reliability import RateLimiter, RateLimitConfig

        config = RateLimitConfig(
            requests_per_minute=60,
            tokens_per_minute=100000
        )
        limiter = RateLimiter(provider="test", config=config)
        test_result("Create RateLimiter", True)

        # Test acquire
        wait_time = await limiter.acquire(tokens=1000)
        test_result("Acquire rate limit", wait_time >= 0)

        # Test success recording
        await limiter.record_success()
        test_result("Record success", True)

    except Exception as e:
        test_result("Rate limiter", False, str(e))

    section("Functional Tests: Circuit Breaker")

    try:
        from agent_orchestrator.reliability import CircuitBreaker

        breaker = CircuitBreaker(
            failure_threshold=3,
            recovery_timeout=0.1
        )
        test_result("Create CircuitBreaker", True)

        # Initial state should be closed
        test_result("Initial state is closed",
                    breaker.state == CircuitBreaker.State.CLOSED)

        # Can execute initially
        can_exec = await breaker.can_execute()
        test_result("Can execute in closed state", can_exec)

        # Record failures to trip breaker
        for _ in range(3):
            await breaker.record_failure()

        test_result("Circuit opens after failures",
                    breaker.state == CircuitBreaker.State.OPEN)

    except Exception as e:
        test_result("Circuit breaker", False, str(e))

    section("Functional Tests: Graceful Shutdown")

    try:
        from agent_orchestrator.reliability import GracefulShutdown, ShutdownConfig, ShutdownState

        shutdown = GracefulShutdown(
            config=ShutdownConfig(grace_period_seconds=5)
        )
        test_result("Create GracefulShutdown", True)

        # Test state
        test_result("Initial state is running",
                    shutdown.state == ShutdownState.RUNNING)

        # Test work tracking
        shutdown.track_work("test-work-1")
        test_result("Track in-flight work",
                    shutdown.get_in_flight_count() == 1)

        shutdown.complete_work("test-work-1")
        test_result("Complete work",
                    shutdown.get_in_flight_count() == 0)

        # Register a shutdown task
        task_ran = False
        async def test_task():
            nonlocal task_ran
            task_ran = True

        shutdown.register_task("test", test_task)
        result = await shutdown.shutdown()
        test_result("Shutdown runs tasks", task_ran and result.success)

    except Exception as e:
        test_result("Graceful shutdown", False, str(e))

    section("Functional Tests: Audit Trail")

    try:
        from agent_orchestrator.observability import AuditTrail, AuditConfig, EventType

        with tempfile.TemporaryDirectory() as tmpdir:
            audit = AuditTrail(config=AuditConfig(log_directory=tmpdir))
            test_result("Create AuditTrail", True)

            # Log an event
            event = audit.log_event(
                event_type=EventType.AGENT_STARTED,
                message="Test agent started",
                agent_id="test-agent"
            )
            test_result("Log audit event", event.event_id is not None)

            # Query events
            events = audit.query_events(event_types=[EventType.AGENT_STARTED])
            test_result("Query audit events", len(events) == 1)

    except Exception as e:
        test_result("Audit trail", False, str(e))

    section("Functional Tests: Budget Enforcer")

    try:
        from agent_orchestrator.budget import AgentBudget, BudgetEnforcer
        from agent_orchestrator.persistence import OrchestratorDB

        with tempfile.TemporaryDirectory() as tmpdir:
            db_path = Path(tmpdir) / "budget.db"
            db = OrchestratorDB(db_path)

            # Create budget with actual API
            budget = AgentBudget(
                agent_id="test-agent",
                daily_input_tokens=10000,
                daily_output_tokens=5000,
                daily_cost_limit=10.0
            )
            test_result("Create AgentBudget", True)

            # Pass budgets dict to constructor
            enforcer = BudgetEnforcer(
                db=db,
                budgets={"test-agent": budget}
            )
            test_result("Create BudgetEnforcer", True)

            # Check if under budget using actual API
            result = enforcer.check_budget("test-agent", estimated_input_tokens=1000, estimated_output_tokens=500)
            test_result("Check budget (under limit)", result.allowed)

    except Exception as e:
        test_result("Budget enforcer", False, str(e))

    section("Functional Tests: Merge Gate")

    try:
        from agent_orchestrator.merge import MergeGate
        from agent_orchestrator.persistence import OrchestratorDB

        with tempfile.TemporaryDirectory() as tmpdir:
            db_path = Path(tmpdir) / "merge.db"
            db = OrchestratorDB(db_path)

            # Create gate with config dict
            gate = MergeGate(
                db=db,
                config={"protected_branches": ["main", "master"]}
            )
            test_result("Create MergeGate", True)

            # Check if branch is protected
            test_result("Main is protected", gate.is_protected("main"))
            test_result("Feature not protected", not gate.is_protected("feature/test"))

    except Exception as e:
        test_result("Merge gate", False, str(e))

    section("Functional Tests: MCP Registry")

    try:
        from agent_orchestrator.budget import MCPRegistry, AccessLevel, RegistryDecision

        registry = MCPRegistry()
        test_result("Create MCPRegistry", True)

        # Register a server with correct API
        registry.register_server(
            server_name="test-server",
            access_level=AccessLevel.ALLOWED,
            allowed_tools={"read", "write"}
        )
        test_result("Register MCP server", True)

        # Check access
        decision = registry.check_access("test-server", "read")
        test_result("Check MCP access", decision.allowed)

    except Exception as e:
        test_result("MCP registry", False, str(e))

    section("Functional Tests: Alert Manager")

    try:
        from agent_orchestrator.observability import AlertManager, AlertRule, AlertSeverity, AlertType
        from agent_orchestrator.persistence import OrchestratorDB

        with tempfile.TemporaryDirectory() as tmpdir:
            db_path = Path(tmpdir) / "alerts.db"
            db = OrchestratorDB(db_path)

            manager = AlertManager(db=db)
            test_result("Create AlertManager", True)

            # Create a rule using correct API
            rule = AlertRule(
                name="test-rule",
                alert_type=AlertType.STUCK_AGENT,
                severity=AlertSeverity.WARNING,
                threshold=100
            )
            manager.add_rule(rule)
            test_result("Add alert rule", True)

    except Exception as e:
        test_result("Alert manager", False, str(e))

    # =========================================================================
    # Summary
    # =========================================================================
    section("Test Summary")

    print()
    print(f"  Total Tests: {passed + failed}")
    print(f"  Passed:      {passed}")
    print(f"  Failed:      {failed}")
    print()

    if failed == 0:
        print("  ╔════════════════════════════════════════╗")
        print("  ║   All Tests Passed! System Ready.      ║")
        print("  ╚════════════════════════════════════════╝")
        return 0
    else:
        print("  ╔════════════════════════════════════════╗")
        print("  ║   Some Tests Failed. Review Above.     ║")
        print("  ╚════════════════════════════════════════╝")
        print()
        print("  Failed tests:")
        for name, success, msg in results:
            if not success:
                print(f"    - {name}: {msg}")
        return 1


if __name__ == "__main__":
    exit_code = asyncio.run(main())
    sys.exit(exit_code)