#!/usr/bin/env python3
"""
Genesis $50/Day Swarm Launcher
================================
Budget: $50 USD/day = ~8,000 MiniMax stories at $0.006/story

Usage:
    python3 launch_swarm_50.py --mission "Build X" --stories 50
    python3 launch_swarm_50.py --mission "Research Y" --agents 20 --model minimax
    python3 launch_swarm_50.py --budget-check      # Show remaining daily budget
    python3 launch_swarm_50.py --list-models       # Show available models and costs
    python3 launch_swarm_50.py --status            # Show all active swarms
"""

import argparse
import json
import os
import sys
import uuid
from datetime import datetime, timezone, timedelta
from pathlib import Path

# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------
REPO_ROOT = Path("/mnt/e/genesis-system")
SECRETS_FILE = REPO_ROOT / "config" / "secrets.env"
DAILY_BUDGET_FILE = REPO_ROOT / "data" / "daily_budget.json"
SWARM_PROGRESS_DIR = REPO_ROOT / "data" / "swarm_progress"

sys.path.insert(0, str(REPO_ROOT))

DAILY_BUDGET_LIMIT_USD = 50.0
KILL_THRESHOLD_USD = 20.0  # Max per single worker before forced kill
AEST_OFFSET = timedelta(hours=10)

# ---------------------------------------------------------------------------
# Available models via OpenRouter
# ---------------------------------------------------------------------------
MODELS: dict[str, dict] = {
    "gemini-flash": {
        "openrouter_id": "google/gemini-2.5-flash",
        "cost_per_call_usd": 0.002,
        "best_for": ["research", "extraction", "bulk"],
        "stories_per_dollar": 500,
    },
    "minimax": {
        "openrouter_id": "minimax/minimax-01",
        "cost_per_call_usd": 0.006,
        "best_for": ["coding", "backend", "api"],
        "stories_per_dollar": 167,
    },
    "deepseek": {
        "openrouter_id": "deepseek/deepseek-chat",
        "cost_per_call_usd": 0.001,
        "best_for": ["extraction", "classification", "kg"],
        "stories_per_dollar": 1000,
    },
    "kimi": {
        "openrouter_id": "moonshot/moonshot-v1-8k",
        "cost_per_call_usd": 0.003,
        "best_for": ["frontend", "ui", "react", "html"],
        "stories_per_dollar": 333,
    },
    "haiku": {
        "openrouter_id": "anthropic/claude-haiku-4-5",
        "cost_per_call_usd": 0.004,
        "best_for": ["testing", "summarization", "voice"],
        "stories_per_dollar": 250,
    },
    "sonnet": {
        "openrouter_id": "anthropic/claude-sonnet-4-5",
        "cost_per_call_usd": 0.012,
        "best_for": ["architecture", "complex", "review"],
        "stories_per_dollar": 83,
    },
}

TASK_TO_MODEL: dict[str, str] = {
    "research": "gemini-flash",
    "find": "gemini-flash",
    "search": "gemini-flash",
    "analyse": "gemini-flash",
    "analyze": "gemini-flash",
    "scrape": "gemini-flash",
    "build": "minimax",
    "code": "minimax",
    "implement": "minimax",
    "fix": "minimax",
    "backend": "minimax",
    "api": "minimax",
    "extract": "deepseek",
    "classify": "deepseek",
    "kg": "deepseek",
    "knowledge": "deepseek",
    "frontend": "kimi",
    "ui": "kimi",
    "react": "kimi",
    "html": "kimi",
    "test": "haiku",
    "summarize": "haiku",
}


# ---------------------------------------------------------------------------
# Secrets loader
# ---------------------------------------------------------------------------

def load_openrouter_key() -> str:
    """Load OPENROUTER_API_KEY from config/secrets.env."""
    if not SECRETS_FILE.exists():
        print(f"[ERROR] Secrets file not found: {SECRETS_FILE}")
        return ""
    for line in SECRETS_FILE.read_text().splitlines():
        line = line.strip()
        if line.startswith("OPENROUTER_API_KEY"):
            parts = line.split("=", 1)
            if len(parts) == 2:
                return parts[1].strip().strip('"').strip("'")
    # Also check environment
    return os.environ.get("OPENROUTER_API_KEY", "")


