"""
Phase 00: Foundation & Infrastructure
Enhanced with Anthropic Agent SDK patterns

This implementation demonstrates the complete agent feedback loop:
gather context → take action → verify work → repeat
"""

import sys
import os

# Add paths
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../../..', 'guardrails'))
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../../..', 'agent_framework'))

from multi_layer_system import MultiLayerGuardrailSystem
from feedback_loop import AgentFeedbackLoop
from context_manager import ContextManager
from subagent_orchestrator import SubagentOrchestrator
from agentic_search import AgenticSearch
from verification_system import MultiMethodVerifier

import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class Phase00Implementation:
    """
    Enhanced implementation for Phase 00: Foundation & Infrastructure

    Implements Anthropic's agent feedback loop pattern:
    1. Gather Context - Use agentic search to understand requirements
    2. Take Action - Implement infrastructure components
    3. Verify Work - Multi-method verification
    4. Repeat - Iterate until verification passes

    Features:
    - Agent feedback loop for self-correction
    - Context compaction for long-running tasks
    - Subagent orchestration for parallel execution
    - Multi-method verification (rules + guardrails + code)
    """

    def __init__(self):
        self.phase_id = 0
        self.phase_name = "Foundation & Infrastructure"

        # Initialize agent framework components
        self.guardrails = MultiLayerGuardrailSystem()
        self.feedback_loop = AgentFeedbackLoop(
            max_iterations=10,
            enable_learning=True,
            log_file=f"phase_{self.phase_id}_feedback_log.json"
        )
        self.context = ContextManager(max_tokens=100000, compact_threshold=0.8)
        self.orchestrator = SubagentOrchestrator(max_parallel=5)
        self.search = AgenticSearch()
        self.verifier = MultiMethodVerifier()

        self.status = "NOT_STARTED"

        logger.info(f"Phase {self.phase_id} initialized with agent framework")

    def gather_context(self, task, iteration_log):
        """
        STEP 1: Gather context for this task

        Uses:
        - Agentic search to explore codebase
        - Analysis of previous phases
        - Documentation search
        - Subagents for parallel context gathering (if needed)
        """
        logger.info(f"Step 1: Gathering context for task: {task.get('goal')}")

        # Add task to context manager
        self.context.add_message("user", f"Task: {task.get('goal')}")

        # Use agentic search to gather comprehensive context
        context = self.search.gather_context_for_phase(self.phase_id)

        # If this is a retry, learn from previous failures
        if iteration_log:
            previous_errors = [
                log.verification.get("message")
                for log in iteration_log
                if not log.success
            ]
            context["previous_errors"] = previous_errors
            context["retry_count"] = len(iteration_log)

            logger.info(f"Retry #{len(iteration_log)}, learned from {len(previous_errors)} previous errors")

        # For complex tasks, spawn subagents for parallel search
        if task.get("requires_extensive_search"):
            logger.info("Task requires extensive search, spawning subagents...")

            search_tasks = [
                {"goal": "search previous phase implementations"},
                {"goal": "analyze dependencies"},
                {"goal": "review documentation"}
            ]

            # Spawn subagents in parallel
            subagent_ids = self.orchestrator.spawn_parallel(
                tasks=search_tasks,
                context_gatherer=lambda t, log: {"task": t},
                action_executor=lambda t, ctx: self.search.gather_context_for_phase(self.phase_id),
                verifier=lambda out, ctx, t: {"passed": True, "message": "Search complete"}
            )

            # Wait for subagents
            results = self.orchestrator.wait_for_subagents(subagent_ids, timeout=60)

            # Merge subagent results
            merged = self.orchestrator.merge_subagent_results(results)
            context["subagent_findings"] = merged["key_findings"]

        # Add to context manager
        self.context.add_message("system", f"Gathered context with {len(context)} components")

        return context

    def take_action(self, task, context):
        """
        STEP 2: Take action based on context

        Implements the actual phase logic:
        - Set up cloud infrastructure
        - Configure Kubernetes
        - Initialize Neo4j with medical ontologies
        """
        logger.info("Step 2: Taking action")

        # Add to context
        self.context.add_message("assistant", "Implementing phase logic...")

        # Get requirements from context
        phase_info = context.get("phase_info", {})
        description = phase_info.get("description", "")

        # Implement phase logic
        output = {
            "phase_id": self.phase_id,
            "phase_name": self.phase_name,
            "status": "implemented",
            "components": {
                "cloud_infrastructure": self._setup_cloud_infrastructure(context),
                "kubernetes": self._configure_kubernetes(context),
                "neo4j": self._initialize_neo4j(context)
            },
            "description": description,
            "iteration": context.get("retry_count", 0) + 1
        }

        # Add result to context
        self.context.add_message("assistant", f"Implementation complete (iteration {output['iteration']})")

        return output

    def verify_work(self, output, context, task):
        """
        STEP 3: Verify the work

        Uses multi-method verification:
        1. Rules-based verification (required fields, structure)
        2. Guardrails verification (medical safety, HIPAA)
        3. Code verification (if code generated)
        """
        logger.info("Step 3: Verifying work")

        # Use multi-method verifier
        verification = self.verifier.verify_output(
            output=output,
            context={
                "input": task.get("goal", ""),
                "phase_id": self.phase_id
            },
            output_type="data",
            task={
                "expected_type": "dict",
                "required_fields": ["phase_id", "status", "components"]
            }
        )

        # Add verification result to context
        result_msg = "✅ Verification passed" if verification["overall_passed"] else "❌ Verification failed"
        self.context.add_message("system", result_msg)

        return {
            "passed": verification["overall_passed"],
            "message": result_msg,
            "details": verification["method_results"],
            "recommendations": verification.get("all_recommendations", [])
        }

    def execute(self, task):
        """
        Main execution with feedback loop

        This orchestrates the complete agent feedback loop:
        - Gathers context (with agentic search + subagents)
        - Takes action (implements phase)
        - Verifies work (multi-method verification)
        - Repeats if verification fails (with learning)

        Args:
            task: dict with:
                - goal: str (what to accomplish)
                - requires_extensive_search: bool (optional)

        Returns:
            FeedbackLoopResult with success status and output
        """
        logger.info(f"Executing Phase {self.phase_id} with agent feedback loop")

        self.status = "IN_PROGRESS"

        # Execute with feedback loop
        result = self.feedback_loop.execute(
            task=task,
            context_gatherer=self.gather_context,
            action_executor=self.take_action,
            verifier=self.verify_work
        )

        # Update phase state based on result
        if result.success:
            self.status = "COMPLETED"
            self._update_phase_state("COMPLETED", result)
            logger.info(f"✅ Phase {self.phase_id} completed successfully after {result.iterations} iterations")
        else:
            self.status = "FAILED"
            self._update_phase_state("FAILED", result)
            logger.error(f"❌ Phase {self.phase_id} failed after {result.iterations} iterations")

        return result

    # Phase-specific implementation methods

    def _setup_cloud_infrastructure(self, context):
        """Set up cloud infrastructure (Azure/AWS/GCP)"""
        logger.info("Setting up cloud infrastructure...")

        # In production, this would:
        # 1. Provision cloud resources
        # 2. Configure networking
        # 3. Set up security groups
        # 4. Configure storage

        return {
            "status": "configured",
            "provider": "Azure",
            "region": "eastus",
            "components": ["compute", "storage", "networking"]
        }

    def _configure_kubernetes(self, context):
        """Configure Kubernetes cluster"""
        logger.info("Configuring Kubernetes...")

        # In production, this would:
        # 1. Create K8s cluster
        # 2. Configure namespaces
        # 3. Set up RBAC
        # 4. Deploy monitoring

        return {
            "status": "configured",
            "version": "1.28",
            "nodes": 3,
            "namespaces": ["swarmcare", "monitoring", "ingress"]
        }

    def _initialize_neo4j(self, context):
        """Initialize Neo4j with medical ontologies"""
        logger.info("Initializing Neo4j with medical ontologies...")

        # In production, this would:
        # 1. Deploy Neo4j instance
        # 2. Load 13 medical ontologies
        # 3. Create indexes
        # 4. Set up graph schema

        return {
            "status": "initialized",
            "version": "5.x",
            "ontologies": [
                "SNOMED CT",
                "ICD-10",
                "LOINC",
                "RxNorm",
                "CPT",
                "UMLS",
                "MeSH",
                "HPO",
                "GO",
                "ChEBI",
                "DrugBank",
                "OMIM",
                "OrphaNet"
            ],
            "relationships": "configured"
        }

    def _update_phase_state(self, status, result):
        """Update phase state tracking"""
        import json
        from pathlib import Path

        state_path = Path(__file__).parent.parent / ".state/phase_state.json"

        state = {
            "phase_id": self.phase_id,
            "status": status,
            "iterations": result.iterations,
            "duration_seconds": result.total_duration_seconds,
            "completed_at": result.iteration_log[-1].timestamp if result.iteration_log else None,
            "success": result.success
        }

        with open(state_path, 'w') as f:
            json.dump(state, f, indent=2)

        logger.info(f"Updated phase state to: {status}")

    def get_statistics(self):
        """Get execution statistics"""
        return {
            "phase_id": self.phase_id,
            "phase_name": self.phase_name,
            "status": self.status,
            "context_stats": self.context.get_statistics(),
            "feedback_loop_stats": self.feedback_loop.get_statistics() if self.feedback_loop.iteration_log else {},
            "orchestrator_stats": self.orchestrator.get_statistics(),
            "search_stats": self.search.get_statistics() if self.search.search_log else {},
            "verifier_stats": self.verifier.get_statistics()
        }


