#!/usr/bin/env python3
"""
RALPH-ENHANCED AIVA ORCHESTRATOR
================================

Integrates Ralph Wiggum methodology for 100x productivity:
1. Atomic stories (one iteration each)
2. Self-verification (acceptance criteria)
3. Fresh context per task
4. 5 phases: Scout → Plan → Build → Validate → Handoff
5. passes: true/false tracking
6. agents.md context files

Based on: /mnt/e/genesis-system/GENESIS_RALPH_BREAKTHROUGH_SYNTHESIS.md
"""

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

# Load secrets
_secrets_path = Path("/mnt/e/genesis-system/config/secrets.env")
if _secrets_path.exists():
    with open(_secrets_path) as _f:
        for _line in _f:
            if "=" in _line and not _line.startswith("#"):
                _key, _value = _line.strip().split("=", 1)
                os.environ[_key] = _value.strip('"')

import redis
import google.generativeai as genai

# Configure Gemini
genai.configure(api_key=os.environ.get("GEMINI_API_KEY", ""))

# Logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(name)s: %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler('/mnt/e/genesis-system/data/ralph_orchestrator.log')
    ]
)
logger = logging.getLogger("RALPH-AIVA")

# Redis config
REDIS_CONFIG = {
    "host": os.environ.get("GENESIS_REDIS_HOST", ""),
    "port": int(os.environ.get("GENESIS_REDIS_PORT", 26379)),
    "password": os.environ.get("GENESIS_REDIS_PASSWORD", ""),
    "decode_responses": True,
}


@dataclass
class RalphStory:
    """An atomic user story following Ralph methodology."""
    id: str
    title: str
    description: str
    acceptance_criteria: List[str] = field(default_factory=list)
    files_affected: List[str] = field(default_factory=list)
    dependencies: List[str] = field(default_factory=list)
    passes: bool = False
    phase: str = "pending"  # pending, scout, plan, build, validate, handoff, complete
    output_path: str = ""
    verification_command: str = ""
    priority: int = 3
    created_at: str = ""
    completed_at: str = ""
    result: str = ""
    error: str = ""
    cost: float = 0.0
    iteration_count: int = 0
    max_iterations: int = 3

    def __post_init__(self):
        if not self.id:
            self.id = hashlib.md5(f"{self.title}{time.time()}".encode()).hexdigest()[:12]
        if not self.created_at:
            self.created_at = datetime.now().isoformat()

    def to_dict(self) -> Dict:
        return asdict(self)

    @classmethod
    def from_dict(cls, data: Dict) -> 'RalphStory':
        valid_fields = {f.name for f in cls.__dataclass_fields__.values()}
        filtered = {k: v for k, v in data.items() if k in valid_fields}
        return cls(**filtered)


