#!/usr/bin/env python3
"""
Universal Ralph Wiggum Loop Runner
Executes any PRD via Gemini Flash 2.0
Fresh API context per story - ZERO context rot
"""

import os
import json
import re
import time
import sys
from datetime import datetime
from pathlib import Path
from dataclasses import dataclass, field
from typing import List, Dict, Optional
import google.generativeai as genai

# Configuration
MAX_ITERATIONS = 30
MODEL_NAME = "gemini-2.0-flash"

@dataclass
class StoryResult:
    story_id: str
    title: str
    status: str
    attempts: int
    duration: float
    files_created: List[str] = field(default_factory=list)
    learnings: str = ""
    error: str = ""

def load_api_key():
    """Load Gemini API key"""
    key = os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")
    if not key:
        # Search paths for both Linux (WSL) and Windows
        search_paths = [
            Path("/mnt/e/genesis-system/.env"), 
            Path("/mnt/e/genesis-system/config/secrets.env"),
            Path("e:/genesis-system/.env"),
            Path("e:/genesis-system/config/secrets.env"),
             Path(".env")
        ]
        for env_path in search_paths:
            if env_path.exists():
                with open(env_path) as f:
                    for line in f:
                        if "GEMINI_API_KEY=" in line or "GOOGLE_API_KEY=" in line:
                            parts = line.split("=", 1)
                            if len(parts) > 1:
                                key = parts[1].strip()
                                if key:
                                    return key
    return key

def parse_prd(prd_path: Path) -> List[Dict]:
    """Parse PRD to extract user stories - flexible pattern matching"""
    content = prd_path.read_text()
    stories = []

    # More flexible pattern
    story_blocks = re.split(r'\n### US-', content)

    for block in story_blocks[1:]:  # Skip first split (before first US-)
        lines = block.strip().split('\n')
        if not lines:
            continue

        # Parse story ID and title from first line
        first_line = lines[0]
        match = re.match(r'(\d+): (.+)', first_line)
        if not match:
            continue

        story_id = f"US-{match.group(1)}"
        title = match.group(2).strip()

        # Extract sections
        goal = ""
        acceptance = ""
        files = []
        deps = ""
        hints = ""

        current_section = None
        section_content = []

        for line in lines[1:]:
            if line.startswith('**Goal:**'):
                goal = line.replace('**Goal:**', '').strip()
            elif line.startswith('**Acceptance Criteria:**'):
                current_section = 'acceptance'
            elif line.startswith('**Files to Create:**'):
                current_section = 'files'
            elif line.startswith('**Dependencies:**'):
                deps = line.replace('**Dependencies:**', '').strip()
                current_section = None
            elif line.startswith('**Technical Hints:**'):
                current_section = 'hints'
            elif current_section == 'acceptance' and line.strip() and not line.startswith('```'):
                acceptance += line + "\n"
            elif current_section == 'files' and line.strip().startswith('-'):
                files.append(line.strip().replace('- ', '').strip('`'))
            elif current_section == 'hints' and line.strip():
                hints += line + "\n"

        if story_id and title:
            stories.append({
                "id": story_id,
                "title": title,
                "goal": goal,
                "acceptance_criteria": acceptance.strip(),
                "files": files,
                "dependencies": deps,
                "hints": hints.strip()
            })

    return stories

def build_prompt(story: Dict, learnings: List[str]) -> str:
    """Build prompt for Gemini"""
    learnings_text = "\n".join(f"- {l}" for l in learnings[-5:]) if learnings else "None"

    return f"""You are an expert developer creating files for a user story.

## TASK
Create the files for: {story['id']}: {story['title']}

## GOAL
{story['goal']}

## ACCEPTANCE CRITERIA
{story['acceptance_criteria']}

## FILES TO CREATE
{chr(10).join('- ' + f for f in story['files'])}

## TECHNICAL HINTS
{story['hints']}

## PREVIOUS LEARNINGS
{learnings_text}

## OUTPUT FORMAT
For each file, output in this exact format:

### FILE: /exact/path/to/file.ext
```
[file contents here - use appropriate language/format]
```

## REQUIREMENTS
1. Create ALL files listed above
2. Use proper formatting for each file type
3. Include comprehensive content
4. Make files production-ready

BEGIN CREATING FILES:"""

def extract_files(response_text: str) -> Dict[str, str]:
    """Extract file contents from response"""
    files = {}
    pattern = r'### FILE: (.+?)\n```(?:\w+)?\n(.*?)```'
    matches = re.findall(pattern, response_text, re.DOTALL)
    for path, content in matches:
        files[path.strip()] = content.strip()
    return files