if __name__ == "__main__":
    # Example usage
    print("=" * 60)
    print("PHASE 00: ENHANCED WITH AGENT FRAMEWORK")
    print("=" * 60)

    impl = Phase00Implementation()

    print(f"\nPhase {impl.phase_id}: {impl.phase_name}")
    print("Status: Ready for implementation")

    # Execute phase
    task = {
        "goal": "Set up foundation infrastructure with cloud, Kubernetes, and Neo4j",
        "requires_extensive_search": False
    }

    print(f"\nExecuting task: {task['goal']}")
    result = impl.execute(task)

    # Display results
    print("\n" + "=" * 60)
    print("EXECUTION RESULTS")
    print("=" * 60)
    print(f"Success: {result.success}")
    print(f"Iterations: {result.iterations}")
    print(f"Duration: {result.total_duration_seconds:.2f}s")

    if result.success:
        print("\n✅ Phase 00 completed successfully!")
        print("\nOutput:")
        import json
        print(json.dumps(result.output, indent=2))
    else:
        print("\n❌ Phase 00 failed")
        print(f"Error: {result.error}")

    # Show statistics
    print("\n" + "=" * 60)
    print("STATISTICS")
    print("=" * 60)
    stats = impl.get_statistics()
    print(json.dumps(stats, indent=2))

    # Cleanup
    impl.orchestrator.cleanup()
