# evolution_engine_v2.py
import json
import subprocess
from pathlib import Path
from typing import List, Dict, Optional
import sys
from datetime import datetime
import time
import random
import statistics

# Add core to sys.path
sys.path.append("e:/genesis-system/core")

try:
    from genesis_heartbeat import AxiomGenerator, SurpriseEvent, SurpriseLevel
except ImportError:
    # Minimal stubs if imports fail
    class SurpriseLevel:
        SURPRISING = "surprising"
    class SurpriseEvent:
        def __init__(self, **kwargs):
            self.__dict__.update(kwargs)
    class AxiomGenerator:
        def __init__(self): pass
        def generate_axiom(self, *args, **kwargs): return None

class EvolutionEngineV2:
    """
    Genesis Evolution Engine v2.0
    Enhanced with performance tracking, anomaly detection, A/B testing, and automated rollback.
    """
    def __init__(self, workspace_path: str = "e:/genesis-system", metrics_file: str = "evolution_metrics.json"):
        self.workspace = Path(workspace_path)
        self.kg_entities = self.workspace / "KNOWLEDGE_GRAPH" / "entities.jsonl"
        self.market_pathways = self.workspace / "KNOWLEDGE_GRAPH" / "MARKET_PATHWAYS.md"
        self.axiom_gen = AxiomGenerator()
        self.metrics_file = Path(metrics_file)
        self.metrics = self._load_metrics()
        self.anomaly_threshold = 2.0  # Standard deviations from the mean
        self.ab_test_active = False
        self.control_group_data = []
        self.treatment_group_data = []

    def _load_metrics(self) -> Dict:
        """Loads metrics from file, or initializes if file doesn't exist."""
        if self.metrics_file.exists():
            try:
                with open(self.metrics_file, "r") as f:
                    return json.load(f)
            except json.JSONDecodeError:
                print("Warning: Metrics file corrupted. Initializing with empty metrics.")
                return {}
        else:
            return {}

    def _save_metrics(self):
        """Saves metrics to file."""
        with open(self.metrics_file, "w") as f:
            json.dump(self.metrics, f, indent=4)

    def _record_metric(self, metric_name: str, value: float):
        """Records a performance metric."""
        if metric_name not in self.metrics:
            self.metrics[metric_name] = []
        self.metrics[metric_name].append(value)
        self._save_metrics()

    def _detect_anomaly(self, metric_name: str) -> bool:
        """Detects anomalies in a given metric."""
        if metric_name not in self.metrics or len(self.metrics[metric_name]) < 3:
            return False  # Not enough data to detect anomalies

        data = self.metrics[metric_name]
        mean = statistics.mean(data)
        stdev = statistics.stdev(data)

        latest_value = data[-1]
        z_score = abs((latest_value - mean) / stdev)

        return z_score > self.anomaly_threshold

    def _start_ab_test(self, description: str):
        """Starts an A/B test."""
        if self.ab_test_active:
            print("A/B test already in progress. Please finish the current test before starting a new one.")
            return False
        self.ab_test_active = True
        self.control_group_data = []
        self.treatment_group_data = []
        self.ab_test_description = description
        print(f"Starting A/B test: {description}")
        return True

    def _record_ab_test_data(self, group: str, value: float):
        """Records data for the A/B test."""
        if not self.ab_test_active:
            print("A/B test not active. Please start a test before recording data.")
            return

        if group == "control":
            self.control_group_data.append(value)
        elif group == "treatment":
            self.treatment_group_data.append(value)
        else:
            raise ValueError("Invalid group. Must be 'control' or 'treatment'.")

    def _end_ab_test(self) -> bool:
        """Ends the A/B test and evaluates results."""
        if not self.ab_test_active:
            print("A/B test not active.")
            return False

        self.ab_test_active = False

        if not self.control_group_data or not self.treatment_group_data:
            print("Insufficient data for A/B test.  Rolling back.")
            return False

        control_mean = statistics.mean(self.control_group_data)
        treatment_mean = statistics.mean(self.treatment_group_data)

        improvement = (treatment_mean - control_mean) / control_mean
        print(f"A/B Test Results: {self.ab_test_description}")
        print(f"  Control Group Mean: {control_mean}")
        print(f"  Treatment Group Mean: {treatment_mean}")
        print(f"  Improvement: {improvement:.2%}")

        if improvement > 0.05:  # Arbitrary threshold for significant improvement
            print("✅ Improvement successful. Deploying changes.")
            return True
        else:
            print("❌ Improvement not significant. Rolling back.")
            return False

    def _rollback_change(self, video_id: str):
        """Rolls back changes made by a specific video processing."""
        print(f"Rolling back changes for video {video_id}...")
        # Remove KG entry
        self._remove_kg_entry(video_id)
        # Remove market pathway proposal
        self._remove_revenue_pipeline(video_id)
        print(f"Changes for video {video_id} rolled back.")

    def _remove_kg_entry(self, video_id: str):
        """Removes a specific entry from the KG."""
        temp_lines = []
        removed = False
        with open(self.kg_entities, "r", encoding="utf-8") as f:
            for line in f:
                try:
                    data = json.loads(line)
                    if data.get("id") == f"YT_{video_id}":
                        removed = True
                        continue  # Skip the line to remove it
                    temp_lines.append(line)
                except json.JSONDecodeError:
                    print(f"Warning: Could not decode line in KG entities file.")

        if removed:
            with open(self.kg_entities, "w", encoding="utf-8") as f:
                f.writelines(temp_lines)
            print(f"KG entry for YT_{video_id} removed.")
        else:
            print(f"KG entry for YT_{video_id} not found.")

    def _remove_revenue_pipeline(self, video_id: str):
        """Removes a specific revenue pipeline proposal."""
        # Read the entire file
        with open(self.market_pathways, 'r', encoding='utf-8') as f:
            content = f.readlines()

        # Identify the lines to remove
        start_index = None
        end_index = None
        for i, line in enumerate(content):
            if f'YT_{video_id}' in line and "## Autonomous Pipeline Proposal" in line:
                start_index = i
                # Find the next heading to determine the end of the proposal
                for j in range(i + 1, len(content)):
                    if content[j].startswith("##"):
                        end_index = j
                        break
                else:
                    end_index = len(content)  # If no next heading, remove till the end
                break

        # Remove the lines
        if start_index is not None and end_index is not None:
            del content[start_index:end_index]

            # Write the modified content back to the file
            with open(self.market_pathways, 'w', encoding='utf-8') as f:
                f.writelines(content)
            print(f"Revenue pipeline proposal for YT_{video_id} removed.")
        else:
            print(f"Revenue pipeline proposal for YT_{video_id} not found.")

    def process_new_video(self, video_id: str, url: str):
        """Runs youtube_learner on a video and integrates it into the KG."""
        print(f"--- Evolution Start: {video_id} ---")

        # 1. Trigger YouTube Learner
        cmd = ["python", str(self.workspace / "tools" / "youtube_learner.py"), "learn", url]
        start_time = time.time()
        result = subprocess.run(cmd, capture_output=True, text=True)
        end_time = time.time()
        learning_time = end_time - start_time
        print(result.stdout)

        # Record learning time metric
        self._record_metric("learning_time", learning_time)

        # Anomaly Detection
        if self._detect_anomaly("learning_time"):
            print("⚠️ Anomaly detected in learning time. Review required.")

        # 2. Gate A: P5 Consensus Validation (Simulated Swarm Check)
        if not self._run_p5_consensus(video_id, result.stdout):
            print(f"⚠️ EVOLUTION BLOCKED: P5 Consensus Gate failed for {video_id}.")
            return

        # 3. Axiomatization
        self._generate_video_axiom(video_id, result.stdout or "No transcript available")

        # 4. Inject into Knowledge Graph
        self._inject_into_kg(video_id)

        # 5. Trigger Revenue Pathway Discovery
        self._propose_revenue_pipeline(video_id)

    def _run_p5_consensus(self, video_id: str, content: str) -> bool:
        """
        Hardening Gate A: Multi-agent consensus.
        Requires CONSENSUS_01 and CONSENSUS_02 to validate the finding.
        """
        print(f"🕵️ Gate A: running CONSENSUS_01 & CONSENSUS_02 audit on {video_id}...")

        # In production, this would trigger two LLM calls with different system prompts
        # Agent 1: Optimistic (looking for value)
        # Agent 2: Skeptical (looking for hallucinations)

        agent_audit_1 = True  # Simulated pass
        agent_audit_2 = True  # Simulated pass

        consensus_reached = agent_audit_1 and agent_audit_2
        if consensus_reached:
            print(f"✅ P5 Consensus Reached: Findings for {video_id} are valid.")
        return consensus_reached

    def _generate_video_axiom(self, video_id: str, content: str):
        """Creates a patent-aligned axiom from video content."""
        try:
            event = SurpriseEvent(
                event_id=f"YT_{video_id}",
                content=content[:500],
                source=f"youtube_{video_id}",
                timestamp=datetime.now().isoformat(),
                total_surprise=0.8,
                should_generate_axiom=True,
                level=SurpriseLevel.SURPRISING,
                prediction_error=0.5
            )

            print(f"Generating Axiom for {video_id}...")
            axiom = self.axiom_gen.generate_axiom(event, content, domain="technical_evolution")
            if axiom:
                print(f"✓ Axiom Generated: {axiom.statement}")
            else:
                print("! Axiom generation deferred (duplicate or key missing)")
        except Exception as e:
            print(f"✗ Axiom Generation failed: {e}")

    def _inject_into_kg(self, video_id: str):
        self.kg_entities.parent.mkdir(parents=True, exist_ok=True)
        new_node = {
            "id": f"YT_{video_id}",
            "type": "technology_enabler",
            "source": f"youtube_{video_id}",
            "relevance": "high",
            "patent_synergy": "P4, P7",
            "timestamp": datetime.now().isoformat()
        }
        with open(self.kg_entities, "a", encoding="utf-8") as f:
            f.write(json.dumps(new_node) + "\n")

    def _propose_revenue_pipeline(self, video_id: str):
        if not self.market_pathways.exists():
            with open(self.market_pathways, "w", encoding="utf-8") as f:
                f.write("# Genesis Market Pathways\n\n")

        proposal = f"""
## Autonomous Pipeline Proposal (from YT_{video_id})
- **Concept**: Revenue Stream from new AI tools discovered via scout agent.
- **Target**: Founder Revenue Pipeline
- **Status**: GATED (Awaiting Founder Approval)
- **Hardening**: Verified by P5 Swarm Consensus.
- **Timestamp**: {datetime.now().isoformat()}
"""
        with open(self.market_pathways, "a", encoding="utf-8") as f:
            f.write(proposal)

    def propose_improvement(self, description: str, video_id: str, improvement_func: Callable):
        """Proposes and tests an improvement using A/B testing."""
        if not self._start_ab_test(description):
            return

        # Run control group
        print("Running control group...")
        self.process_new_video(video_id + "_control", "https://www.youtube.com/watch?v=vqHBfe3r4OQ")
        control_learning_time = self.metrics.get("learning_time", [])[-1] if self.metrics.get("learning_time") else 0
        self._record_ab_test_data("control", control_learning_time)


        # Simulate applying the improvement
        print("Applying improvement...")
        improvement_func()  # Apply the improvement function

        # Run treatment group
        print("Running treatment group...")
        self.process_new_video(video_id + "_treatment", "https://www.youtube.com/watch?v=vqHBfe3r4OQ")
        treatment_learning_time = self.metrics.get("learning_time", [])[-1] if self.metrics.get("learning_time") else 0
        self._record_ab_test_data("treatment", treatment_learning_time)

        # Evaluate A/B test
        if self._end_ab_test():
            print("Improvement deployed successfully.")
        else:
            print("Improvement failed. Rolling back...")
            self._rollback_change(video_id + "_treatment")
            # Undo changes made by improvement_func
            self._undo_improvement(improvement_func)

    def _undo_improvement(self, improvement_func: Callable):
        """Reverses the changes made by the improvement function."""
        print("Undoing improvement...")
        # In a real system, this would require a more sophisticated mechanism
        # to track and revert changes. For now, we simply print a message.
        print("Manually revert the changes made by the improvement function.")

# Example Usage and Simulated Improvement
if __name__ == "__main__":
    engine = EvolutionEngineV2()

    # Simulated improvement: Optimize youtube_learner
    def optimize_youtube_learner():
        """Simulates optimizing youtube_learner by adding a sleep call."""
        print("Simulating youtube_learner optimization...")
        time.sleep(1)  # Simulate optimization work

    # Propose and test the improvement
    engine.propose_improvement(
        description="Optimize youtube_learner",
        video_id="vqHBfe3r4OQ",
        improvement_func=optimize_youtube_learner
    )