#!/usr/bin/env python3
"""
OpenRouter Swarm Orchestrator - Wave 2
Fires 200 agents (100 Kimi K2.5 + 100 MiniMax M2.5) at fresh stories.
Reads from SWARM_MISSIONS_WAVE2.md.

USAGE:
    PYTHONUNBUFFERED=1 python scripts/openrouter_swarm_wave2.py

OUTPUT:
    /mnt/e/genesis-system/swarm-output/session20/kimi_results.jsonl
    /mnt/e/genesis-system/swarm-output/session20/minimax_results.jsonl
    /mnt/e/genesis-system/swarm-output/session20/WAVE2_REPORT.md
    /mnt/e/genesis-system/hive/agents_active_wave2.json

Author: Genesis System
Date: 2026-02-16
"""

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

try:
    import aiohttp
except ImportError:
    print("ERROR: aiohttp not installed. Install with: pip install aiohttp")
    sys.exit(1)

# ── Configuration ──────────────────────────────────────────────────────────
OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1/chat/completions"
OPENROUTER_API_KEY = "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"

MISSIONS_FILE = "/mnt/e/genesis-system/hive/SWARM_MISSIONS_WAVE2.md"
OUTPUT_DIR = "/mnt/e/genesis-system/swarm-output/session20"
KIMI_RESULTS_FILE = f"{OUTPUT_DIR}/kimi_results.jsonl"
MINIMAX_RESULTS_FILE = f"{OUTPUT_DIR}/minimax_results.jsonl"
REPORT_FILE = f"{OUTPUT_DIR}/WAVE2_REPORT.md"
MANIFEST_FILE = "/mnt/e/genesis-system/hive/agents_active_wave2.json"

# Models
MINIMAX_MODEL = "minimax/minimax-m2.5"
KIMI_MODEL = "moonshotai/kimi-k2.5"

# Batch sizes
BATCH_SIZE_MINIMAX = 20
BATCH_SIZE_KIMI = 5

# Timeouts
REQUEST_TIMEOUT = 300   # 5 min per request
BATCH_TIMEOUT = 600     # 10 min per batch


# ── Data Classes ───────────────────────────────────────────────────────────

@dataclass
class Story:
    story_id: str
    title: str
    role: str
    need: str
    benefit: str
    acceptance_criteria: List[str]
    black_box_tests: str
    white_box_tests: str
    estimated_tokens: int
    model: str
    prd_name: str

    def to_prompt(self) -> str:
        criteria = "\n".join(f"- {c}" for c in self.acceptance_criteria)
        return f"""# Story: {self.story_id} - {self.title}

**User Story:**
As a {self.role}, I need {self.need}, so that {self.benefit}.

**Acceptance Criteria:**
{criteria}

**Black Box Tests:**
{self.black_box_tests}

**White Box Tests:**
{self.white_box_tests}

**Instructions:**
Implement this story following the acceptance criteria. Provide production-ready code/content with inline comments, error handling, and best practices. Include example usage or test scenarios where applicable.
For code stories: Write clean, type-hinted Python (FastAPI) with comprehensive error handling.
For content stories: Write clear, compelling, accurate content optimized for conversion.
"""

    def system_prompt(self) -> str:
        if self.model.lower() == "kimi":
            return (
                "You are an expert software engineer building production-grade Python code for a SaaS product called Sunaiva "
                "(AI voice widgets for small businesses). Write clean, well-documented, type-hinted Python code with "
                "comprehensive error handling. Use FastAPI, asyncpg, Redis, Pydantic v2. Follow PEP 8. "
                "The backend runs on Elestio VPS with PostgreSQL, Redis, and serves widget.js to customer websites."
            )
        return (
            "You are an expert marketing strategist and conversion copywriter for Sunaiva — an AI voice widget product "
            "that answers business calls 24/7 using AI. Target market: Australian small businesses (plumbers, dentists, "
            "lawyers, real estate agents, restaurants, HVAC, auto repair, vets, accountants, property managers). "
            "Pricing: Starter $197/mo, Pro $397/mo, Growth $597/mo. 30% agency commission. "
            "Key selling points: never miss a call, 24/7 AI receptionist, 5-minute setup, works on any website. "
            "Write persuasive, authentic content. Use Australian English where appropriate."
        )


