main.py

#!/usr/bin/env python3
"""Pipeline orchestrator — Obsidian vault-based architecture.

Usage:
    python main.py <project_dir> <vault_path> "<task description>"
    python main.py <project_dir> <vault_path> "<task>" --rebuild
"""

import sys
import os
import json
import argparse
import subprocess
from pathlib import Path

# Add project root to path so agents can import each other
script_dir = Path(__file__).parent.resolve()
sys.path.insert(0, str(script_dir))

from agents.vault_loader import Vault
from agents.stack_detector import StackDetector
from agents.context_agent import ContextAgent
from agents.po_agent import POAgent
from agents.ba_agent import BAAgent
from agents.architect_agent import ArchitectAgent
from agents.devflow_agent import DevflowAgent
from agents.dev_agent import DevAgent
from agents.qa_agent import QAAgent
from agents.testing_agent import TestingAgent


def load_env(env_path: Path):
    """Load .env file if present."""
    if not env_path.exists():
        return
    with open(env_path) as f:
        for line in f:
            line = line.strip()
            if line and not line.startswith("#") and "=" in line:
                key, value = line.split("=", 1)
                os.environ.setdefault(key, value)


def run_context_creator(project_dir: Path, vault_project_dir: Path, rebuild: bool):
    """Phase 0: Run context-creator.py to build the vault.

    Skips if vault exists and --rebuild not set.
    """
    if Vault.exists(vault_project_dir) and not rebuild:
        print(f"[INFO] Vault already exists at {vault_project_dir}, skipping context-creator.", file=sys.stderr)
        return True

    print(f"[INFO] Building vault at {vault_project_dir}...", file=sys.stderr)
    context_creator = script_dir / "agents" / "context-creator.py"
    if not context_creator.exists():
        print(f"[ERROR] context-creator.py not found at {context_creator}", file=sys.stderr)
        return False

    venv_python = script_dir / ".venv" / "bin" / "python"
    python_exe = str(venv_python) if venv_python.exists() else sys.executable

    # Generate markdown vault
    result = subprocess.run(
        [python_exe, str(context_creator), str(project_dir), "--out", str(vault_project_dir)],
        capture_output=True,
        text=True,
        timeout=600,
    )
    if result.stdout:
        print(result.stdout, file=sys.stderr)
    if result.stderr:
        print(result.stderr, file=sys.stderr)
    if result.returncode != 0:
        print(f"[ERROR] context-creator.py failed with exit code {result.returncode}", file=sys.stderr)
        return False

    print(f"[INFO] Vault built at {vault_project_dir}", file=sys.stderr)
    return True


def run_dev_context_creator(project_dir: Path, req_path: Path, out_path: Path):
    """Run context-creator.py --dev-context to build the consolidated source context."""
    print(f"[INFO] Generating dev context at {out_path}...", file=sys.stderr)
    context_creator = script_dir / "agents" / "context-creator.py"
    if not context_creator.exists():
        print(f"[ERROR] context-creator.py not found at {context_creator}", file=sys.stderr)
        return False

    venv_python = script_dir / ".venv" / "bin" / "python"
    python_exe = str(venv_python) if venv_python.exists() else sys.executable

    result = subprocess.run(
        [
            python_exe,
            str(context_creator),
            str(project_dir),
            "--dev-context",
            "--requirements",
            str(req_path),
            "--out",
            str(out_path),
        ],
        capture_output=True,
        text=True,
    )
    if result.stdout:
        print(result.stdout, file=sys.stderr)
    if result.stderr:
        print(result.stderr, file=sys.stderr)
    if result.returncode != 0:
        print(f"[ERROR] dev context generation failed with exit code {result.returncode}", file=sys.stderr)
        return False

    return True


