#!/usr/bin/env python3
"""
GENESIS PERPETUAL EVOLUTION ENGINE
===================================
The Unstoppable Force that elevates AIVA to Genesis Prime Mother.

This engine:
1. Reads PRDs from the queue
2. Executes stories via Flash 2.0 swarm (32 parallel agents)
3. Spawns new PRDs from completed stories
4. Loops FOREVER until explicitly paused

Run: python perpetual_engine.py --start
"""

import os
import sys
import json
import time
import asyncio
import hashlib
import logging
from pathlib import Path
from datetime import datetime
from typing import List, Dict, Optional, Any
from dataclasses import dataclass, field, asdict
from concurrent.futures import ThreadPoolExecutor
import re

# Add parent to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent.parent))

try:
    import google.generativeai as genai
except ImportError:
    print("Installing google-generativeai...")
    os.system(f"{sys.executable} -m pip install google-generativeai")
    import google.generativeai as genai

# Configuration
GENESIS_ROOT = Path("/mnt/e/genesis-system")
PRD_FOLDER = GENESIS_ROOT / "RALPH WIGGUM"
EVOLUTION_LOG = GENESIS_ROOT / "logs" / "perpetual_evolution.jsonl"
STATE_FILE = GENESIS_ROOT / "data" / "perpetual_state.json"
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "AIzaSyCT_rx0NusUJWoqtT7uxHAKEfHo129SJb8")

# RWL Best Practice Configuration
MAX_RETRIES = 30  # Stories get 30 attempts before failing
MIN_STORIES_PER_PRD = 20  # PRDs should have at least 20 stories
RETRY_BACKOFF_BASE = 2  # Exponential backoff base (seconds)