def write_files(files: Dict[str, str]) -> List[str]:
    """Write files to disk"""
    written = []
    for path, content in files.items():
        try:
            file_path = Path(path)
            file_path.parent.mkdir(parents=True, exist_ok=True)
            file_path.write_text(content, encoding='utf-8')
            written.append(str(file_path))
            print(f"  ✅ {path}")
        except Exception as e:
            print(f"  ❌ {path}: {e}")
    return written

def execute_story(story: Dict, model, learnings: List[str]) -> StoryResult:
    """Execute single story with fresh context"""
    print(f"\n{'='*60}")
    print(f"📋 {story['id']}: {story['title']}")
    print(f"{'='*60}")

    start_time = time.time()

    for attempt in range(1, MAX_ITERATIONS + 1):
        print(f"  Attempt {attempt}/{MAX_ITERATIONS}...")

        try:
            prompt = build_prompt(story, learnings)
            response = model.generate_content(prompt)
            files = extract_files(response.text)

            if not files:
                learnings.append(f"{story['id']} attempt {attempt}: No files extracted")
                time.sleep(2)
                continue

            files_created = write_files(files)

            if len(files_created) >= len(story['files']) * 0.5:  # 50% threshold
                return StoryResult(
                    story_id=story['id'],
                    title=story['title'],
                    status="success",
                    attempts=attempt,
                    duration=time.time() - start_time,
                    files_created=files_created,
                    learnings=f"Completed in {attempt} attempt(s)"
                )

            learnings.append(f"{story['id']} attempt {attempt}: Only {len(files_created)}/{len(story['files'])} files")

        except Exception as e:
            error_msg = str(e)
            if "429" in error_msg:
                wait_time = 60
                print(f"  ⏳ Rate limited, waiting {wait_time}s...")
                time.sleep(wait_time)
            else:
                learnings.append(f"{story['id']} attempt {attempt}: {error_msg[:100]}")
                time.sleep(2)

    return StoryResult(
        story_id=story['id'],
        title=story['title'],
        status="failed",
        attempts=MAX_ITERATIONS,
        duration=time.time() - start_time,
        error=f"Failed after {MAX_ITERATIONS} attempts"
    )

def generate_report(prd_name: str, results: List[StoryResult], duration: float) -> str:
    """Generate execution report"""
    success = sum(1 for r in results if r.status == "success")
    report = f"""# Ralph Wiggum Loop Report: {prd_name}
**Generated:** {datetime.now().isoformat()}
**Model:** {MODEL_NAME}

## Summary
| Metric | Value |
|--------|-------|
| Stories | {len(results)} |
| Successful | {success} |
| Failed | {len(results) - success} |
| Success Rate | {success/len(results)*100:.1f}% |
| Duration | {duration:.1f}s |

## Results
"""
    for r in results:
        emoji = "✅" if r.status == "success" else "❌"
        report += f"\n### {emoji} {r.story_id}: {r.title}\n"
        report += f"- Status: {r.status}\n- Attempts: {r.attempts}\n- Duration: {r.duration:.1f}s\n"
        if r.files_created:
            report += f"- Files: {len(r.files_created)}\n"
    return report

def run_prd(prd_path: Path, output_dir: Path):
    """Execute a PRD via Ralph Loop"""
    print(f"\n🚀 RALPH WIGGUM LOOP: {prd_path.name}")
    print("=" * 60)

    api_key = load_api_key()
    if not api_key:
        print("❌ No Gemini API key found!")
        return

    genai.configure(api_key=api_key)
    model = genai.GenerativeModel(MODEL_NAME)

    stories = parse_prd(prd_path)
    print(f"📄 Found {len(stories)} stories")

    if not stories:
        print("❌ No stories parsed!")
        return

    output_dir.mkdir(parents=True, exist_ok=True)
    results = []
    learnings = []
    start = time.time()

    for story in stories:
        result = execute_story(story, model, learnings)
        results.append(result)
        if result.learnings:
            learnings.append(result.learnings)

    duration = time.time() - start

    report = generate_report(prd_path.stem, results, duration)
    report_path = output_dir / f"{prd_path.stem}_REPORT.md"
    report_path.write_text(report, encoding='utf-8')

    success = sum(1 for r in results if r.status == "success")
    print(f"\n🎯 FINAL: {success}/{len(results)} successful")
    print(f"📄 Report: {report_path}")

if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: universal_rwl_runner.py <prd_path> <output_dir>")
        sys.exit(1)

    prd_path = Path(sys.argv[1])
    output_dir = Path(sys.argv[2])
    run_prd(prd_path, output_dir)