def select_ai_provider():
    """Interactive AI provider selection."""
    env_provider = os.environ.get("AI_PROVIDER", "")
    if env_provider in ("gemini", "groq"):
        print(f"[INFO] Using AI provider '{env_provider}' from environment.", file=sys.stderr)
        return env_provider

    groq_key = os.environ.get("GROQ_API_KEY", "")
    print("\n=== AI Provider Selection ===", file=sys.stderr)
    print("1. Gemini (Google)", file=sys.stderr)
    print("2. Groq (Llama-3.3-70b-versatile)", file=sys.stderr)
    print("==============================\n", file=sys.stderr)

    choice = input("Enter your choice (1-2), or press Enter to use Gemini: ").strip()
    if choice == "2":
        provider = "groq"
        if not groq_key:
            print("[WARNING] GROQ_API_KEY not set!", file=sys.stderr)
    else:
        provider = "gemini"

    print(f"[INFO] Selected provider: {provider}", file=sys.stderr)
    os.environ["AI_PROVIDER"] = provider
    return provider


def print_pipeline_summary(req_path: Path):
    """Print a summary of pipeline-requirements.json for user review."""
    if not req_path.exists():
        return
    req = json.loads(req_path.read_text())

    print("\n" + "=" * 50, file=sys.stderr)
    print("=== Phase 1 Complete ===", file=sys.stderr)
    print(f"Task: {req.get('task', 'N/A')}", file=sys.stderr)
    print(f"User stories: {len(req.get('po', {}).get('userStories', []))}", file=sys.stderr)
    print(f"Acceptance criteria: {len(req.get('po', {}).get('acceptanceCriteria', []))}", file=sys.stderr)
    print(f"Edge cases: {len(req.get('ba', {}).get('edgeCases', []))}", file=sys.stderr)
    print(f"Test scenarios: {len(req.get('ba', {}).get('testScenarios', []))}", file=sys.stderr)

    br = req.get("architect", {}).get("blastRadius", {})
    must = len(br.get("mustChange", []))
    may = len(br.get("mayAffect", []))
    print(f"Blast radius: {must} files must change, {may} may affect", file=sys.stderr)
    if br.get("mustChange"):
        for f in br["mustChange"]:
            print(f"  - {f.get('path', '?')}: {f.get('reason', '')}", file=sys.stderr)
    if br.get("mayAffect"):
        for f in br["mayAffect"]:
            print(f"  ~ {f.get('path', '?')}: {f.get('reason', '')}", file=sys.stderr)

    devflow = req.get("devflow", [])
    print(f"Dev tasks: {len(devflow)}", file=sys.stderr)
    for t in devflow:
        print(f"  [{t.get('type', '?')}] {t.get('file', '?')} — {t.get('task', '')}", file=sys.stderr)
    print("=" * 50, file=sys.stderr)


