#!/usr/bin/env python3
"""
GEMINI RALPH ORCHESTRATOR - Queen Elevation Edition
Executes PRD user stories via Gemini Flash 2.0 with max 30 iterations per story.
Pure Gemini throughput - no escalation.
"""

import os
import re
import json
import time
import logging
import requests
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Any, Optional
from dataclasses import dataclass, asdict

# Configuration
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY", "AIzaSyCT_rx0NusUJWoqtT7uxHAKEfHo129SJb8")
MODEL = "gemini-2.0-flash-exp"
MAX_ITERATIONS = 30
OUTPUT_DIR = Path("/mnt/e/genesis-system/logs/queen_elevation")
PRD_PATH = Path("/mnt/e/genesis-system/RALPH WIGGUM/AIVA_QUEEN_ELEVATION_PRD.md")

# Gemini API endpoint
GEMINI_API_URL = f"https://generativelanguage.googleapis.com/v1beta/models/{MODEL}:generateContent"

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s | %(levelname)s | %(message)s',
    handlers=[
        logging.FileHandler(OUTPUT_DIR / "orchestrator.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)


@dataclass
class StoryResult:
    story_id: str
    title: str
    status: str  # completed, failed
    attempts: int
    duration_seconds: float
    output_file: Optional[str]
    error: Optional[str]
    learnings: List[str]


@dataclass
class RunStats:
    run_id: str
    start_time: str
    end_time: str
    total_stories: int
    completed: int
    failed: int
    total_attempts: int
    total_duration_seconds: float
    estimated_cost: float


class GeminiRalphOrchestrator:
    """Execute PRD stories via Gemini Flash 2.0 with fresh sessions."""

    def __init__(self, prd_path: Path = PRD_PATH):
        self.prd_path = prd_path
        self.output_dir = OUTPUT_DIR
        self.output_dir.mkdir(parents=True, exist_ok=True)

        self.stories: List[Dict[str, Any]] = []
        self.results: List[StoryResult] = []
        self.total_attempts = 0

        logger.info("=" * 60)
        logger.info("GEMINI RALPH ORCHESTRATOR - Queen Elevation")
        logger.info(f"Model: {MODEL}")
        logger.info(f"Max iterations per story: {MAX_ITERATIONS}")
        logger.info("=" * 60)

    def load_prd(self) -> int:
        """Load and parse PRD to extract user stories."""
        logger.info(f"Loading PRD from {self.prd_path}...")

        content = self.prd_path.read_text(encoding='utf-8')

        # Parse user stories using regex
        story_pattern = r'#### \*\*US-(\d+): ([^*]+)\*\*\s*\n\*\*Goal:\*\* ([^\n]+)\s*\n\n\*\*Acceptance Criteria:\*\*\s*\n((?:- \[[ x]\] [^\n]+\n)+)\s*\n\*\*Files to Create:\*\*\s*\n\s*\n((?:[^\n]+\n)+?)'

        matches = re.findall(story_pattern, content, re.MULTILINE)

        for match in matches:
            story_id = f"US-{match[0]}"
            title = match[1].strip()
            goal = match[2].strip()
            criteria_raw = match[3].strip()
            files_raw = match[4].strip()

            # Parse acceptance criteria
            criteria = [
                c.replace("- [ ] ", "").replace("- [x] ", "").strip()
                for c in criteria_raw.split("\n")
                if c.strip()
            ]

            # Parse files to create
            files = [f.strip() for f in files_raw.split("\n") if f.strip()]

            self.stories.append({
                "id": story_id,
                "title": title,
                "goal": goal,
                "acceptance_criteria": criteria,
                "files_to_create": files
            })

        logger.info(f"Loaded {len(self.stories)} user stories from PRD")
        return len(self.stories)

    def build_prompt(self, story: Dict[str, Any], learnings: List[str] = None) -> str:
        """Build fresh prompt for story execution."""
        prompt = f"""You are a Genesis system developer implementing user story {story['id']}.

## STORY: {story['id']} - {story['title']}

### Goal
{story['goal']}

### Acceptance Criteria
{chr(10).join(f"- {c}" for c in story['acceptance_criteria'])}

### Files to Create
{chr(10).join(f"- {f}" for f in story['files_to_create'])}

### Context
- Working directory: /mnt/e/genesis-system
- Python version: 3.12
- Database: PostgreSQL at postgresql-genesis-u50607.vm.elestio.app:25432
- Redis: redis-genesis-u50607.vm.elestio.app:26379
- Qdrant: qdrant-b3knu-u50607.vm.elestio.app:6333
- AIVA Ollama: localhost:23405 (internal)

"""
        if learnings:
            prompt += f"""### Previous Attempt Learnings
{chr(10).join(f"- {l}" for l in learnings)}

Consider these learnings and try a different approach.

"""

        prompt += """### Instructions
1. Generate COMPLETE, PRODUCTION-READY code for all files listed
2. Include all imports, error handling, logging, and type hints
3. Follow Python best practices and Genesis conventions
4. Make the code immediately executable
5. Include docstrings for all classes and functions

### Output Format
For each file, output:

# FILE: /mnt/e/genesis-system/path/to/file.py
<complete file content>


Begin implementation:"""

        return prompt

    def call_gemini(self, prompt: str) -> Dict[str, Any]:
        """Make fresh API call to Gemini Flash 2.0."""
        headers = {
            "Content-Type": "application/json"
        }

        payload = {
            "contents": [{
                "parts": [{
                    "text": prompt
                }]
            }],
            "generationConfig": {
                "temperature": 0.7,
                "maxOutputTokens": 8192,
                "topP": 0.95
            }
        }

        url = f"{GEMINI_API_URL}?key={GEMINI_API_KEY}"

        response = requests.post(url, headers=headers, json=payload, timeout=120)

        if response.status_code == 200:
            return response.json()
        else:
            raise Exception(f"Gemini API error {response.status_code}: {response.text[:500]}")

    def extract_code_blocks(self, response_text: str) -> List[Dict[str, str]]:
        """Extract file paths and code from response."""
        files = []

        # Pattern to match FILE: path followed by code
        pattern = r'# FILE: ([^\n]+)\n(.*?)(?=# FILE:|$)'

        # Also try markdown code blocks with triple backticks
        backticks = "```"
        md_pattern = backticks + r'\n# FILE: ([^\n]+)\n(.*?)' + backticks

        matches = re.findall(pattern, response_text, re.DOTALL)
        if not matches:
            matches = re.findall(md_pattern, response_text, re.DOTALL)

        for match in matches:
            file_path = match[0].strip()
            code = match[1].strip()

            # Clean up code - remove markdown code block markers
            if code.startswith(backticks):
                parts = code.split(backticks)
                code = parts[1] if len(parts) > 1 else code
                if code.startswith("python"):
                    code = code[6:]
            # Also handle case where backticks appear elsewhere
            code = code.replace(backticks + "python", "").replace(backticks, "")

            files.append({
                "path": file_path,
                "code": code.strip()
            })

        return files

    def validate_output(self, files: List[Dict[str, str]], story: Dict[str, Any]) -> tuple[bool, str]:
        """Validate that output meets acceptance criteria."""
        if not files:
            return False, "No files extracted from response"

        # Check all required files are present
        required_files = set(story['files_to_create'])
        generated_files = set(f['path'] for f in files)

        missing = required_files - generated_files
        if missing:
            return False, f"Missing files: {missing}"

        # Check code has minimum structure
        for f in files:
            if len(f['code']) < 50:
                return False, f"File {f['path']} is too short ({len(f['code'])} chars)"

            # SQL files just need CREATE/INSERT/etc, not def/class
            if f['path'].endswith('.sql'):
                if 'CREATE' not in f['code'].upper() and 'INSERT' not in f['code'].upper():
                    return False, f"File {f['path']} missing SQL statements"
            # HTML files need basic structure
            elif f['path'].endswith('.html'):
                if '<html' not in f['code'].lower() and '<!DOCTYPE' not in f['code'].upper():
                    return False, f"File {f['path']} missing HTML structure"
            # Python files need def/class
            elif f['path'].endswith('.py'):
                if 'def ' not in f['code'] and 'class ' not in f['code']:
                    return False, f"File {f['path']} missing function/class definitions"

        return True, "Validation passed"

    def save_output(self, story: Dict[str, Any], files: List[Dict[str, str]], attempt: int) -> str:
        """Save generated code to output directory."""
        story_dir = self.output_dir / story['id']
        story_dir.mkdir(parents=True, exist_ok=True)

        output_file = story_dir / f"output_attempt_{attempt}.md"

        content = f"""# {story['id']}: {story['title']}
**Generated:** {datetime.now().isoformat()}
**Attempt:** {attempt}
**Model:** {MODEL}

## Goal
{story['goal']}

## Acceptance Criteria
{chr(10).join(f"- {c}" for c in story['acceptance_criteria'])}

## Generated Files

"""
        for f in files:
            content += f"""### {f['path']}

{f['code']}


"""

        output_file.write_text(content, encoding='utf-8')

        # Also save individual Python files to their intended locations
        for f in files:
            target_path = Path(f['path'])
            target_path.parent.mkdir(parents=True, exist_ok=True)
            target_path.write_text(f['code'], encoding='utf-8')
            logger.info(f"  Wrote: {f['path']}")

        return str(output_file)

    def execute_story(self, story: Dict[str, Any]) -> StoryResult:
        """Execute single story with up to MAX_ITERATIONS attempts."""
        logger.info("")
        logger.info("=" * 60)
        logger.info(f"EXECUTING: {story['id']} - {story['title']}")
        logger.info("=" * 60)

        start_time = time.time()
        learnings = []

        for attempt in range(1, MAX_ITERATIONS + 1):
            self.total_attempts += 1
            logger.info(f"  Attempt {attempt}/{MAX_ITERATIONS}...")

            try:
                # Build fresh prompt
                prompt = self.build_prompt(story, learnings if learnings else None)

                # Call Gemini
                response = self.call_gemini(prompt)

                # Extract response text
                response_text = response['candidates'][0]['content']['parts'][0]['text']

                # Extract code blocks
                files = self.extract_code_blocks(response_text)

                # Validate output
                valid, message = self.validate_output(files, story)

                if valid:
                    # Save output
                    output_file = self.save_output(story, files, attempt)

                    duration = time.time() - start_time
                    logger.info(f"  SUCCESS after {attempt} attempts ({duration:.1f}s)")

                    return StoryResult(
                        story_id=story['id'],
                        title=story['title'],
                        status="completed",
                        attempts=attempt,
                        duration_seconds=duration,
                        output_file=output_file,
                        error=None,
                        learnings=learnings
                    )
                else:
                    learnings.append(f"Attempt {attempt}: {message}")
                    logger.info(f"    Validation failed: {message}")

            except Exception as e:
                learnings.append(f"Attempt {attempt}: {type(e).__name__}: {str(e)[:100]}")
                logger.error(f"    Error: {e}")

            # Small delay between attempts
            time.sleep(1)

        # All attempts exhausted
        duration = time.time() - start_time
        logger.error(f"  FAILED after {MAX_ITERATIONS} attempts ({duration:.1f}s)")

        return StoryResult(
            story_id=story['id'],
            title=story['title'],
            status="failed",
            attempts=MAX_ITERATIONS,
            duration_seconds=duration,
            output_file=None,
            error=f"Exhausted {MAX_ITERATIONS} attempts",
            learnings=learnings
        )

    def execute_range(self, start: int, end: int) -> RunStats:
        """Execute stories in specified range (1-indexed)."""
        run_id = f"run_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        start_time = datetime.now()

        # Filter stories
        stories_to_run = [s for s in self.stories if start <= int(s['id'].split('-')[1]) <= end]

        logger.info("")
        logger.info("=" * 60)
        logger.info(f"RALPH LOOP RUN: {run_id}")
        logger.info(f"Stories: US-{start:03d} to US-{end:03d} ({len(stories_to_run)} stories)")
        logger.info("=" * 60)

        for story in stories_to_run:
            result = self.execute_story(story)
            self.results.append(result)

        end_time = datetime.now()

        # Calculate stats
        completed = sum(1 for r in self.results if r.status == "completed")
        failed = sum(1 for r in self.results if r.status == "failed")
        total_duration = (end_time - start_time).total_seconds()
        estimated_cost = self.total_attempts * 0.01

        stats = RunStats(
            run_id=run_id,
            start_time=start_time.isoformat(),
            end_time=end_time.isoformat(),
            total_stories=len(stories_to_run),
            completed=completed,
            failed=failed,
            total_attempts=self.total_attempts,
            total_duration_seconds=total_duration,
            estimated_cost=estimated_cost
        )

        # Save run report
        self.save_run_report(stats)

        return stats

    def save_run_report(self, stats: RunStats):
        """Save run report to file."""
        report_file = self.output_dir / f"{stats.run_id}_report.md"

        content = f"""# Ralph Loop Run Report
**Run ID:** {stats.run_id}
**Start:** {stats.start_time}
**End:** {stats.end_time}

## Summary
| Metric | Value |
|--------|-------|
| Total Stories | {stats.total_stories} |
| Completed | {stats.completed} |
| Failed | {stats.failed} |
| Success Rate | {stats.completed / stats.total_stories * 100:.1f}% |
| Total Attempts | {stats.total_attempts} |
| Total Duration | {stats.total_duration_seconds:.1f}s |
| Estimated Cost | ${stats.estimated_cost:.2f} |

## Story Results

| Story | Status | Attempts | Duration |
|-------|--------|----------|----------|
"""

        for result in self.results:
            status_icon = "✅" if result.status == "completed" else "❌"
            content += f"| {result.story_id} | {status_icon} {result.status} | {result.attempts} | {result.duration_seconds:.1f}s |\n"

        content += f"""

## Failed Stories Details

"""
        for result in self.results:
            if result.status == "failed":
                content += f"""### {result.story_id}: {result.title}
**Error:** {result.error}
**Learnings:**
{chr(10).join(f"- {l}" for l in result.learnings[-5:])}

"""

        report_file.write_text(content, encoding='utf-8')
        logger.info(f"Run report saved: {report_file}")

        # Also save JSON for programmatic access
        json_file = self.output_dir / f"{stats.run_id}_results.json"
        json_data = {
            "stats": asdict(stats),
            "results": [asdict(r) for r in self.results]
        }
        json_file.write_text(json.dumps(json_data, indent=2), encoding='utf-8')

    def print_summary(self):
        """Print execution summary."""
        completed = sum(1 for r in self.results if r.status == "completed")
        failed = sum(1 for r in self.results if r.status == "failed")

        logger.info("")
        logger.info("=" * 60)
        logger.info("EXECUTION SUMMARY")
        logger.info("=" * 60)
        logger.info(f"Total Stories: {len(self.results)}")
        logger.info(f"Completed: {completed}")
        logger.info(f"Failed: {failed}")
        logger.info(f"Success Rate: {completed / len(self.results) * 100:.1f}%")
        logger.info(f"Total Attempts: {self.total_attempts}")
        logger.info(f"Estimated Cost: ${self.total_attempts * 0.01:.2f}")
        logger.info("=" * 60)


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

    parser = argparse.ArgumentParser(description="Gemini Ralph Orchestrator")
    parser.add_argument("--start", type=int, default=1, help="Start story number")
    parser.add_argument("--end", type=int, default=18, help="End story number")
    parser.add_argument("--prd", type=str, default=str(PRD_PATH), help="PRD file path")

    args = parser.parse_args()

    # Create output directory
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    # Initialize orchestrator
    orchestrator = GeminiRalphOrchestrator(Path(args.prd))

    # Load PRD
    orchestrator.load_prd()

    # Execute stories
    stats = orchestrator.execute_range(args.start, args.end)

    # Print summary
    orchestrator.print_summary()

    return 0 if stats.failed == 0 else 1


if __name__ == "__main__":
    exit(main())