class RalphGeminiExecutor:
    """Gemini executor with Ralph 5-phase methodology."""

    RALPH_SYSTEM_PROMPT = """You are a Genesis system worker operating under RALPH WIGGUM METHODOLOGY.

## CORE PRINCIPLES
1. **Atomic Execution**: Complete this ONE story in ONE iteration
2. **Self-Verification**: Test your own work against acceptance criteria
3. **Proof Required**: Every claim needs command output or test result
4. **Fresh Context**: You have full context for just this task

## 5-PHASE APPROACH
1. SCOUT (5%): Understand what exists, identify integration points
2. PLAN (10%): Design approach, identify files to modify
3. BUILD (60%): Implement with continuous testing
4. VALIDATE (20%): Prove completion with evidence
5. HANDOFF (5%): Document what was done

## OUTPUT FORMAT
Your response MUST include:

```python
# === IMPLEMENTATION START ===
[Complete, runnable Python code here]
# === IMPLEMENTATION END ===
```

Then verification:
```
## VERIFICATION
- Acceptance Criteria 1: [PASS/FAIL] - [Evidence]
- Acceptance Criteria 2: [PASS/FAIL] - [Evidence]

## FILES CREATED/MODIFIED
- /path/to/file.py - [Description of changes]

## NEXT STEPS (if any)
- [What the next story should address]
```

## QUALITY STANDARDS
- Include ALL imports at the top
- Add type hints to all functions
- Add docstrings explaining purpose
- Include error handling with logging
- Make code production-ready, not prototype
"""

    def __init__(self):
        self.model = genai.GenerativeModel("gemini-2.0-flash-exp")
        self.total_cost = 0.0
        self.total_tokens = 0

    async def execute_story(self, story: RalphStory) -> Dict:
        """Execute a single Ralph story through all phases."""
        start_time = time.time()

        # Build context from agents.md files
        context = self._gather_context(story)

        # Build the execution prompt
        prompt = f"""{self.RALPH_SYSTEM_PROMPT}

## CURRENT STORY
**Title**: {story.title}
**Description**: {story.description}

## ACCEPTANCE CRITERIA (You must verify each)
{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_affected) if story.files_affected else '- Determine based on task'}

## OUTPUT PATH
{story.output_path if story.output_path else 'Determine appropriate location in /mnt/e/genesis-system/'}

## EXISTING CONTEXT
{context}

## VERIFICATION COMMAND (if provided)
{story.verification_command or 'Self-verify against acceptance criteria'}

---

Execute this story now. Follow the 5-phase approach. Output complete, working code.
"""

        try:
            response = self.model.generate_content(prompt)

            # Track costs
            input_tokens = response.usage_metadata.prompt_token_count
            output_tokens = response.usage_metadata.candidates_token_count
            cost = (input_tokens * 0.10 + output_tokens * 0.40) / 1_000_000

            self.total_cost += cost
            self.total_tokens += input_tokens + output_tokens

            # Extract code from response
            result_text = response.text
            code = self._extract_code(result_text)

            # Determine output path
            output_path = story.output_path or self._infer_output_path(story, result_text)

            # Save code if extracted
            saved = False
            if code and output_path:
                saved = self._save_code(code, output_path)

            # Check verification status
            passes = self._check_verification(result_text)

            elapsed = time.time() - start_time

            logger.info(f"✅ Story '{story.title[:30]}' | Phase: BUILD→VALIDATE | Cost: ${cost:.4f} | Passes: {passes}")

            return {
                "success": True,
                "passes": passes,
                "result": result_text,
                "code": code,
                "output_path": output_path if saved else None,
                "cost": cost,
                "tokens": input_tokens + output_tokens,
                "elapsed": elapsed,
            }

        except Exception as e:
            logger.error(f"❌ Story failed: {story.title[:30]} - {e}")
            return {
                "success": False,
                "passes": False,
                "error": str(e),
                "cost": 0,
            }

    def _gather_context(self, story: RalphStory) -> str:
        """Gather context from agents.md files and related code."""
        context_parts = []

        # Read root agents.md if exists
        agents_md = Path("/mnt/e/genesis-system/agents.md")
        if agents_md.exists():
            context_parts.append(f"## ROOT CONTEXT (agents.md)\n{agents_md.read_text()[:2000]}")

        # Read folder-specific agents.md based on output path
        if story.output_path:
            folder = Path(story.output_path).parent
            folder_agents = folder / "agents.md"
            if folder_agents.exists():
                context_parts.append(f"## FOLDER CONTEXT ({folder_agents})\n{folder_agents.read_text()[:1000]}")

        # Read related files
        for file_path in story.files_affected[:3]:  # Limit to 3 files
            fp = Path(file_path)
            if fp.exists() and fp.suffix == '.py':
                content = fp.read_text()[:2000]
                context_parts.append(f"## EXISTING FILE: {file_path}\n```python\n{content}\n```")

        return "\n\n".join(context_parts) if context_parts else "No existing context files found."

    def _extract_code(self, text: str) -> Optional[str]:
        """Extract Python code from response."""
        # Look for implementation block
        if "# === IMPLEMENTATION START ===" in text:
            match = re.search(r'# === IMPLEMENTATION START ===\n(.+?)# === IMPLEMENTATION END ===', text, re.DOTALL)
            if match:
                return match.group(1).strip()

        # Fallback to ```python blocks
        if "```python" in text:
            match = re.search(r'```python\n(.+?)```', text, re.DOTALL)
            if match:
                return match.group(1).strip()

        return None

    def _infer_output_path(self, story: RalphStory, result: str) -> Optional[str]:
        """Infer output path from story or result."""
        # Check result for file path
        match = re.search(r'(/mnt/e/genesis-system/[^\s]+\.py)', result)
        if match:
            return match.group(1)

        # Infer from title
        title_slug = re.sub(r'[^a-z0-9]+', '_', story.title.lower())[:30]

        # Categorize based on keywords
        title_lower = story.title.lower()
        if any(w in title_lower for w in ['test', 'benchmark', 'chaos']):
            return f"/mnt/e/genesis-system/tests/{title_slug}.py"
        elif any(w in title_lower for w in ['knowledge', 'entity', 'axiom', 'pattern']):
            return f"/mnt/e/genesis-system/core/knowledge/{title_slug}.py"
        elif any(w in title_lower for w in ['skill', 'tool']):
            return f"/mnt/e/genesis-system/skills/{title_slug}.py"
        elif any(w in title_lower for w in ['api', 'webhook', 'integration']):
            return f"/mnt/e/genesis-system/core/integrations/{title_slug}.py"
        else:
            return f"/mnt/e/genesis-system/core/{title_slug}.py"

    def _save_code(self, code: str, path: str) -> bool:
        """Save code to file."""
        try:
            output_path = Path(path)
            output_path.parent.mkdir(parents=True, exist_ok=True)
            output_path.write_text(code)
            logger.info(f"📁 Saved: {path} ({len(code)} bytes)")
            return True
        except Exception as e:
            logger.error(f"Failed to save {path}: {e}")
            return False

    def _check_verification(self, result: str) -> bool:
        """Check if verification section indicates PASS."""
        result_lower = result.lower()

        # Look for verification section
        if "## verification" in result_lower:
            verif_section = result_lower.split("## verification")[1][:500]
            # Count passes vs fails
            passes = verif_section.count("pass")
            fails = verif_section.count("fail")
            return passes > 0 and fails == 0

        # Fallback: look for success indicators
        success_indicators = ["successfully", "completed", "✅", "pass"]
        return any(ind in result_lower for ind in success_indicators)