# ---------------------------------------------------------------------------
# Budget tracking
# ---------------------------------------------------------------------------

def _today_aest() -> str:
    return (datetime.now(timezone.utc) + AEST_OFFSET).strftime("%Y-%m-%d")


def _now_ts() -> str:
    return datetime.now(timezone.utc).isoformat()


def load_budget() -> dict:
    """Load or initialise daily budget JSON."""
    DAILY_BUDGET_FILE.parent.mkdir(parents=True, exist_ok=True)
    today = _today_aest()

    if DAILY_BUDGET_FILE.exists():
        try:
            data = json.loads(DAILY_BUDGET_FILE.read_text())
            if data.get("date_aest") == today:
                return data
        except (json.JSONDecodeError, OSError):
            pass

    # New day or corrupt — reset
    fresh = {
        "date_aest": today,
        "total_spent_usd": 0.0,
        "story_count": 0,
        "agent_runs": [],
        "last_updated": _now_ts(),
    }
    DAILY_BUDGET_FILE.write_text(json.dumps(fresh, indent=2))
    print(f"[BUDGET] Initialised daily budget for {today} — $0.00 / ${DAILY_BUDGET_LIMIT_USD:.2f}")
    return fresh


def save_budget(budget: dict) -> None:
    budget["last_updated"] = _now_ts()
    DAILY_BUDGET_FILE.write_text(json.dumps(budget, indent=2))


def check_budget(estimated_cost: float) -> tuple[bool, float]:
    """
    Check if we have budget for estimated_cost.

    Returns
    -------
    (ok: bool, remaining_usd: float)
    """
    budget = load_budget()
    spent = budget.get("total_spent_usd", 0.0)
    remaining = DAILY_BUDGET_LIMIT_USD - spent
    return estimated_cost <= remaining, remaining


def record_spend(run_id: str, model: str, agent_count: int, cost: float, mission: str) -> None:
    """Record a spend event in daily_budget.json."""
    budget = load_budget()
    budget["total_spent_usd"] = budget.get("total_spent_usd", 0.0) + cost
    budget["story_count"] = budget.get("story_count", 0) + agent_count
    budget.setdefault("agent_runs", []).append({
        "run_id": run_id,
        "mission": mission[:80],
        "model": model,
        "agents": agent_count,
        "cost_usd": cost,
        "timestamp": _now_ts(),
    })
    save_budget(budget)


# ---------------------------------------------------------------------------
# Auto model selection
# ---------------------------------------------------------------------------

def auto_select_model(mission: str) -> str:
    """Select best model based on mission keywords."""
    mission_lower = mission.lower()
    for keyword, model_key in TASK_TO_MODEL.items():
        if keyword in mission_lower:
            return model_key
    return "minimax"  # Default: best coding model


# ---------------------------------------------------------------------------
# Swarm launcher
# ---------------------------------------------------------------------------