def main():
    parser = argparse.ArgumentParser(description="AI-driven autonomous development pipeline")
    parser.add_argument("project_dir", help="Path to the target project directory")
    parser.add_argument("task", help="Task description (e.g. 'Add user authentication endpoint')")
    parser.add_argument(
        "vault_path",
        nargs="?",
        default=None,
        help="Path to the Obsidian vault root. Defaults to ~/contx/<project_name>.",
    )
    parser.add_argument("--rebuild", action="store_true", help="Force vault rebuild")
    args = parser.parse_args()

    project_dir = Path(args.project_dir).resolve()
    if not project_dir.exists():
        print(f"[ERROR] Project directory does not exist: {project_dir}", file=sys.stderr)
        sys.exit(1)

    artifacts_dir = project_dir / "artifacts"
    rebuild_context = args.rebuild
    
    if artifacts_dir.exists() and any(artifacts_dir.iterdir()):
        print(f"\n[INFO] Previous artifacts found in {artifacts_dir}", file=sys.stderr)
        try:
            choice = input("Do you want to (r)esume previous session or (s)tart fresh? [r/s]: ").strip().lower()
        except EOFError:
            choice = 'r'
            
        if choice == 's':
            import shutil
            print(f"[INFO] Clearing previous artifacts...", file=sys.stderr)
            shutil.rmtree(artifacts_dir)
            artifacts_dir.mkdir(parents=True, exist_ok=True)
            rebuild_context = True
        else:
            print(f"[INFO] Resuming session...", file=sys.stderr)
    else:
        artifacts_dir.mkdir(parents=True, exist_ok=True)
        rebuild_context = True

    project_name = project_dir.name
    vault_project_dir = Path(args.vault_path).expanduser().resolve() if args.vault_path else (Path.home() / "contx" / project_name)
    
    # Clear vault if forced rebuild
    if rebuild_context and vault_project_dir.exists():
        import shutil
        print(f"[INFO] Clearing vault at {vault_project_dir}...", file=sys.stderr)
        shutil.rmtree(vault_project_dir)
    
    vault_project_dir.mkdir(parents=True, exist_ok=True)

    # Load .env
    load_env(script_dir / ".env")

    # Select AI provider
    provider = select_ai_provider()

    # Set common env
    os.environ["APP_ROOT"] = str(project_dir)
    os.environ["OUTPUT_DIR"] = str(project_dir)
    os.environ["AI_PROVIDER"] = provider

    # ─── Phase 0: Context Analysis (no AI calls) ──────────────────────────
    print("\n=== Phase 0: Context Analysis ===", file=sys.stderr)

    # Build vault
    if not run_context_creator(project_dir, vault_project_dir, rebuild_context):
        print("[FATAL] Vault creation failed.", file=sys.stderr)
        sys.exit(1)

    # Move code-index and code-summary to artifacts/ if they were generated in the vault
    for filename in ["code-index.json", "code-summary.json"]:
        vault_file = vault_project_dir / filename
        target_file = artifacts_dir / filename
        if vault_file.exists():
            if target_file.exists():
                target_file.unlink()
            vault_file.rename(target_file)

    vault = Vault(vault_project_dir, project_root=project_dir)

    # Stack detector + context agent
    for name, agent_factory in [
        ("Stack Detector", lambda: StackDetector(vault, str(project_dir))),
        ("Context Agent", lambda: ContextAgent(vault, str(project_dir))),
    ]:
        print(f"\n--- Running: {name} ---", file=sys.stderr)
        agent = agent_factory()
        exit_code = agent.run()
        if exit_code != 0:
            print(f"[FATAL] {name} failed with exit code {exit_code}", file=sys.stderr)
            sys.exit(1)

    # ─── Phase 1: Planning Agents (AI calls) ──────────────────────────────
    print("\n=== Phase 1: Planning ===", file=sys.stderr)

    for name, agent_factory in [
        ("Product Owner", lambda: POAgent(vault, str(project_dir), args.task)),
        ("Business Analyst", lambda: BAAgent(vault, str(project_dir), args.task)),
        ("Architect", lambda: ArchitectAgent(vault, str(project_dir), args.task)),
        ("Devflow", lambda: DevflowAgent(vault, str(project_dir), args.task)),
    ]:
        print(f"\n--- Running: {name} ---", file=sys.stderr)
        agent = agent_factory()
        exit_code = agent.run()
        if exit_code != 0:
            print(f"[FATAL] {name} failed with exit code {exit_code}", file=sys.stderr)
            sys.exit(1)

    # ─── Phase 1.5: Dev Context Generation ────────────────────────────────
    req_path = artifacts_dir / "pipeline-requirements.json"
    dev_context_path = artifacts_dir / "dev_context" / "dev-context.md"
    if not run_dev_context_creator(project_dir, req_path, dev_context_path):
        print("[WARN] Dev context generation failed, continuing with vault only.", file=sys.stderr)

    # ─── User Review Gate ──────────────────────────────────────────────────
    print_pipeline_summary(req_path)

    print("\nReview pipeline-requirements.json manually if needed.", file=sys.stderr)
    try:
        confirm = input("\nProceed with Phase 2 (dev/test loop)? [y/n]: ").strip().lower()
    except EOFError:
        confirm = "n"

    if confirm != "y":
        print("[INFO] Pipeline cancelled by user. Edit pipeline-requirements.json and rerun.", file=sys.stderr)
        sys.exit(0)

    # ─── Phase 2: Dev/Test Loop ────────────────────────────────────────────
    req = json.loads(req_path.read_text())
    if req.get("completed"):
        print("\n[SUCCESS] Pipeline already completed successfully. Skipping Phase 2.", file=sys.stderr)
        sys.exit(0)

    print("\n=== Phase 2: Dev/Test Loop ===", file=sys.stderr)
    loop_count = req.get("iteration", 0)

    while True:
        display_iteration = loop_count + 1
        print(f"\n--- Iteration {display_iteration} ---", file=sys.stderr)
        
        # Dev agent
        print("--- Running: Dev Agent ---", file=sys.stderr)
        dev = DevAgent(vault, str(project_dir), args.task)
        exit_code = dev.run()
        if exit_code != 0:
            print(f"[FATAL] Dev agent failed with exit code {exit_code}", file=sys.stderr)
            sys.exit(1)

        # QA agent
        print("--- Running: QA Agent ---", file=sys.stderr)
        qa = QAAgent(vault, str(project_dir), args.task)
        exit_code = qa.run()
        if exit_code != 0:
            print(f"[FATAL] QA agent failed with exit code {exit_code}", file=sys.stderr)
            sys.exit(1)

        # ─── QA Review Gate ────────────────────────────────────────────────
        qa_report_path = artifacts_dir / "qa-report.json"
        if qa_report_path.exists():
            try:
                qa_report = json.loads(qa_report_path.read_text())
                bugs = qa_report.get("bugs", [])
                print("\n" + "=" * 50, file=sys.stderr)
                print("=== QA Report Review ===", file=sys.stderr)
                print(f"Summary: {qa_report.get('summary', 'N/A')}", file=sys.stderr)
                print(f"Bugs found: {len(bugs)}", file=sys.stderr)
                for i, bug in enumerate(bugs, 1):
                    edit_status = "[FIX]" if bug.get("edit", True) else "[SKIP]"
                    print(f"  {i}. {edit_status} [{bug.get('severity', '???')}] {bug.get('title')}", file=sys.stderr)
                print("=" * 50, file=sys.stderr)
                
                if bugs:
                    print(f"\nReview {qa_report_path} to toggle 'edit' values if needed.", file=sys.stderr)
                    try:
                        qa_confirm = input("Loop back to Dev Agent for fixes? [y/n]: ").strip().lower()
                    except EOFError:
                        qa_confirm = "n"
                    
                    if qa_confirm != "y":
                        print("[INFO] User declined loop back. Running final verification...", file=sys.stderr)
                        tester = TestingAgent(vault, str(project_dir), args.task)
                        try:
                            tester.run()
                            print("\n[SUCCESS] Final verification passed.", file=sys.stderr)
                        except SystemExit as e:
                            if e.code == 0:
                                print("\n[SUCCESS] Final verification passed.", file=sys.stderr)
                            else:
                                print(f"\n[WARN] Final verification found issues (Exit {e.code}), but loop terminated by user.", file=sys.stderr)
                        
                        req["completed"] = True
                        req_path.write_text(json.dumps(req, indent=2))
                        print("\n[INFO] Pipeline finished.", file=sys.stderr)
                        sys.exit(0)
                else:
                    print("\n[INFO] No QA bugs found. Proceeding to final verification...", file=sys.stderr)
            except Exception as e:
                print(f"[WARN] Error reading QA report for review: {e}", file=sys.stderr)

        # Test agent (Normal flow: runs and handles loop back)
        print("--- Running: Testing Agent ---", file=sys.stderr)
        tester = TestingAgent(vault, str(project_dir), args.task)
        try:
            tester.run()
            # If we get here, tests passed (exit 0)
            print("\n[SUCCESS] Pipeline completed. All tests passed.", file=sys.stderr)
            req["completed"] = True
            req["iteration"] = loop_count + 1
            req_path.write_text(json.dumps(req, indent=2))
            break
        except SystemExit as e:
            if e.code == 0:
                print("\n[SUCCESS] Pipeline completed. All tests passed.", file=sys.stderr)
                req["completed"] = True
                req["iteration"] = loop_count + 1
                req_path.write_text(json.dumps(req, indent=2))
                break
            elif e.code == 2:
                print("[INFO] Tests failed, fix task added. Looping back...", file=sys.stderr)
                # Increment iteration only after full cycle
                loop_count += 1
                req["iteration"] = loop_count
                req_path.write_text(json.dumps(req, indent=2))
                # Refresh req as it was updated by TestingAgent (devflow)
                req = json.loads(req_path.read_text())
                continue
            else:
                print(f"[FATAL] Testing agent failed with exit code {e.code}", file=sys.stderr)
                sys.exit(1)

    print("\n[INFO] Pipeline finished.", file=sys.stderr)


if __name__ == "__main__":
    main()