@dataclass
class AgentResult:
    agent_id: str
    team: str
    model: str
    story_id: str
    story_title: str
    prd_name: str
    status: str
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int
    cost_usd: float
    response_time_ms: int
    response: Optional[str] = None
    error: Optional[str] = None
    timestamp: str = ""

    def __post_init__(self):
        if not self.timestamp:
            self.timestamp = datetime.now().isoformat()


# ── Parser ─────────────────────────────────────────────────────────────────

def parse_missions(file_path: str) -> List[Story]:
    """Parse SWARM_MISSIONS_WAVE2.md into Story objects."""
    if not os.path.exists(file_path):
        print(f"ERROR: {file_path} not found")
        sys.exit(1)

    with open(file_path, 'r', encoding='utf-8') as f:
        content = f.read()

    stories = []
    current_prd = "Unknown"

    # Split on story headers
    sections = re.split(r'(?=^#### )', content, flags=re.MULTILINE)

    # Track PRD names
    for line in content.split('\n'):
        prd_match = re.match(r'^# PRD \d+:\s+(.+?)(?:\s*\(|$)', line)
        if prd_match:
            current_prd = prd_match.group(1).strip()

    # Reset PRD tracking during actual parsing
    current_prd = "Unknown"
    full_lines = content.split('\n')
    prd_line_map = {}
    for i, line in enumerate(full_lines):
        prd_match = re.match(r'^# PRD \d+:\s+(.+?)(?:\s*\(|$)', line)
        if prd_match:
            current_prd = prd_match.group(1).strip()
        prd_line_map[i] = current_prd

    for section in sections:
        if not section.strip().startswith('####'):
            continue

        # Find this section's position in the original content to get PRD context
        section_first_line = section.strip().split('\n')[0]
        for i, line in enumerate(full_lines):
            if line.strip() == section_first_line.strip():
                current_prd = prd_line_map.get(i, "Unknown")
                break

        story = _parse_story(section, current_prd)
        if story:
            stories.append(story)

    print(f"Parsed {len(stories)} stories from {file_path}")
    return stories


def _parse_story(text: str, prd_name: str) -> Optional[Story]:
    """Parse a single story section."""
    try:
        # Story ID and title
        m = re.search(r'####\s+([\w-]+):\s+(.+)', text)
        if not m:
            return None
        story_id, title = m.group(1), m.group(2).strip()

        # User story
        role_m = re.search(r'\*\*As a\*\*\s+(.+?),', text)
        need_m = re.search(r'\*\*I need\*\*\s+(.+?),\s+\*\*so that\*\*', text)
        benefit_m = re.search(r'\*\*so that\*\*\s+(.+?)\.', text)
        if not (role_m and need_m and benefit_m):
            return None

        # Acceptance criteria - handle both **Key:** and **Key**: formats
        criteria_m = re.search(r'\*\*Acceptance Criteria(?::\*\*|\*\*:)\n(.*?)\*\*Black Box Tests', text, re.DOTALL)
        criteria = []
        if criteria_m:
            for line in criteria_m.group(1).split('\n'):
                line = line.strip()
                if line.startswith('-'):
                    criteria.append(line[1:].strip())

        # Tests - handle both **Key:** and **Key**: formats
        bb_m = re.search(r'\*\*Black Box Tests(?::\*\*|\*\*:)\s*(.+?)(?:\n\*\*White Box Tests)', text, re.DOTALL)
        wb_m = re.search(r'\*\*White Box Tests(?::\*\*|\*\*:)\s*(.+?)(?:\n\*\*Estimated Tokens)', text, re.DOTALL)

        # Tokens and model - handle both **Key:** and **Key**: formats
        tok_m = re.search(r'\*\*Estimated Tokens(?::\*\*|\*\*:)\s*([\d,]+)', text)
        mod_m = re.search(r'\*\*Model(?::\*\*|\*\*:)\s*(\w+)', text)

        return Story(
            story_id=story_id,
            title=title,
            role=role_m.group(1).strip(),
            need=need_m.group(1).strip(),
            benefit=benefit_m.group(1).strip(),
            acceptance_criteria=criteria,
            black_box_tests=bb_m.group(1).strip() if bb_m else "",
            white_box_tests=wb_m.group(1).strip() if wb_m else "",
            estimated_tokens=int(tok_m.group(1).replace(',', '')) if tok_m else 5000,
            model=mod_m.group(1) if mod_m else "Kimi",
            prd_name=prd_name
        )
    except Exception as e:
        print(f"  WARNING: Failed to parse story: {e}")
        return None


