#!/usr/bin/env python3
import argparse
import os
import subprocess
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
WORKFLOW = ROOT / 'harness' / 'workflow.stages.yml'
STATUS_EXT = ROOT / '.pi' / 'extensions' / 'orquestra-status'
WEB_FETCH_EXT = ROOT / '.pi' / 'extensions' / 'orquestra-web-fetch.ts'
ENGRAM_EXT = Path.home() / '.pi' / 'agent' / 'npm' / 'node_modules' / 'gentle-engram' / 'index.ts'


def parse_stages():
    stages = {}
    current = None
    section = None
    for raw in WORKFLOW.read_text(encoding='utf-8').splitlines():
        line = raw.rstrip()
        if line.startswith('  - name:'):
            current = {'input': [], 'output': [], 'post_actions': []}
            current['name'] = line.split(':', 1)[1].strip()
            stages[current['name']] = current
            section = None
            continue
        if current is None:
            continue
        if line.startswith('    owner:'):
            current['owner'] = line.split(':', 1)[1].strip()
        elif line.startswith('    input:'):
            section = 'input'
        elif line.startswith('    output:'):
            section = 'output'
        elif line.startswith('    post_actions:'):
            section = 'post_actions'
        elif section and line.startswith('      - '):
            current[section].append(line.split('- ', 1)[1].strip())
        elif line.startswith('    '):
            section = None
    return stages


def expand(paths, feature_id):
    return [p.replace('<feature_id>', feature_id) for p in paths]


def prompt_for(stage, feature_id):
    inputs = expand(stage.get('input', []), feature_id)
    outputs = expand(stage.get('output', []), feature_id)
    post_actions = expand(stage.get('post_actions', []), feature_id)
    owner = stage.get('owner', 'leader')
    
    post_actions_text = ""
    if post_actions:
        post_actions_text = f"""
Post-actions (execute after completing the stage):
{chr(10).join(f'- {a}' for a in post_actions)}
"""
    
    return f"""You are the Orquestra stage agent '{owner}' for feature '{feature_id}'.

Fresh-process rule: do not rely on previous chat/session context. Use only the files listed here, explicit Engram memories you choose after a narrow search, and repository evidence you read yourself.

First read these harness rules:
- AGENTS.md
- harness/agents.matrix.yml
- harness/workflow.stages.yml
- harness/contracts/handoff.md

Stage input paths:
{chr(10).join(f'- {p}' for p in inputs) or '- none'}

Stage output paths:
{chr(10).join(f'- {p}' for p in outputs) or '- none'}
{post_actions_text}
Rules:
- Before work, run: python3 scripts/agent_status.py set --feature-id {feature_id} --stage {stage['name']} --agent {owner} --state running --action "Running {stage['name']}"
- Read only the stage inputs that exist. If a required input is missing, write a blocked artifact when possible and answer blocked -> <path>.
- Do not carry or summarize previous chat. Previous stage context is on disk in the declared input artifacts.
- Write evidence to the declared output path(s). Do not return code in chat.
    - NEVER edit backlog/features.json directly. Use scripts/close_feature.py to close features.
- Before finishing, save to Engram only if this stage produced durable knowledge: a decision, bugfix, non-obvious discovery, reusable convention, or configuration change. Do not save routine progress, command output, or artifact summaries; those belong in the output artifact.
- Finish by updating runtime status to done or blocked.
- Final response must be exactly: done -> <path> or blocked -> <path>.
"""


def main():
    parser = argparse.ArgumentParser(description='Run one Orquestra stage in a fresh Pi process with minimal declared context.')
    parser.add_argument('stage')
    parser.add_argument('--feature-id', required=True)
    parser.add_argument('--dry-run', action='store_true')
    args = parser.parse_args()

    stages = parse_stages()
    if args.stage not in stages:
        print(f"[FAIL] Unknown stage: {args.stage}", file=sys.stderr)
        print(f"Allowed: {', '.join(stages)}", file=sys.stderr)
        return 2

    cmd = [
        'pi', '--no-session', '--no-context-files', '--no-extensions', '--no-skills',
        '-e', str(ENGRAM_EXT), '-e', str(STATUS_EXT), '-e', str(WEB_FETCH_EXT),
        '-p', prompt_for(stages[args.stage], args.feature_id),
    ]
    if args.dry_run:
        print(' '.join(cmd[:-1]))
        print('\n--- prompt ---')
        print(cmd[-1])
        return 0

    missing = []
    if not ENGRAM_EXT.is_file():
        missing.append(str(ENGRAM_EXT))
    if not (STATUS_EXT / 'index.ts').is_file():
        missing.append(str(STATUS_EXT / 'index.ts'))
    if not WEB_FETCH_EXT.is_file():
        missing.append(str(WEB_FETCH_EXT))
    if missing:
        print('[FAIL] Missing Pi runtime extension(s):', file=sys.stderr)
        for path in missing:
            print(f'  - {path}', file=sys.stderr)
        print('Run scripts/install.sh for this project and install gentle-engram.', file=sys.stderr)
        return 1
    return subprocess.run(cmd, cwd=ROOT, env=os.environ.copy()).returncode


if __name__ == '__main__':
    raise SystemExit(main())