class RalphOrchestrator:
    """Orchestrator following Ralph methodology."""

    def __init__(self):
        self.redis = redis.Redis(**REDIS_CONFIG)
        self.executor = RalphGeminiExecutor()
        self.running = False

        # Redis keys
        self.queue_key = "genesis:ralph_queue"
        self.active_key = "genesis:ralph_active"
        self.completed_key = "genesis:ralph_completed"
        self.failed_key = "genesis:ralph_failed"

    def enqueue_story(self, story: RalphStory) -> str:
        """Add story to queue."""
        score = story.priority * 1000000 + time.time()
        self.redis.zadd(self.queue_key, {json.dumps(story.to_dict()): score})
        logger.info(f"📥 Enqueued: {story.id} - {story.title}")
        return story.id

    def dequeue_story(self) -> Optional[RalphStory]:
        """Get next story from queue."""
        result = self.redis.zpopmin(self.queue_key, count=1)
        if result:
            story_json, _ = result[0]
            story = RalphStory.from_dict(json.loads(story_json))
            self.redis.hset(self.active_key, story.id, json.dumps(story.to_dict()))
            return story
        return None

    def complete_story(self, story: RalphStory, result: Dict):
        """Mark story as complete."""
        story.passes = result.get("passes", False)
        story.phase = "complete"
        story.completed_at = datetime.now().isoformat()
        story.result = result.get("output_path", "completed")
        story.cost = result.get("cost", 0)

        self.redis.hdel(self.active_key, story.id)
        self.redis.hset(self.completed_key, story.id, json.dumps(story.to_dict()))

        logger.info(f"✅ Complete: {story.id} | Passes: {story.passes} | Cost: ${story.cost:.4f}")

    def fail_story(self, story: RalphStory, error: str):
        """Mark story as failed."""
        story.phase = "failed"
        story.error = error
        story.completed_at = datetime.now().isoformat()

        self.redis.hdel(self.active_key, story.id)
        self.redis.hset(self.failed_key, story.id, json.dumps(story.to_dict()))

        logger.warning(f"❌ Failed: {story.id} - {error[:50]}")

    async def run(self, max_stories: int = 100):
        """Main execution loop."""
        self.running = True
        processed = 0

        logger.info("=" * 60)
        logger.info("RALPH-ENHANCED ORCHESTRATOR STARTING")
        logger.info("=" * 60)

        while self.running and processed < max_stories:
            # Get next story
            story = self.dequeue_story()
            if not story:
                logger.debug("Queue empty, waiting...")
                await asyncio.sleep(10)
                continue

            # Execute story
            story.phase = "build"
            story.iteration_count += 1

            result = await self.executor.execute_story(story)

            if result.get("success"):
                self.complete_story(story, result)
            else:
                # Retry logic
                if story.iteration_count < story.max_iterations:
                    story.phase = "pending"
                    self.enqueue_story(story)
                    logger.info(f"🔄 Retry {story.iteration_count}/{story.max_iterations}: {story.id}")
                else:
                    self.fail_story(story, result.get("error", "Max iterations exceeded"))

            processed += 1
            await asyncio.sleep(2)  # Brief pause between stories

        # Summary
        logger.info("=" * 60)
        logger.info(f"Session complete: {processed} stories")
        logger.info(f"Total cost: ${self.executor.total_cost:.4f}")
        logger.info(f"Total tokens: {self.executor.total_tokens:,}")
        logger.info("=" * 60)

    def get_stats(self) -> Dict:
        """Get queue statistics."""
        return {
            "pending": self.redis.zcard(self.queue_key),
            "active": self.redis.hlen(self.active_key),
            "completed": self.redis.hlen(self.completed_key),
            "failed": self.redis.hlen(self.failed_key),
            "total_cost": round(self.executor.total_cost, 4),
            "total_tokens": self.executor.total_tokens,
        }