# ── Agent Execution ────────────────────────────────────────────────────────

async def execute_agent(
    session: aiohttp.ClientSession,
    story: Story,
    agent_id: str,
    model_id: str
) -> AgentResult:
    """Fire a single agent via OpenRouter."""
    start = time.time()
    team = "KIMI" if "kimi" in model_id else "MINIMAX"

    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://github.com/genesis-system",
        "X-Title": "Genesis Swarm Wave2"
    }

    payload = {
        "model": model_id,
        "messages": [
            {"role": "system", "content": story.system_prompt()},
            {"role": "user", "content": story.to_prompt()}
        ],
        "temperature": 0.7,
        "max_tokens": min(story.estimated_tokens + 1000, 8192)
    }

    try:
        timeout = aiohttp.ClientTimeout(total=REQUEST_TIMEOUT)
        async with session.post(OPENROUTER_BASE_URL, json=payload, headers=headers, timeout=timeout) as resp:
            elapsed = int((time.time() - start) * 1000)

            if resp.status != 200:
                err = await resp.text()
                return AgentResult(
                    agent_id=agent_id, team=team, model=model_id,
                    story_id=story.story_id, story_title=story.title,
                    prd_name=story.prd_name, status="fail",
                    prompt_tokens=0, completion_tokens=0, total_tokens=0,
                    cost_usd=0.0, response_time_ms=elapsed,
                    error=f"HTTP {resp.status}: {err[:300]}"
                )

            data = await resp.json()
            usage = data.get('usage', {})
            pt = usage.get('prompt_tokens', 0)
            ct = usage.get('completion_tokens', 0)
            tt = usage.get('total_tokens', 0)
            rate = 1.07 if 'kimi' in model_id else 1.00
            cost = (tt / 1_000_000) * rate
            text = data['choices'][0]['message']['content']

            return AgentResult(
                agent_id=agent_id, team=team, model=model_id,
                story_id=story.story_id, story_title=story.title,
                prd_name=story.prd_name, status="success",
                prompt_tokens=pt, completion_tokens=ct, total_tokens=tt,
                cost_usd=cost, response_time_ms=elapsed, response=text
            )

    except asyncio.TimeoutError:
        elapsed = int((time.time() - start) * 1000)
        return AgentResult(
            agent_id=agent_id, team=team, model=model_id,
            story_id=story.story_id, story_title=story.title,
            prd_name=story.prd_name, status="fail",
            prompt_tokens=0, completion_tokens=0, total_tokens=0,
            cost_usd=0.0, response_time_ms=elapsed,
            error=f"Timeout after {REQUEST_TIMEOUT}s"
        )
    except Exception as e:
        elapsed = int((time.time() - start) * 1000)
        return AgentResult(
            agent_id=agent_id, team=team, model=model_id,
            story_id=story.story_id, story_title=story.title,
            prd_name=story.prd_name, status="fail",
            prompt_tokens=0, completion_tokens=0, total_tokens=0,
            cost_usd=0.0, response_time_ms=elapsed,
            error=str(e)[:300]
        )