def estimate_cost(model_key: str, agent_count: int, stories: int) -> float:
    """Estimate total cost for a swarm run."""
    model = MODELS.get(model_key, MODELS["minimax"])
    calls_per_agent = max(1, stories // max(1, agent_count))
    return model["cost_per_call_usd"] * agent_count * calls_per_agent


def launch_swarm(
    mission: str,
    stories: int = 50,
    agents: int = 10,
    model_key: str = "",
    dry_run: bool = False,
) -> dict:
    """
    Launch a swarm for a given mission.

    Parameters
    ----------
    mission : str
        High-level task description
    stories : int
        Target number of stories/tasks to complete
    agents : int
        Number of parallel agents
    model_key : str
        Model shortname (minimax, gemini-flash, deepseek, kimi, haiku, sonnet)
        Auto-selected if empty.
    dry_run : bool
        If True, show estimate but don't launch

    Returns
    -------
    dict with run metadata
    """
    SWARM_PROGRESS_DIR.mkdir(parents=True, exist_ok=True)

    # Auto-select model
    if not model_key or model_key not in MODELS:
        model_key = auto_select_model(mission)

    model_info = MODELS[model_key]
    estimated_cost = estimate_cost(model_key, agents, stories)

    print(f"\n[SWARM LAUNCHER]")
    print(f"  Mission   : {mission[:80]}")
    print(f"  Stories   : {stories}")
    print(f"  Agents    : {agents}")
    print(f"  Model     : {model_key} ({model_info['openrouter_id']})")
    print(f"  Est. Cost : ${estimated_cost:.4f}")
    print(f"  Best for  : {', '.join(model_info['best_for'])}")

    # Budget check
    budget_ok, remaining = check_budget(estimated_cost)
    print(f"  Remaining : ${remaining:.2f} / ${DAILY_BUDGET_LIMIT_USD:.2f}")

    if not budget_ok:
        print(f"\n[KILL SWITCH] Daily budget exceeded! ${remaining:.2f} remaining < ${estimated_cost:.4f} needed.")
        print("[KILL SWITCH] Refusing to launch new agents. Wait for budget reset at midnight AEST.")
        return {"status": "rejected", "reason": "budget_exceeded", "remaining_usd": remaining}

    if estimated_cost > KILL_THRESHOLD_USD:
        print(f"\n[WARNING] Single-run cost ${estimated_cost:.2f} exceeds kill threshold ${KILL_THRESHOLD_USD:.2f}")
        print("[WARNING] Reduce agent count or story count to proceed safely.")
        return {"status": "rejected", "reason": "per_run_kill_threshold", "estimated_cost": estimated_cost}

    if dry_run:
        print("\n[DRY RUN] Estimation complete. Pass --execute to launch.")
        return {"status": "dry_run", "estimated_cost": estimated_cost, "model": model_key}

    # Launch
    run_id = f"swarm_{uuid.uuid4().hex[:8]}"
    api_key = load_openrouter_key()

    if not api_key:
        print("[ERROR] OPENROUTER_API_KEY not found in config/secrets.env")
        print("[ERROR] Add: OPENROUTER_API_KEY=sk-or-...")
        return {"status": "error", "reason": "missing_api_key"}

    # Write swarm state
    swarm_state = {
        "run_id": run_id,
        "mission": mission,
        "model": model_info["openrouter_id"],
        "model_key": model_key,
        "agent_count": agents,
        "stories_total": stories,
        "stories_completed": 0,
        "stories_completed_this_hour": 0,
        "status": "running",
        "estimated_cost_usd": estimated_cost,
        "started_at": _now_ts(),
        "output_dir": str(SWARM_PROGRESS_DIR / run_id),
    }
    state_file = SWARM_PROGRESS_DIR / f"{run_id}.json"
    state_file.write_text(json.dumps(swarm_state, indent=2))

    # Record spend
    record_spend(run_id, model_info["openrouter_id"], agents, estimated_cost, mission)

    print(f"\n[LAUNCHED] Swarm {run_id}")
    print(f"[LAUNCHED] {agents} agents x {stories} stories via {model_key}")
    print(f"[LAUNCHED] Framework-only mode — attach live API calls in agent_runner.py")
    print(f"[LAUNCHED] State: {state_file}")

    # NOTE: Live API calls would go here:
    # from core.openrouter_client import OpenRouterClient
    # client = OpenRouterClient(api_key=api_key)
    # For each story batch:
    #   response = client.chat(model=model_info["openrouter_id"], messages=[...])

    return swarm_state


# ---------------------------------------------------------------------------
# Status display
# ---------------------------------------------------------------------------

def show_budget() -> None:
    budget = load_budget()
    spent = budget.get("total_spent_usd", 0.0)
    remaining = DAILY_BUDGET_LIMIT_USD - spent
    pct = (spent / DAILY_BUDGET_LIMIT_USD) * 100
    bar_filled = int(20 * pct / 100)
    bar = "#" * bar_filled + "-" * (20 - bar_filled)
    runs = budget.get("agent_runs", [])

    print(f"\n[BUDGET — {budget.get('date_aest', 'unknown')} AEST]")
    print(f"  Spent    : ${spent:.4f} / ${DAILY_BUDGET_LIMIT_USD:.2f}")
    print(f"  Progress : [{bar}] {pct:.1f}%")
    print(f"  Remaining: ${remaining:.4f}")
    print(f"  Stories  : {budget.get('story_count', 0)}")
    print(f"  Runs     : {len(runs)}")
    if runs:
        print(f"\n  Recent runs:")
        for run in runs[-5:]:
            print(f"    [{run['run_id']}] {run['model'][:30]} x{run['agents']} — ${run['cost_usd']:.4f} — {run['mission'][:40]}")


def show_status() -> None:
    if not SWARM_PROGRESS_DIR.exists():
        print("[STATUS] No swarm progress directory found.")
        return

    files = sorted(SWARM_PROGRESS_DIR.glob("*.json"))
    if not files:
        print("[STATUS] No swarms found.")
        return

    print(f"\n[SWARM STATUS — {len(files)} total]")
    print(f"{'RUN ID':<20} {'GENERAL':<12} {'MODEL':<30} {'AGENTS':>6} {'DONE':>5} {'STATUS':<12}")
    print("-" * 90)
    for f in files[-20:]:  # Show last 20
        try:
            d = json.loads(f.read_text())
            run_id = d.get("run_id") or d.get("swarm_id", f.stem)
            general = d.get("general", "?")
            model = d.get("model", "?")[:28]
            agents = d.get("agent_count", 0)
            done = d.get("stories_completed", 0)
            status = d.get("status", "?")
            print(f"{run_id:<20} {general:<12} {model:<30} {agents:>6} {done:>5} {status:<12}")
        except (json.JSONDecodeError, OSError):
            print(f"{f.stem:<20} [corrupt state file]")


def list_models() -> None:
    print(f"\n[AVAILABLE MODELS via OpenRouter]")
    print(f"{'KEY':<14} {'OPENROUTER ID':<35} {'$/CALL':>8}  BEST FOR")
    print("-" * 80)
    for key, info in MODELS.items():
        print(f"{key:<14} {info['openrouter_id']:<35} ${info['cost_per_call_usd']:>6.4f}  {', '.join(info['best_for'])}")
    print(f"\nOpenRouter endpoint: https://openrouter.ai/api/v1/chat/completions")


# ---------------------------------------------------------------------------
# CLI entry point
# ---------------------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(
        description="Genesis $50/Day Swarm Launcher",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument("--mission", "-m", help="Task/mission description")
    parser.add_argument("--stories", "-s", type=int, default=50, help="Target story count (default 50)")
    parser.add_argument("--agents", "-a", type=int, default=10, help="Parallel agents (default 10)")
    parser.add_argument("--model", help=f"Model key: {', '.join(MODELS.keys())}")
    parser.add_argument("--execute", action="store_true", help="Actually launch (default: dry-run)")
    parser.add_argument("--budget-check", action="store_true", help="Show daily budget status")
    parser.add_argument("--status", action="store_true", help="Show all swarm statuses")
    parser.add_argument("--list-models", action="store_true", help="List available models")

    args = parser.parse_args()

    if args.budget_check:
        show_budget()
        return

    if args.status:
        show_status()
        show_budget()
        return

    if args.list_models:
        list_models()
        return

    if not args.mission:
        parser.print_help()
        print("\n[ERROR] --mission is required. Example:")
        print("  python3 launch_swarm_50.py --mission 'Build the voice widget checkout' --execute")
        sys.exit(1)

    result = launch_swarm(
        mission=args.mission,
        stories=args.stories,
        agents=args.agents,
        model_key=args.model or "",
        dry_run=not args.execute,
    )

    if result.get("status") == "rejected":
        sys.exit(2)


if __name__ == "__main__":
    # Ensure daily_budget.json is initialised
    load_budget()
    main()