def convert_existing_tasks_to_ralph():
    """Convert existing AIVA tasks to Ralph format."""
    r = redis.Redis(**REDIS_CONFIG)
    orchestrator = RalphOrchestrator()

    # Get existing tasks
    existing = r.zrange("genesis:task_queue", 0, -1)
    converted = 0

    for task_json in existing:
        task = json.loads(task_json)

        # Convert to Ralph story
        story = RalphStory(
            id=task.get("id", ""),
            title=task.get("title", ""),
            description=task.get("description", ""),
            priority=task.get("priority", 3),
            acceptance_criteria=[
                f"Code compiles without errors",
                f"All functions have docstrings and type hints",
                f"Error handling is implemented",
                f"File is saved to appropriate location",
            ],
            files_affected=task.get("context_files", []),
        )

        orchestrator.enqueue_story(story)
        converted += 1

    # Clear old queue
    r.delete("genesis:task_queue")

    logger.info(f"✅ Converted {converted} tasks to Ralph stories")
    return converted


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

    parser = argparse.ArgumentParser(description="Ralph-Enhanced AIVA Orchestrator")
    parser.add_argument("command", choices=["run", "stats", "convert", "add"])
    parser.add_argument("--max", type=int, default=100, help="Max stories to process")
    parser.add_argument("--title", help="Story title (for add)")
    parser.add_argument("--description", help="Story description (for add)")

    args = parser.parse_args()

    if args.command == "run":
        orchestrator = RalphOrchestrator()
        await orchestrator.run(max_stories=args.max)

    elif args.command == "stats":
        orchestrator = RalphOrchestrator()
        stats = orchestrator.get_stats()
        print(json.dumps(stats, indent=2))

    elif args.command == "convert":
        count = convert_existing_tasks_to_ralph()
        print(f"Converted {count} tasks to Ralph stories")

    elif args.command == "add":
        if not args.title:
            print("Error: --title required")
            return

        story = RalphStory(
            id="",
            title=args.title,
            description=args.description or args.title,
            acceptance_criteria=[
                "Code compiles without errors",
                "Functions have docstrings and type hints",
                "Error handling implemented",
            ],
        )

        orchestrator = RalphOrchestrator()
        story_id = orchestrator.enqueue_story(story)
        print(f"Added story: {story_id}")


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