async def fire_team(
    stories: List[Story],
    model_id: str,
    batch_size: int,
    team_name: str
) -> List[AgentResult]:
    """Fire all stories for a team in batches."""
    results = []
    total = len(stories)
    print(f"\n{'='*60}")
    print(f"  TEAM {team_name}: {total} agents | Model: {model_id}")
    print(f"  Batch size: {batch_size} | Timeout: {REQUEST_TIMEOUT}s")
    print(f"{'='*60}")

    async with aiohttp.ClientSession() as session:
        for i in range(0, total, batch_size):
            batch = stories[i:i+batch_size]
            batch_num = (i // batch_size) + 1
            total_batches = (total + batch_size - 1) // batch_size

            print(f"  Batch {batch_num}/{total_batches}: Firing {len(batch)} agents...")

            # Generate agent IDs
            agent_ids = [
                hashlib.md5(f"{s.story_id}-{datetime.now().isoformat()}".encode()).hexdigest()[:8]
                for s in batch
            ]

            try:
                batch_results = await asyncio.wait_for(
                    asyncio.gather(*[
                        execute_agent(session, story, aid, model_id)
                        for story, aid in zip(batch, agent_ids)
                    ]),
                    timeout=BATCH_TIMEOUT
                )
                results.extend(batch_results)

                ok = sum(1 for r in batch_results if r.status == "success")
                fail = sum(1 for r in batch_results if r.status == "fail")
                cost = sum(r.cost_usd for r in batch_results)
                print(f"    OK: {ok} | FAIL: {fail} | Cost: ${cost:.4f}")

                # Delay between batches for Kimi rate limits
                if 'kimi' in model_id and i + batch_size < total:
                    print(f"    Waiting 10s (Kimi rate limit)...")
                    await asyncio.sleep(10)

            except asyncio.TimeoutError:
                print(f"    BATCH TIMEOUT after {BATCH_TIMEOUT}s — marking all failed")
                for s, aid in zip(batch, agent_ids):
                    results.append(AgentResult(
                        agent_id=aid, team=team_name, model=model_id,
                        story_id=s.story_id, story_title=s.title,
                        prd_name=s.prd_name, status="fail",
                        prompt_tokens=0, completion_tokens=0, total_tokens=0,
                        cost_usd=0.0, response_time_ms=BATCH_TIMEOUT * 1000,
                        error=f"Batch timeout {BATCH_TIMEOUT}s"
                    ))

    return results


def save_results(results: List[AgentResult], filepath: str):
    """Save results to JSONL."""
    os.makedirs(os.path.dirname(filepath), exist_ok=True)
    with open(filepath, 'w') as f:
        for r in results:
            f.write(json.dumps(asdict(r)) + '\n')
    print(f"  Saved {len(results)} results -> {filepath}")


def save_manifest(kimi_results: List[AgentResult], minimax_results: List[AgentResult]):
    """Write agent manifest."""
    manifest = {
        "wave": 2,
        "session": 20,
        "timestamp": datetime.now().isoformat(),
        "total_agents": len(kimi_results) + len(minimax_results),
        "teams": {
            "kimi": {
                "model": KIMI_MODEL,
                "total": len(kimi_results),
                "success": sum(1 for r in kimi_results if r.status == "success"),
                "failed": sum(1 for r in kimi_results if r.status == "fail"),
                "total_tokens": sum(r.total_tokens for r in kimi_results),
                "total_cost": sum(r.cost_usd for r in kimi_results),
                "agents": [
                    {
                        "agent_id": r.agent_id,
                        "story_id": r.story_id,
                        "status": r.status,
                        "tokens": r.total_tokens,
                        "cost": r.cost_usd,
                        "response_time_ms": r.response_time_ms
                    }
                    for r in kimi_results
                ]
            },
            "minimax": {
                "model": MINIMAX_MODEL,
                "total": len(minimax_results),
                "success": sum(1 for r in minimax_results if r.status == "success"),
                "failed": sum(1 for r in minimax_results if r.status == "fail"),
                "total_tokens": sum(r.total_tokens for r in minimax_results),
                "total_cost": sum(r.cost_usd for r in minimax_results),
                "agents": [
                    {
                        "agent_id": r.agent_id,
                        "story_id": r.story_id,
                        "status": r.status,
                        "tokens": r.total_tokens,
                        "cost": r.cost_usd,
                        "response_time_ms": r.response_time_ms
                    }
                    for r in minimax_results
                ]
            }
        }
    }

    os.makedirs(os.path.dirname(MANIFEST_FILE), exist_ok=True)
    with open(MANIFEST_FILE, 'w') as f:
        json.dump(manifest, f, indent=2)
    print(f"\nAgent manifest -> {MANIFEST_FILE}")


def generate_report(kimi_results: List[AgentResult], minimax_results: List[AgentResult]):
    """Generate execution report."""
    all_r = kimi_results + minimax_results
    total = len(all_r)
    ok = sum(1 for r in all_r if r.status == "success")
    fail = sum(1 for r in all_r if r.status == "fail")
    tokens = sum(r.total_tokens for r in all_r)
    cost = sum(r.cost_usd for r in all_r)
    avg_time = sum(r.response_time_ms for r in all_r) / total if total else 0

    k_ok = sum(1 for r in kimi_results if r.status == "success")
    k_fail = sum(1 for r in kimi_results if r.status == "fail")
    k_cost = sum(r.cost_usd for r in kimi_results)
    k_tokens = sum(r.total_tokens for r in kimi_results)

    m_ok = sum(1 for r in minimax_results if r.status == "success")
    m_fail = sum(1 for r in minimax_results if r.status == "fail")
    m_cost = sum(r.cost_usd for r in minimax_results)
    m_tokens = sum(r.total_tokens for r in minimax_results)

    report = f"""# Wave 2 Swarm Execution Report
**Session 20 | {datetime.now().strftime('%Y-%m-%d %H:%M')}**

## Summary

| Metric | Value |
|--------|-------|
| Total Agents | {total} |
| Successful | {ok} ({ok/total*100:.1f}%) |
| Failed | {fail} |
| Total Tokens | {tokens:,} |
| Total Cost | ${cost:.4f} |
| Avg Response Time | {avg_time:.0f}ms |

## Team Kimi K2.5 (Code & Infrastructure)

| Metric | Value |
|--------|-------|
| Stories | {len(kimi_results)} |
| Success | {k_ok} |
| Failed | {k_fail} |
| Tokens | {k_tokens:,} |
| Cost | ${k_cost:.4f} |

### Kimi Stories
"""
    for r in kimi_results:
        icon = "OK" if r.status == "success" else "FAIL"
        report += f"- [{icon}] **{r.story_id}**: {r.story_title} (${r.cost_usd:.4f}, {r.response_time_ms}ms)\n"

    report += f"""
## Team MiniMax M2.5 (Marketing & Content)

| Metric | Value |
|--------|-------|
| Stories | {len(minimax_results)} |
| Success | {m_ok} |
| Failed | {m_fail} |
| Tokens | {m_tokens:,} |
| Cost | ${m_cost:.4f} |

### MiniMax Stories
"""
    for r in minimax_results:
        icon = "OK" if r.status == "success" else "FAIL"
        report += f"- [{icon}] **{r.story_id}**: {r.story_title} (${r.cost_usd:.4f}, {r.response_time_ms}ms)\n"

    # Failed stories detail
    failed = [r for r in all_r if r.status == "fail"]
    if failed:
        report += f"\n## Failed Stories ({len(failed)})\n\n"
        for r in failed:
            report += f"- **{r.story_id}** ({r.team}): {r.error}\n"

    report += f"""
## Output Files

- Kimi results: `{KIMI_RESULTS_FILE}`
- MiniMax results: `{MINIMAX_RESULTS_FILE}`
- Agent manifest: `{MANIFEST_FILE}`
- This report: `{REPORT_FILE}`

---
Generated by Genesis OpenRouter Swarm Wave 2 | {datetime.now().isoformat()}
"""

    with open(REPORT_FILE, 'w') as f:
        f.write(report)
    print(f"Report -> {REPORT_FILE}")


# ── Main ───────────────────────────────────────────────────────────────────

async def main():
    print(f"\n{'='*60}")
    print("  GENESIS SWARM WAVE 2 — 200 AGENTS — SHIP TODAY")
    print(f"  {datetime.now().isoformat()}")
    print(f"{'='*60}")

    # Parse stories
    stories = parse_missions(MISSIONS_FILE)
    if not stories:
        print("ERROR: No stories found!")
        sys.exit(1)

    kimi_stories = [s for s in stories if s.model.lower() == "kimi"]
    minimax_stories = [s for s in stories if s.model.lower() == "minimax"]

    print(f"\nKimi stories: {len(kimi_stories)}")
    print(f"MiniMax stories: {len(minimax_stories)}")
    print(f"Total: {len(stories)}")

    # Fire both teams
    kimi_results = await fire_team(kimi_stories, KIMI_MODEL, BATCH_SIZE_KIMI, "KIMI")
    save_results(kimi_results, KIMI_RESULTS_FILE)

    minimax_results = await fire_team(minimax_stories, MINIMAX_MODEL, BATCH_SIZE_MINIMAX, "MINIMAX")
    save_results(minimax_results, MINIMAX_RESULTS_FILE)

    # Generate report and manifest
    generate_report(kimi_results, minimax_results)
    save_manifest(kimi_results, minimax_results)

    # Final summary
    all_r = kimi_results + minimax_results
    ok = sum(1 for r in all_r if r.status == "success")
    cost = sum(r.cost_usd for r in all_r)
    print(f"\n{'='*60}")
    print(f"  WAVE 2 COMPLETE")
    print(f"  {ok}/{len(all_r)} agents succeeded | ${cost:.4f} total cost")
    print(f"{'='*60}\n")


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