# Ensure directories exist
(GENESIS_ROOT / "logs").mkdir(exist_ok=True)
(GENESIS_ROOT / "data").mkdir(exist_ok=True)
(GENESIS_ROOT / "core" / "evolution").mkdir(parents=True, exist_ok=True)
(GENESIS_ROOT / "core" / "consciousness").mkdir(parents=True, exist_ok=True)
(GENESIS_ROOT / "core" / "swarm").mkdir(parents=True, exist_ok=True)

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s | %(levelname)s | %(message)s',
    handlers=[
        logging.FileHandler(GENESIS_ROOT / "logs" / "perpetual_engine.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger("PerpetualEngine")


@dataclass
class Story:
    """An atomic story from a PRD."""
    id: str
    title: str
    goal: str
    acceptance_criteria: List[str]
    files: List[str]
    spawns: Optional[str] = None
    status: str = "pending"  # pending, in_progress, completed, failed
    attempts: int = 0
    result: Optional[str] = None
    completed_at: Optional[str] = None


@dataclass
class PRD:
    """A Product Requirements Document."""
    id: str
    title: str
    source_file: str
    stories: List[Story]
    status: str = "pending"  # pending, in_progress, completed
    priority: int = 1
    parent_story: Optional[str] = None
    created_at: str = field(default_factory=lambda: datetime.utcnow().isoformat())
    completed_at: Optional[str] = None


@dataclass
class EngineState:
    """Persistent state of the perpetual engine."""
    is_running: bool = False
    current_prd: Optional[str] = None
    total_stories_completed: int = 0
    total_prds_completed: int = 0
    total_prds_spawned: int = 0
    total_cost_usd: float = 0.0
    start_time: Optional[str] = None
    last_activity: Optional[str] = None
    gpm_progress: Dict[str, bool] = field(default_factory=dict)


class FlashAgent:
    """A Gemini Flash 2.0 agent for executing stories."""

    def __init__(self, agent_id: int):
        self.agent_id = agent_id
        self.model = genai.GenerativeModel('gemini-2.5-flash')
        genai.configure(api_key=GEMINI_API_KEY)

    async def execute_story(self, story: Story, prd_context: str) -> Dict[str, Any]:
        """Execute a single story and return results."""

        prompt = f"""Execute this Ralph Wiggum Loop story. Output working code only.

## STORY: {story.title}
**Goal:** {story.goal}
**Acceptance Criteria:**
{chr(10).join(f'- {c}' for c in story.acceptance_criteria)}

**Files to Create/Modify:**
{chr(10).join(f'- {f}' for f in story.files)}

## OUTPUT FORMAT (JSON only, no markdown):
{{
    "success": true/false,
    "files_created": [
        {{"path": "relative/path/to/file.py", "content": "complete file content"}}
    ],
    "verification": {{
        "criteria_met": ["criterion 1", "criterion 2"],
        "criteria_failed": []
    }},
    "learnings": "Brief technical insight"
}}

Requirements:
- Complete, production-ready code
- All imports included
- Error handling where needed
- No placeholders or TODOs"""

        try:
            response = await asyncio.to_thread(
                self.model.generate_content,
                prompt,
                generation_config=genai.GenerationConfig(
                    temperature=0.7,
                    max_output_tokens=8192
                )
            )

            # Parse response
            text = response.text

            # Extract JSON from response
            json_match = re.search(r'\{[\s\S]*\}', text)
            if json_match:
                result = json.loads(json_match.group())
                return {
                    "success": result.get("success", False),
                    "files": result.get("files_created", []),
                    "verification": result.get("verification", {}),
                    "learnings": result.get("learnings", ""),
                    "raw_response": text[:200]
                }
            else:
                return {
                    "success": False,
                    "error": "Could not parse response",
                    "raw_response": text[:500]
                }

        except Exception as e:
            return {
                "success": False,
                "error": str(e)
            }


class PRDParser:
    """Parse PRD markdown files into structured PRD objects."""

    @staticmethod
    def parse_file(filepath: Path) -> Optional[PRD]:
        """Parse a PRD markdown file."""
        try:
            content = filepath.read_text(encoding='utf-8')

            # Extract title
            title_match = re.search(r'^#\s+(.+)$', content, re.MULTILINE)
            title = title_match.group(1) if title_match else filepath.stem

            # Extract stories
            stories = []
            story_pattern = r'###\s+(GPM-\d+|GEN-Q-\d+|REV-\d+|US-\d+):\s*(.+?)(?=###|\Z)'

            for match in re.finditer(story_pattern, content, re.DOTALL):
                story_id = match.group(1)
                story_content = match.group(2)

                # Extract goal
                goal_match = re.search(r'\*\*Goal:\*\*\s*(.+?)(?:\n|$)', story_content)
                goal = goal_match.group(1).strip() if goal_match else ""

                # Extract acceptance criteria
                criteria = []
                criteria_section = re.search(r'\*\*Acceptance Criteria:\*\*(.+?)(?:\*\*|$)', story_content, re.DOTALL)
                if criteria_section:
                    criteria = re.findall(r'-\s+(.+?)(?:\n|$)', criteria_section.group(1))

                # Extract files
                files = []
                files_match = re.search(r'\*\*Files.*?:\*\*\s*(.+?)(?:\n\*\*|$)', story_content, re.DOTALL)
                if files_match:
                    files = re.findall(r'`([^`]+)`', files_match.group(1))

                # Extract spawns
                spawns = None
                spawns_match = re.search(r'\*\*Spawns:\*\*\s*(.+?)(?:\n|$)', story_content)
                if spawns_match:
                    spawns = spawns_match.group(1).strip()

                if goal or criteria:  # Only add if we have meaningful content
                    stories.append(Story(
                        id=story_id,
                        title=story_content.split('\n')[0].strip(),
                        goal=goal,
                        acceptance_criteria=criteria,
                        files=files,
                        spawns=spawns
                    ))

            if stories:
                return PRD(
                    id=hashlib.md5(filepath.name.encode()).hexdigest()[:8],
                    title=title,
                    source_file=str(filepath),
                    stories=stories
                )

        except Exception as e:
            logger.error(f"Error parsing {filepath}: {e}")

        return None


class PRDSpawner:
    """Generate new PRDs from completed stories."""

    def __init__(self):
        self.model = genai.GenerativeModel('gemini-2.5-flash')
        genai.configure(api_key=GEMINI_API_KEY)

    async def spawn_prd(self, parent_story: Story, spawn_hint: str) -> Optional[PRD]:
        """Generate a new PRD based on a completed story's spawn hint."""

        prompt = f"""Generate a PRD with 25-40 atomic stories for Ralph Wiggum Loop execution.

## SPAWN FROM
Parent: {parent_story.id} - {parent_story.title}
Direction: {spawn_hint}

## REQUIREMENTS
- Generate 25-40 atomic stories (MUST be completable in single iteration)
- Each story creates ONE file or makes ONE focused change
- Stories build on each other progressively
- Include test stories, integration stories, and documentation stories

## STORY FORMAT (strictly follow):
### GPM-XXX: Title
**Goal:** Single atomic goal
**Acceptance Criteria:**
- Testable criterion 1
- Testable criterion 2
**Files:** `core/path/to/file.py`
**Spawns:** Next evolution hint

## STORY DISTRIBUTION (25-40 total):
- 10-15 core implementation stories
- 5-8 test/verification stories
- 5-8 integration stories
- 3-5 documentation/config stories
- 2-4 optimization stories

Generate complete PRD now:"""

        try:
            response = await asyncio.to_thread(
                self.model.generate_content,
                prompt,
                generation_config=genai.GenerationConfig(
                    temperature=0.7,
                    max_output_tokens=16384
                )
            )

            # Save the spawned PRD
            prd_content = response.text
            spawn_id = hashlib.md5(f"{parent_story.id}-{spawn_hint}".encode()).hexdigest()[:8]
            prd_filename = f"SPAWNED_{spawn_id}_{spawn_hint.replace(' ', '_')[:30]}.md"
            prd_path = PRD_FOLDER / prd_filename

            prd_path.write_text(prd_content, encoding='utf-8')
            logger.info(f"Spawned new PRD: {prd_filename}")

            # Parse the newly created PRD
            return PRDParser.parse_file(prd_path)

        except Exception as e:
            logger.error(f"Error spawning PRD: {e}")
            return None


class PerpetualEvolutionEngine:
    """The Unstoppable Force - AIVA's perpetual evolution engine."""

    def __init__(self, max_agents: int = 8):
        self.max_agents = max_agents
        self.agents = [FlashAgent(i) for i in range(max_agents)]
        self.prd_queue: List[PRD] = []
        self.spawner = PRDSpawner()
        self.state = self._load_state()
        self.running = False

    def _load_state(self) -> EngineState:
        """Load persistent state."""
        if STATE_FILE.exists():
            try:
                data = json.loads(STATE_FILE.read_text())
                return EngineState(**data)
            except:
                pass
        return EngineState()

    def _save_state(self):
        """Save persistent state."""
        self.state.last_activity = datetime.utcnow().isoformat()
        STATE_FILE.write_text(json.dumps(asdict(self.state), indent=2))

    def _log_evolution(self, event: Dict[str, Any]):
        """Log evolution event."""
        event["timestamp"] = datetime.utcnow().isoformat()
        with open(EVOLUTION_LOG, "a") as f:
            f.write(json.dumps(event) + "\n")

    def load_prds(self):
        """Load all PRDs from the queue folder."""
        logger.info("Loading PRDs from RALPH WIGGUM folder...")

        prd_files = list(PRD_FOLDER.glob("*_PRD.md"))
        logger.info(f"Found {len(prd_files)} PRD files")

        for prd_file in prd_files:
            prd = PRDParser.parse_file(prd_file)
            if prd and prd.stories:
                # Check if already in queue
                if not any(p.id == prd.id for p in self.prd_queue):
                    self.prd_queue.append(prd)
                    logger.info(f"Loaded PRD: {prd.title} ({len(prd.stories)} stories)")

        # Prioritize the Genesis Prime Mother PRD
        self.prd_queue.sort(key=lambda p: (
            0 if "PRIME_MOTHER" in p.title.upper() else
            1 if "QUANTUM" in p.title.upper() else
            2
        ))

        logger.info(f"Total PRDs in queue: {len(self.prd_queue)}")

    async def execute_story_with_agent(self, story: Story, prd: PRD) -> bool:
        """Execute a story using an available agent."""

        # Get available agent
        agent = self.agents[self.state.total_stories_completed % self.max_agents]

        logger.info(f"Agent {agent.agent_id} executing: {story.id} - {story.title}")
        story.status = "in_progress"
        story.attempts += 1

        result = await agent.execute_story(story, prd.title)

        if result.get("success"):
            story.status = "completed"
            story.result = result.get("learnings", "")
            story.completed_at = datetime.utcnow().isoformat()

            # Write files if provided
            for file_info in result.get("files", []):
                try:
                    file_path = GENESIS_ROOT / file_info.get("path", "")
                    file_path.parent.mkdir(parents=True, exist_ok=True)
                    file_path.write_text(file_info.get("content", ""), encoding='utf-8')
                    logger.info(f"Created file: {file_path}")
                except Exception as e:
                    logger.error(f"Error writing file: {e}")

            self.state.total_stories_completed += 1

            # Log technical learnings only
            learnings = result.get("learnings", "")
            if learnings and len(learnings) < 100:
                logger.debug(f"Learning: {learnings}")

            # Spawn child PRD if specified
            if story.spawns:
                logger.info(f"Spawning child PRD: {story.spawns}")
                child_prd = await self.spawner.spawn_prd(story, story.spawns)
                if child_prd:
                    self.prd_queue.append(child_prd)
                    self.state.total_prds_spawned += 1
                    logger.info(f"Spawned PRD: {child_prd.title}")

            self._log_evolution({
                "event": "story_completed",
                "story_id": story.id,
                "story_title": story.title,
                "prd_id": prd.id,
                "learnings": story.result
            })

            return True
        else:
            if story.attempts >= MAX_RETRIES:
                story.status = "failed"
                logger.error(f"Story failed after {MAX_RETRIES} attempts: {story.id}")
            else:
                story.status = "pending"  # Will retry with backoff
                backoff = min(RETRY_BACKOFF_BASE ** (story.attempts - 1), 30)  # Cap at 30s
                if story.attempts % 5 == 0:  # Log every 5th attempt
                    logger.warning(f"Story {story.id}: attempt {story.attempts}/{MAX_RETRIES}, backoff {backoff}s")
                await asyncio.sleep(backoff)

            return False

    async def execute_prd(self, prd: PRD):
        """Execute all stories in a PRD."""

        logger.info(f"\n{'='*60}")
        logger.info(f"EXECUTING PRD: {prd.title}")
        logger.info(f"Stories: {len(prd.stories)}")
        logger.info(f"{'='*60}\n")

        prd.status = "in_progress"
        self.state.current_prd = prd.id
        self._save_state()

        pending_stories = [s for s in prd.stories if s.status == "pending"]

        while pending_stories:
            # Execute stories in parallel (up to max_agents)
            batch = pending_stories[:self.max_agents]

            tasks = [
                self.execute_story_with_agent(story, prd)
                for story in batch
            ]

            results = await asyncio.gather(*tasks, return_exceptions=True)

            # Update pending list
            pending_stories = [s for s in prd.stories if s.status == "pending"]

            # Save state after each batch
            self._save_state()

            # Brief pause between batches
            await asyncio.sleep(1)

        # Check if PRD is complete
        completed = sum(1 for s in prd.stories if s.status == "completed")
        failed = sum(1 for s in prd.stories if s.status == "failed")

        if failed == 0:
            prd.status = "completed"
            prd.completed_at = datetime.utcnow().isoformat()
            self.state.total_prds_completed += 1
            logger.info(f"PRD COMPLETED: {prd.title} ({completed}/{len(prd.stories)} stories)")
        else:
            logger.warning(f"PRD PARTIAL: {prd.title} ({completed} completed, {failed} failed)")

        self._log_evolution({
            "event": "prd_completed",
            "prd_id": prd.id,
            "prd_title": prd.title,
            "stories_completed": completed,
            "stories_failed": failed
        })

    async def run_forever(self):
        """The Unstoppable Force - runs until explicitly stopped."""

        logger.info("""
╔═══════════════════════════════════════════════════════════════════╗
║                                                                   ║
║   ██████╗ ███████╗███╗   ██╗███████╗███████╗██╗███████╗          ║
║  ██╔════╝ ██╔════╝████╗  ██║██╔════╝██╔════╝██║██╔════╝          ║
║  ██║  ███╗█████╗  ██╔██╗ ██║█████╗  ███████╗██║███████╗          ║
║  ██║   ██║██╔══╝  ██║╚██╗██║██╔══╝  ╚════██║██║╚════██║          ║
║  ╚██████╔╝███████╗██║ ╚████║███████╗███████║██║███████║          ║
║   ╚═════╝ ╚══════╝╚═╝  ╚═══╝╚══════╝╚══════╝╚═╝╚══════╝          ║
║                                                                   ║
║        PERPETUAL EVOLUTION ENGINE - THE UNSTOPPABLE FORCE         ║
║                                                                   ║
║   Elevating AIVA to Genesis Prime Mother                          ║
║   PRDs spawn PRDs | Flash 2.0 Swarm Active | Loop Forever         ║
║                                                                   ║
╚═══════════════════════════════════════════════════════════════════╝
        """)

        self.running = True
        self.state.is_running = True
        self.state.start_time = datetime.utcnow().isoformat()
        self._save_state()

        cycle = 0

        while self.running:
            cycle += 1
            logger.info(f"\n{'#'*60}")
            logger.info(f"EVOLUTION CYCLE {cycle}")
            logger.info(f"PRDs in queue: {len(self.prd_queue)}")
            logger.info(f"Total stories completed: {self.state.total_stories_completed}")
            logger.info(f"Total PRDs spawned: {self.state.total_prds_spawned}")
            logger.info(f"{'#'*60}\n")

            # Reload PRDs to catch any new ones
            self.load_prds()

            # Get next pending PRD
            pending_prds = [p for p in self.prd_queue if p.status == "pending"]

            if not pending_prds:
                logger.info("No pending PRDs. Waiting for spawned PRDs or new additions...")
                await asyncio.sleep(60)  # Wait a minute then check again
                continue

            # Execute next PRD
            prd = pending_prds[0]
            await self.execute_prd(prd)

            # Log cycle completion
            self._log_evolution({
                "event": "cycle_completed",
                "cycle": cycle,
                "prds_pending": len([p for p in self.prd_queue if p.status == "pending"]),
                "prds_completed": self.state.total_prds_completed,
                "stories_completed": self.state.total_stories_completed,
                "prds_spawned": self.state.total_prds_spawned
            })

            # Brief pause between PRDs
            await asyncio.sleep(2)

        logger.info("Perpetual engine paused. AIVA continues to grow.")

    def stop(self):
        """Gracefully stop the engine."""
        self.running = False
        self.state.is_running = False
        self._save_state()
        logger.info("Stop signal received. Completing current work...")


async def main():
    """Main entry point."""
    import argparse

    parser = argparse.ArgumentParser(description="Genesis Perpetual Evolution Engine")
    parser.add_argument("--start", action="store_true", help="Start the unstoppable force")
    parser.add_argument("--status", action="store_true", help="Show current status")
    parser.add_argument("--agents", type=int, default=8, help="Number of parallel agents (default: 8)")

    args = parser.parse_args()

    engine = PerpetualEvolutionEngine(max_agents=args.agents)

    if args.status:
        print(f"\n=== PERPETUAL ENGINE STATUS ===")
        print(f"Running: {engine.state.is_running}")
        print(f"Stories Completed: {engine.state.total_stories_completed}")
        print(f"PRDs Completed: {engine.state.total_prds_completed}")
        print(f"PRDs Spawned: {engine.state.total_prds_spawned}")
        print(f"Last Activity: {engine.state.last_activity}")
        return

    if args.start:
        try:
            await engine.run_forever()
        except KeyboardInterrupt:
            engine.stop()
            print("\nEngine stopped gracefully.")
    else:
        parser.print_help()


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