#!/usr/bin/env python3
"""
GPT Agent Orchestrator - CLI Entry Point
==========================================

A standalone CLI tool that runs a queue of tasks through an AI coding
agent using a hook-driven chain. Tasks can be manual markdown files or
auto-generated by GPT.

Usage:
  Manual mode:  chainsmith ~/projects/my-app
  GPT mode:     chainsmith ~/projects/my-app --gpt -d "Build a REST API"
"""

import argparse
import json
import os
import signal
import subprocess
import sys
import time
from datetime import datetime, timezone
from pathlib import Path

from .state import State
from .config import (
    load_config,
    save_config,
    merge_env_overrides,
    mask_secrets,
    get_config_path,
    set_config_value,
    get_config_value,
)
from .logging import Logger, is_rich_available

# Only import GPT modules if needed (they require openai)
GPTPlanner = None


def lazy_import_gpt():
    global GPTPlanner
    if GPTPlanner is None:
        from .gpt_planner import GPTPlanner as _GPTPlanner
        GPTPlanner = _GPTPlanner


# ──────────────────────────────────────────────
# CLI
# ──────────────────────────────────────────────


def parse_args(args=None):
    parser = argparse.ArgumentParser(
        description="GPT Agent Orchestrator: Run tasks through an AI coding agent with hook-driven chaining.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Manual mode — run pre-created task files
  chainsmith ~/projects/my-app

  # GPT mode — GPT generates and manages tasks
  chainsmith ~/projects/my-app --gpt -d "Build a todo app with auth"

  # Resume after interruption
  chainsmith ~/projects/my-app --resume

  # Configuration commands
  chainsmith init                    # Interactive setup
  chainsmith config show             # Show current config
  chainsmith config set planner.model gpt-4  # Set config value
        """,
    )

    # Subparsers for config commands
    subparsers = parser.add_subparsers(dest="command", help="Available commands")

    # Main run command (default)
    run_parser = subparsers.add_parser("run", help="Run orchestration (default)", add_help=False)
    run_parser.add_argument("project_dir", help="Path to the target project directory")

    # GPT mode
    run_parser.add_argument("--gpt", action="store_true", help="Enable GPT orchestration mode")
    run_parser.add_argument("-d", "--description", help="Project description (required with --gpt)")
    run_parser.add_argument("--gpt-model", help="GPT model (overrides config) [deprecated: use --planner-model]")
    run_parser.add_argument("--planner-provider", help="Planner provider (openai, anthropic, azure_openai, deepseek, openai_compatible)")
    run_parser.add_argument("--planner-model", help="Planner model (overrides config)")
    run_parser.add_argument("--planner-base-url", help="Base URL for openai_compatible or deepseek provider")
    run_parser.add_argument("--azure-endpoint", help="Azure OpenAI endpoint URL")
    run_parser.add_argument("--azure-deployment", help="Azure OpenAI deployment name")
    run_parser.add_argument("--azure-api-version", help="Azure OpenAI API version (default: 2024-02-01)")
    run_parser.add_argument("--max-batches", type=int, default=5, help="Max GPT planning rounds (default: 5)")
    run_parser.add_argument("--batch-size", type=int, default=5, help="Tasks per batch (default: 5)")

    # Agent options
    run_parser.add_argument("--agent-model", help="AI agent model (overrides config)")
    run_parser.add_argument("--max-cost-usd", type=float, help="Cost ceiling in USD (overrides config)")
    run_parser.add_argument("--task-timeout", type=int, default=600, help="Max seconds per task (default: 600)")

    # Other
    run_parser.add_argument("--resume", action="store_true", help="Resume from existing state")
    run_parser.add_argument("--verbose", action="store_true", help="Detailed logging")
    run_parser.add_argument("--openai-api-key", help="OpenAI API key (or use env var)")
    run_parser.add_argument("--provider", help="Planner provider (openai, anthropic, azure_openai, deepseek, openai_compatible)")
    run_parser.add_argument("--base-url", help="Base URL for openai_compatible provider")

    # Init command
    init_parser = subparsers.add_parser("init", help="Interactive configuration setup")

    # Config command
    config_parser = subparsers.add_parser("config", help="Manage configuration")
    config_subparsers = config_parser.add_subparsers(dest="config_command", help="Config operations")

    config_show_parser = config_subparsers.add_parser("show", help="Show current configuration")
    config_show_parser.add_argument("--no-mask", action="store_true", help="Show secrets unmasked")

    config_set_parser = config_subparsers.add_parser("set", help="Set configuration value")
    config_set_parser.add_argument("key", help="Config key in dot notation (e.g., planner.model)")
    config_set_parser.add_argument("value", help="Value to set")

    # Parse args
    parsed = parser.parse_args(args)

    # If no command specified and first arg is not a command, assume "run"
    if parsed.command is None and args and len(args) > 0 and not args[0] in ("init", "config"):
        # Reparse with "run" command injected
        new_args = ["run"] + (args if args else [])
        parsed = parser.parse_args(new_args)

    return parsed


# ──────────────────────────────────────────────
# CLAUDE.md management
# ──────────────────────────────────────────────


class DocsManager:
    """Manages prework.md and postwork.md template files in the target project."""

    PREWORK_TEMPLATE = """\
# Pre-work Documentation

Before starting any task, create or update a documentation file at:
  docs/<area>/<feature>.md

## What to document

1. **Task Overview**: What you are about to do
2. **Plan**: Step-by-step approach you intend to take
3. **Files Affected**: Which files you expect to create or modify
4. **Dependencies**: Any libraries, APIs, or other components involved
5. **Assumptions**: Any assumptions you are making

## Format

```markdown
# <Feature Name>

## Pre-work

**Date**: YYYY-MM-DD
**Task**: <task title>

### Plan
- Step 1: ...
- Step 2: ...

### Files to be affected
- path/to/file1
- path/to/file2

### Dependencies
- ...

### Assumptions
- ...
```
"""

    POSTWORK_TEMPLATE = """\
# Post-work Documentation

After completing any task, update the same documentation file at:
  docs/<area>/<feature>.md

Append a Post-work section below the existing Pre-work section.

## What to document

1. **What was done**: Actual changes made
2. **Files Modified**: Exact files created or modified with brief descriptions
3. **Key Decisions**: Any architectural or design decisions made during implementation
4. **Issues Encountered**: Problems faced and how they were resolved
5. **Status**: Whether the task was fully completed or has remaining work

## Format

Append this to the existing doc file:

```markdown
## Post-work

**Completed**: YYYY-MM-DD
**Status**: Completed | Partial | Failed

### Changes Made
- Description of change 1
- Description of change 2

### Files Modified
- `path/to/file1` — what was changed
- `path/to/file2` — what was changed

### Key Decisions
- Decision 1 and reasoning
- Decision 2 and reasoning

### Issues & Resolutions
- Issue: ... → Resolution: ...
```
"""

    def __init__(self, project_dir: Path):
        self.project_dir = project_dir
        self.prework_path = project_dir / "prework.md"
        self.postwork_path = project_dir / "postwork.md"
        self.docs_dir = project_dir / "docs"

    def setup(self):
        """Create prework.md, postwork.md, and docs/ directory if they don't exist."""
        try:
            self.docs_dir.mkdir(parents=True, exist_ok=True)
        except OSError as e:
            print(f"  WARNING: Could not create docs/ directory: {e}")
            return

        if not self.prework_path.exists():
            try:
                self.prework_path.write_text(self.PREWORK_TEMPLATE)
                print(f"  Created prework.md")
            except OSError as e:
                print(f"  WARNING: Could not create prework.md: {e}")
        else:
            print(f"  Found existing prework.md")

        if not self.postwork_path.exists():
            try:
                self.postwork_path.write_text(self.POSTWORK_TEMPLATE)
                print(f"  Created postwork.md")
            except OSError as e:
                print(f"  WARNING: Could not create postwork.md: {e}")
        else:
            print(f"  Found existing postwork.md")


class ContextMDManager:
    """Manages the context file (CLAUDE.md) in the target project directory."""

    def __init__(self, project_dir: Path):
        self.context_md = project_dir / "CLAUDE.md"
        self.backup_path = project_dir / ".claude" / "CLAUDE.md.orchestrator_backup"
        self.original_existed = self.context_md.exists()

    def backup(self):
        """Back up existing context file before orchestration."""
        if self.original_existed:
            self.backup_path.parent.mkdir(parents=True, exist_ok=True)
            self.backup_path.write_text(self.context_md.read_text())
            print(f"  Backed up existing context file")

    def restore(self):
        """Restore original context file after orchestration."""
        if self.backup_path.exists():
            self.context_md.write_text(self.backup_path.read_text())
            self.backup_path.unlink()
            print(f"  Restored original CLAUDE.md")
        elif not self.original_existed and self.context_md.exists():
            self.context_md.unlink()
            print(f"  Removed orchestrator CLAUDE.md")

    def write_initial(self, description: str = "", cumulative_summary: str = ""):
        """Write initial CLAUDE.md for the first task."""
        lines = [
            "# Orchestrator Context",
            "",
            "You are executing a task as part of an automated pipeline.",
            "Complete the task in your prompt thoroughly and precisely.",
            "Do not ask clarifying questions — make reasonable assumptions.",
            "Focus on writing working code.",
            "",
            "## Documentation Requirements",
            "",
            "IMPORTANT: Before and after each task, you MUST create documentation.",
            "",
            "1. **Before starting work**: Read `prework.md` in the project root for instructions.",
            "   Create a doc file at `docs/<area>/<feature>.md` with a Pre-work section",
            "   documenting your plan. Choose `<area>` based on the domain (e.g., frontend,",
            "   backend, auth, database, api, config) and `<feature>` based on what you're building.",
            "",
            "2. **After completing work**: Read `postwork.md` in the project root for instructions.",
            "   Update the SAME doc file with a Post-work section documenting what was actually done.",
            "",
            "If a doc file already exists for this feature, append to it rather than overwriting.",
            "",
        ]

        if description:
            lines.extend([
                "## Project Description",
                description,
                "",
            ])

        lines.extend([
            "## Progress So Far",
            "(No tasks completed yet)",
            "",
        ])

        if cumulative_summary:
            lines.extend([
                "## Cumulative Summary",
                cumulative_summary,
                "",
            ])

        # Preserve original content
        if self.backup_path.exists():
            lines.extend([
                "---",
                "## Original Project Instructions",
                self.backup_path.read_text(),
            ])

        self.context_md.write_text("\n".join(lines))

    def rewrite_with_summary(self, description: str, cumulative_summary: str):
        """Rewrite CLAUDE.md with updated cumulative summary (between batches)."""
        self.write_initial(description=description, cumulative_summary=cumulative_summary)


# ──────────────────────────────────────────────
# Hook installation
# ──────────────────────────────────────────────


def get_orchestrator_dir():
    """Get the orchestrator installation directory (where this package is installed)."""
    return Path(__file__).parent.resolve()


def install_hook(project_dir: Path):
    """Install the SessionEnd hook into the project's settings.local.json."""
    config_dir = project_dir / ".claude"
    config_dir.mkdir(exist_ok=True)

    settings_path = config_dir / "settings.local.json"
    hook_handler = str(get_orchestrator_dir() / "hook_handler.py")

    # Read existing settings if present
    existing = {}
    if settings_path.exists():
        try:
            existing = json.loads(settings_path.read_text())
        except json.JSONDecodeError:
            existing = {}

    # Merge permissions
    permissions = existing.get("permissions", {})
    allow = set(permissions.get("allow", []))
    allow.update(["Bash", "Read", "Write", "Edit", "Glob", "Grep"])
    permissions["allow"] = sorted(allow)

    # Merge hook (preserve existing hooks for other events)
    existing["permissions"] = permissions
    hooks = existing.get("hooks", {})
    hooks["SessionEnd"] = [
        {
            "hooks": [
                {
                    "type": "command",
                    "command": hook_handler,
                    "timeout": 30,
                }
            ]
        }
    ]
    existing["hooks"] = hooks

    settings_path.write_text(json.dumps(existing, indent=2))
    print(f"  Installed SessionEnd hook → {hook_handler}")

    # Write reference file so hook_handler can find state.json
    ref_path = config_dir / "orchestrator_ref.json"
    ref_path.write_text(
        json.dumps(
            {
                "state_file": str(get_orchestrator_dir() / "state.json"),
                "orchestrator_dir": str(get_orchestrator_dir()),
            },
            indent=2,
        )
    )


def uninstall_hook(project_dir: Path):
    """Remove the orchestrator hook from settings.local.json."""
    settings_path = project_dir / ".claude" / "settings.local.json"
    ref_path = project_dir / ".claude" / "orchestrator_ref.json"

    if settings_path.exists():
        try:
            settings = json.loads(settings_path.read_text())
            settings.pop("hooks", None)
            settings_path.write_text(json.dumps(settings, indent=2))
            print(f"  Removed hook from settings.local.json")
        except json.JSONDecodeError:
            pass

    if ref_path.exists():
        ref_path.unlink()


# ──────────────────────────────────────────────
# Task discovery
# ──────────────────────────────────────────────


def discover_manual_tasks() -> list[dict]:
    """Discover task files in the tasks/ directory, sorted numerically."""
    orchestrator_dir = get_orchestrator_dir()
    tasks_dir = orchestrator_dir / "tasks"
    if not tasks_dir.exists():
        return []

    task_files = []
    for f in tasks_dir.iterdir():
        if f.suffix == ".md" and f.stem.isdigit():
            task_files.append((int(f.stem), f))

    # Sort by number
    task_files.sort(key=lambda x: x[0])

    tasks = []
    for num, filepath in task_files:
        content = filepath.read_text().strip()
        # Extract title from first heading or filename
        title = f"Task {num}"
        for line in content.split("\n"):
            line = line.strip()
            if line.startswith("# "):
                title = line[2:].strip()
                break

        tasks.append({
            "title": title,
            "file": f"tasks/{filepath.name}",
        })

    return tasks


# ──────────────────────────────────────────────
# Agent process management
# ──────────────────────────────────────────────


def launch_agent_task(project_dir: str, task_file: str, config: dict):
    """Launch an AI agent process for a task."""
    orchestrator_dir = get_orchestrator_dir()
    task_path = orchestrator_dir / task_file
    if not task_path.exists():
        print(f"  ERROR: Task file not found: {task_path}")
        return None

    prompt = task_path.read_text().strip()
    if not prompt:
        print(f"  ERROR: Task file is empty: {task_path}")
        return None

    cmd = [
        "claude",
        "-p",
        prompt,
        "--output-format",
        "json",
    ]

    agent_model = config.get("agent_model")
    if agent_model:
        cmd.extend(["--model", agent_model])

    proc = subprocess.Popen(
        cmd,
        cwd=project_dir,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
        start_new_session=True,
    )

    return proc


# ──────────────────────────────────────────────
# Main orchestration loop
# ──────────────────────────────────────────────


class Orchestrator:
    def __init__(self, args):
        self.args = args
        self.project_dir = Path(args.project_dir).resolve()
        orchestrator_dir = get_orchestrator_dir()
        self.state_path = str(orchestrator_dir / "state.json")
        self.state = State(self.state_path)
        self.context_md = ContextMDManager(self.project_dir)
        self.docs_manager = DocsManager(self.project_dir)
        self.gpt = None
        self.shutting_down = False

        # Initialize logger with log directory in .claude/logs
        log_dir = self.project_dir / ".claude" / "logs"
        self.logger = Logger(log_dir=log_dir, enable_rich=not args.verbose if hasattr(args, 'verbose') else True)

        # Load and merge config
        self.config = load_config()
        self.config = merge_env_overrides(self.config)

        # Apply CLI flag overrides (flags have highest priority)
        self._apply_flag_overrides()

        # Set up signal handler for graceful shutdown
        signal.signal(signal.SIGINT, self._handle_shutdown)
        signal.signal(signal.SIGTERM, self._handle_shutdown)

    def _apply_flag_overrides(self):
        """Apply CLI flag overrides to config (flags have highest priority)."""
        # Agent model
        if hasattr(self.args, 'agent_model') and self.args.agent_model:
            self.config["agent"]["model"] = self.args.agent_model

        # Max cost
        if hasattr(self.args, 'max_cost_usd') and self.args.max_cost_usd is not None:
            self.config["planner"]["max_cost_usd"] = self.args.max_cost_usd

        # GPT model (deprecated, but still supported)
        if hasattr(self.args, 'gpt_model') and self.args.gpt_model:
            self.config["planner"]["model"] = self.args.gpt_model

        # Planner provider
        if hasattr(self.args, 'planner_provider') and self.args.planner_provider:
            self.config["planner"]["provider"] = self.args.planner_provider
        # Legacy provider flag
        elif hasattr(self.args, 'provider') and self.args.provider:
            self.config["planner"]["provider"] = self.args.provider

        # Planner model
        if hasattr(self.args, 'planner_model') and self.args.planner_model:
            self.config["planner"]["model"] = self.args.planner_model

        # Planner base URL
        if hasattr(self.args, 'planner_base_url') and self.args.planner_base_url:
            self.config["planner"]["base_url"] = self.args.planner_base_url
        # Legacy base_url flag
        elif hasattr(self.args, 'base_url') and self.args.base_url:
            self.config["planner"]["base_url"] = self.args.base_url

        # API key (legacy flag)
        if hasattr(self.args, 'openai_api_key') and self.args.openai_api_key:
            self.config["planner"]["api_key"] = self.args.openai_api_key

        # Azure-specific flags
        if hasattr(self.args, 'azure_endpoint') and self.args.azure_endpoint:
            if "azure" not in self.config:
                self.config["azure"] = {}
            self.config["azure"]["endpoint"] = self.args.azure_endpoint

        if hasattr(self.args, 'azure_deployment') and self.args.azure_deployment:
            if "azure" not in self.config:
                self.config["azure"] = {}
            self.config["azure"]["deployment"] = self.args.azure_deployment

        if hasattr(self.args, 'azure_api_version') and self.args.azure_api_version:
            if "azure" not in self.config:
                self.config["azure"] = {}
            self.config["azure"]["api_version"] = self.args.azure_api_version

    def _handle_shutdown(self, signum, frame):
        """Graceful shutdown: set flag, let the polling loop handle cleanup.

        Avoids file I/O in signal handler to prevent deadlocks if signal
        arrives while a file lock is held.
        """
        if self.shutting_down:
            # Second interrupt — force exit
            print("\nForce quitting...")
            sys.exit(1)

        self.shutting_down = True
        print("\n\nShutting down gracefully...")
        print("  (Current task will finish but no new tasks will start)")

    def validate(self):
        """Validate prerequisites."""
        if not self.project_dir.is_dir():
            self.logger.error(f"Project directory not found: {self.project_dir}")
            sys.exit(1)

        # Check agent CLI
        try:
            result = subprocess.run(
                ["claude", "--version"],
                capture_output=True,
                text=True,
                timeout=10,
            )
            if result.returncode != 0:
                self.logger.error("AI agent CLI not found or not working")
                sys.exit(1)
        except (FileNotFoundError, subprocess.TimeoutExpired):
            self.logger.error("AI agent CLI not found. Install it first.")
            sys.exit(1)

        if self.args.gpt:
            if not self.args.description:
                self.logger.error("--description is required with --gpt mode")
                sys.exit(1)

            # Get API key from config
            api_key = self.config["planner"].get("api_key")
            if not api_key:
                self.logger.error("API key not found in config or environment")
                self.logger.info("Run 'chainsmith init' to configure, or set environment variable")
                sys.exit(1)

            lazy_import_gpt()
            provider = self.config["planner"]["provider"]
            model = self.config["planner"]["model"]

            # Initialize GPTPlanner with provider-specific settings
            if provider == "azure_openai":
                azure_config = self.config.get("azure", {})
                self.gpt = GPTPlanner(
                    api_key=api_key,
                    model=model,
                    provider=provider,
                    azure_endpoint=azure_config.get("endpoint"),
                    azure_deployment=azure_config.get("deployment"),
                    azure_api_version=azure_config.get("api_version"),
                )
            elif provider == "openai_compatible":
                self.gpt = GPTPlanner(
                    api_key=api_key,
                    model=model,
                    provider=provider,
                    base_url=self.config["planner"].get("base_url"),
                )
            else:
                self.gpt = GPTPlanner(
                    api_key=api_key,
                    model=model,
                    provider=provider,
                )

    def run(self):
        """Main orchestration loop."""
        self.validate()

        # Display header
        self.logger.header("GPT Agent Orchestrator")

        # Display configuration
        config_lines = [
            f"Project:  {self.project_dir}",
            f"Mode:     {'GPT' if self.args.gpt else 'Manual'}",
            f"Agent:    {self.config['agent']['model']}",
        ]

        if self.args.gpt:
            config_lines.extend([
                f"Provider: {self.config['planner']['provider']}",
                f"Model:    {self.config['planner']['model']}",
                f"Batches:  {self.args.max_batches} x {self.args.batch_size} tasks",
            ])

        config_lines.append(f"Cost cap: ${self.config['planner']['max_cost_usd']:.2f}")

        if is_rich_available():
            config_lines.append("Rich UI:  Enabled")

        for line in config_lines:
            self.logger.print(f"  {line}")

        print()  # Add spacing

        if self.args.resume:
            self._resume()
        elif self.args.gpt:
            self._run_gpt_mode()
        else:
            self._run_manual_mode()

    def _run_manual_mode(self):
        """Run tasks from tasks/ directory."""
        tasks = discover_manual_tasks()
        if not tasks:
            self.logger.error("No task files found in tasks/ directory.")
            self.logger.info("Create numbered markdown files: tasks/1.md, tasks/2.md, ...")
            sys.exit(1)

        self.logger.info(f"Found {len(tasks)} task(s):")

        # Display tasks in a table if rich is available
        if is_rich_available():
            rows = [[str(i + 1), t["title"], t["file"]] for i, t in enumerate(tasks)]
            self.logger.table("Planned Tasks", ["#", "Title", "File"], rows)
        else:
            for t in tasks:
                print(f"  • {t['file']}: {t['title']}")
            print()

        config = {
            "agent_model": self.config["agent"]["model"],
            "max_cost_usd": self.config["planner"]["max_cost_usd"],
            "task_timeout": self.args.task_timeout,
        }

        orchestrator_dir = get_orchestrator_dir()

        # Initialize state
        State.create_initial(
            path=self.state_path,
            project_dir=str(self.project_dir),
            orchestrator_dir=str(orchestrator_dir),
            tasks=tasks,
            config=config,
            mode="manual",
        )

        # Set up project
        self.context_md.backup()
        self.context_md.write_initial()
        self.docs_manager.setup()
        install_hook(self.project_dir)

        # Mark first task as running with timestamp
        self.state.mark_task_running(0)

        # Launch first task
        first_task = tasks[0]
        self.logger.info(f"Starting task 1/{len(tasks)}: {first_task['title']}")
        launch_agent_task(str(self.project_dir), first_task["file"], config)

        # Poll until done
        self._poll_until_complete()

        # Cleanup
        self._cleanup()

    def _run_gpt_mode(self):
        """Run with GPT generating batches of tasks."""
        description = self.args.description

        # Read description from file if it's a path
        if os.path.isfile(description):
            description = Path(description).read_text().strip()

        config = {
            "agent_model": self.config["agent"]["model"],
            "max_cost_usd": self.config["planner"]["max_cost_usd"],
            "task_timeout": self.args.task_timeout,
            "gpt_model": self.config["planner"]["model"],
        }

        self.context_md.backup()
        self.docs_manager.setup()
        install_hook(self.project_dir)

        cumulative_summary = ""
        previous_results = []
        total_task_index = 0

        orchestrator_dir = get_orchestrator_dir()

        for batch_num in range(1, self.args.max_batches + 1):
            if self.shutting_down:
                break

            self.logger.section(f"Batch {batch_num}/{self.args.max_batches}")

            # Ask GPT to plan tasks
            self.logger.info(f"Asking GPT to plan batch {batch_num}...")
            gpt_tasks = self.gpt.plan_batch(
                description=description,
                batch_num=batch_num,
                max_batches=self.args.max_batches,
                batch_size=self.args.batch_size,
                previous_results=previous_results,
                cumulative_summary=cumulative_summary,
            )

            if not gpt_tasks:
                self.logger.success("GPT says project is complete!")
                break

            # Write tasks to files
            task_entries = self.gpt.write_tasks_to_files(
                gpt_tasks, str(orchestrator_dir), start_index=total_task_index
            )

            self.logger.info(f"Generated {len(task_entries)} tasks:")

            # Display generated tasks in a table if rich is available
            if is_rich_available():
                rows = [[str(i + 1), t["title"]] for i, t in enumerate(task_entries)]
                self.logger.table("Generated Tasks", ["#", "Title"], rows)
            else:
                for t in task_entries:
                    print(f"    • {t['title']}")

            # Create or update state
            if batch_num == 1:
                State.create_initial(
                    path=self.state_path,
                    project_dir=str(self.project_dir),
                    orchestrator_dir=str(orchestrator_dir),
                    tasks=task_entries,
                    config=config,
                    mode="gpt",
                )
            else:
                # Add new batch tasks to existing state
                state_data = self.state.read()
                new_indices = []
                for t in task_entries:
                    idx = len(state_data["tasks"])
                    new_indices.append(idx)
                    state_data["tasks"].append({
                        "index": idx,
                        "title": t["title"],
                        "file": t["file"],
                        "status": "pending",
                        "session_id": None,
                        "transcript_path": None,
                        "result_summary": None,
                        "files_modified": [],
                        "cost_usd": 0,
                    })
                state_data["current_batch"] = batch_num
                state_data["current_task_index"] = new_indices[0]
                state_data["batches"][str(batch_num)] = {
                    "task_indices": new_indices,
                    "status": "running",
                }
                self.state.write(state_data)

            # Update CLAUDE.md
            self.context_md.rewrite_with_summary(description, cumulative_summary)

            # Mark first task of batch as running
            state_data = self.state.read()
            first_idx = state_data["current_task_index"]
            # Mark first task as running with timestamp
            self.state.mark_task_running(first_idx)

            # Launch first task
            first_task = task_entries[0]
            self.logger.info(f"Starting: {first_task['title']}")
            launch_agent_task(str(self.project_dir), first_task["file"], config)

            # Poll until batch complete
            self._poll_until_batch_complete(batch_num)

            if self.shutting_down:
                break

            # Gather batch results for GPT assessment
            state_data = self.state.read()
            batch_info = state_data["batches"][str(batch_num)]
            batch_results = []
            for idx in batch_info["task_indices"]:
                task = state_data["tasks"][idx]
                batch_results.append({
                    "title": task["title"],
                    "success": task["status"] == "completed",
                    "result_summary": task.get("result_summary", ""),
                    "files_modified": task.get("files_modified", []),
                    "errors": [],
                })

            # GPT assessment
            self.logger.info(f"Asking GPT to assess batch {batch_num}...")
            assessment = self.gpt.assess_batch(
                description=description,
                batch_num=batch_num,
                batch_results=batch_results,
                previous_summary=cumulative_summary,
            )

            cumulative_summary = assessment.get("cumulative_summary", cumulative_summary)
            previous_results = batch_results
            total_task_index += len(task_entries)

            complete_pct = assessment.get("completion_percentage", 0)
            self.logger.info(f"Assessment: {complete_pct}% complete")
            if assessment.get("issues"):
                for issue in assessment["issues"]:
                    self.logger.warning(issue)

            if assessment.get("project_complete"):
                self.logger.success("GPT says project is complete!")
                break

            # Check cost
            if state_data.get("total_cost_usd", 0) >= self.config["planner"]["max_cost_usd"]:
                self.logger.warning(f"Cost ceiling reached: ${state_data['total_cost_usd']:.2f}")
                break

        self._cleanup()

    def _resume(self):
        """Resume from existing state file with stall detection."""
        if not self.state.exists():
            print("No state file found. Nothing to resume.")
            sys.exit(1)

        state_data = self.state.read()
        print(f"Resuming from state: {state_data.get('status')}")
        print(f"  Tasks: {len(state_data.get('tasks', []))}")
        print(f"  Cost so far: ${state_data.get('total_cost_usd', 0):.2f}")

        # Detect stalled tasks
        tasks = state_data.get("tasks", [])
        stalled_tasks = []
        task_timeout = self.args.task_timeout if hasattr(self.args, 'task_timeout') else 600

        now = datetime.now(timezone.utc)
        for i, task in enumerate(tasks):
            if task["status"] == "running" and task.get("started_at"):
                try:
                    started_at = datetime.fromisoformat(task["started_at"].replace("Z", "+00:00"))
                    elapsed = (now - started_at).total_seconds()
                    if elapsed > task_timeout:
                        stalled_tasks.append((i, task, elapsed))
                except (ValueError, TypeError):
                    # Invalid timestamp, treat as stalled
                    stalled_tasks.append((i, task, 0))

        # Report stalled tasks
        if stalled_tasks:
            print("\n  WARNING: Found stalled tasks:")
            for idx, task, elapsed in stalled_tasks:
                print(f"    Task {idx + 1}: {task['title']} (stalled for {elapsed:.0f}s)")
                # Mark as failed so they can be retried
                self.state.update_task(idx, {
                    "status": "failed",
                    "result_summary": f"Task stalled (exceeded {task_timeout}s timeout)"
                })

        # Find next pending task
        next_idx = None
        for i, t in enumerate(tasks):
            if t["status"] in ("pending", "running"):
                # Re-check after marking stalled tasks as failed
                fresh_state = self.state.read()
                fresh_task = fresh_state["tasks"][i]
                if fresh_task["status"] in ("pending", "running"):
                    next_idx = i
                    break

        if next_idx is None:
            print("All tasks already completed.")
            self._cleanup()
            return

        # Re-install hook
        self.project_dir = Path(state_data["project_dir"])
        self.context_md = ContextMDManager(self.project_dir)
        install_hook(self.project_dir)

        # Reset status and persist current_task_index
        state_data["status"] = "running"
        state_data["current_task_index"] = next_idx
        self.state.write(state_data)

        # Mark task as running with timestamp
        self.state.mark_task_running(next_idx)

        task = tasks[next_idx]
        print(f"\n  > Resuming task {next_idx + 1}/{len(tasks)}: {task['title']}")
        launch_agent_task(
            str(self.project_dir),
            task["file"],
            state_data.get("config", {}),
        )

        self._poll_until_complete()
        self._cleanup()

    def _poll_until_complete(self, poll_interval: float = 3.0):
        """Poll state.json until all tasks are done or aborted."""
        last_index = -1

        while True:
            if self.shutting_down:
                # Signal handler set the flag — do the I/O cleanup here (safe)
                try:
                    if self.state.exists():
                        self.state.set_status("aborted", "User interrupted")
                except Exception:
                    pass
                break

            time.sleep(poll_interval)

            try:
                state_data = self.state.read()
            except (json.JSONDecodeError, OSError):
                continue

            status = state_data.get("status", "")
            if status in ("completed", "aborted"):
                break

            current_idx = state_data.get("current_task_index", 0)
            tasks = state_data.get("tasks", [])
            total = len(tasks)

            # Print progress
            if current_idx != last_index:
                completed = sum(1 for t in tasks if t["status"] in ("completed", "failed"))
                running = sum(1 for t in tasks if t["status"] == "running")
                cost = state_data.get("total_cost_usd", 0)

                if running > 0:
                    running_task = next(
                        (t for t in tasks if t["status"] == "running"), None
                    )
                    name = running_task["title"] if running_task else "?"
                    print(f"  [{completed}/{total}] Running: {name} (${cost:.2f} spent)")

                last_index = current_idx

            # Check if all done
            all_done = all(t["status"] in ("completed", "failed") for t in tasks)
            if all_done:
                break

            # Stall detection
            for t in tasks:
                if t["status"] == "running" and t.get("started_at"):
                    try:
                        started = datetime.fromisoformat(t["started_at"])
                        elapsed = (
                            datetime.now(timezone.utc) - started
                        ).total_seconds()
                        if elapsed > self.args.task_timeout:
                            print(f"  WARNING: Task {t['index']+1} stalled ({elapsed:.0f}s)")
                            self.state.update_task(
                                t["index"],
                                {"status": "failed", "result_summary": "Timeout / stall detected"},
                            )
                    except (ValueError, TypeError):
                        pass

    def _poll_until_batch_complete(self, batch_num: int, poll_interval: float = 3.0):
        """Poll until a specific batch is done."""
        last_index = -1

        while True:
            if self.shutting_down:
                try:
                    if self.state.exists():
                        self.state.set_status("aborted", "User interrupted")
                except Exception:
                    pass
                break

            time.sleep(poll_interval)

            try:
                state_data = self.state.read()
            except (json.JSONDecodeError, OSError):
                continue

            status = state_data.get("status", "")
            if status in ("aborted",):
                break

            batch = state_data.get("batches", {}).get(str(batch_num), {})
            if batch.get("status") == "completed":
                break

            # Print progress
            current_idx = state_data.get("current_task_index", 0)
            if current_idx != last_index:
                batch_indices = batch.get("task_indices", [])
                tasks = state_data.get("tasks", [])
                completed = sum(
                    1
                    for i in batch_indices
                    if i < len(tasks) and tasks[i]["status"] in ("completed", "failed")
                )
                total = len(batch_indices)
                cost = state_data.get("total_cost_usd", 0)

                running_task = next(
                    (tasks[i] for i in batch_indices if i < len(tasks) and tasks[i]["status"] == "running"),
                    None,
                )
                if running_task:
                    print(f"  [{completed}/{total}] Running: {running_task['title']} (${cost:.2f} spent)")

                last_index = current_idx

            # Check if all batch tasks done
            batch_indices = batch.get("task_indices", [])
            tasks = state_data.get("tasks", [])
            all_done = all(
                tasks[i]["status"] in ("completed", "failed")
                for i in batch_indices
                if i < len(tasks)
            )
            if all_done:
                break

    def _cleanup(self):
        """Clean up after orchestration."""
        self.logger.section("Results")

        if self.state.exists():
            state_data = self.state.read()
            tasks = state_data.get("tasks", [])
            completed = sum(1 for t in tasks if t["status"] == "completed")
            failed = sum(1 for t in tasks if t["status"] == "failed")
            cost = state_data.get("total_cost_usd", 0)

            self.logger.info(f"Completed: {completed}/{len(tasks)} tasks")
            if failed:
                self.logger.warning(f"Failed: {failed} tasks")
            self.logger.info(f"Cost: ${cost:.2f}")
            print()

            # Display task summary using rich table if available
            self.logger.task_summary(tasks, cost)

            if state_data.get("status") != "aborted":
                state_data["status"] = "completed"
                self.state.write(state_data)

        # Restore project
        try:
            self.context_md.restore()
        except Exception:
            pass

        try:
            uninstall_hook(self.project_dir)
        except Exception:
            pass

        self.logger.success(f"Done. State saved to: {self.state_path}")

        # Close logger
        self.logger.close()


# ──────────────────────────────────────────────
# Config commands
# ──────────────────────────────────────────────


def cmd_init():
    """Interactive configuration setup wizard."""
    import json

    print("\n" + "="*60)
    print("  GPT Agent Orchestrator - Configuration Setup")
    print("="*60 + "\n")

    config = load_config()

    # Ask provider
    print("Select planner provider:")
    print("  1. OpenAI (default)")
    print("  2. Anthropic")
    print("  3. Azure OpenAI")
    print("  4. DeepSeek")
    print("  5. OpenAI-compatible API")

    provider_choice = input("\nChoice [1]: ").strip() or "1"
    provider_map = {
        "1": "openai",
        "2": "anthropic",
        "3": "azure_openai",
        "4": "deepseek",
        "5": "openai_compatible",
    }
    provider = provider_map.get(provider_choice, "openai")
    config["planner"]["provider"] = provider

    # Ask model
    default_models = {
        "openai": "gpt-5.2",
        "anthropic": "claude-sonnet-4-5-20250929",
        "azure_openai": "gpt-4",
        "deepseek": "deepseek-chat",
        "openai_compatible": "gpt-3.5-turbo",
    }
    default_model = default_models.get(provider, "gpt-5.2")
    model = input(f"Model name [{default_model}]: ").strip() or default_model
    config["planner"]["model"] = model

    # Ask about API key storage
    print("\nWhere should the API key be stored?")
    print("  1. Environment variable (recommended)")
    print("  2. Config file")

    storage_choice = input("\nChoice [1]: ").strip() or "1"

    if storage_choice == "2":
        api_key = input(f"Enter API key: ").strip()
        if api_key:
            config["planner"]["api_key"] = api_key
    else:
        # Provide guidance on env var
        from .config import PROVIDER_API_KEY_ENV_VARS
        env_var = PROVIDER_API_KEY_ENV_VARS.get(provider, "SHOEMAKER_ELVES_OPENAI_API_KEY")
        print(f"\nSet your API key as environment variable: {env_var}")
        print(f"Example: export {env_var}=your-api-key-here")

    # Azure-specific settings
    if provider == "azure_openai":
        endpoint = input("\nAzure endpoint URL: ").strip()
        if endpoint:
            if "azure" not in config:
                config["azure"] = {}
            config["azure"]["endpoint"] = endpoint

        deployment = input("Azure deployment name: ").strip()
        if deployment:
            config["azure"]["deployment"] = deployment

        api_version = input("Azure API version [2024-02-01]: ").strip() or "2024-02-01"
        config["azure"]["api_version"] = api_version

    # OpenAI-compatible base URL
    if provider == "openai_compatible":
        base_url = input("\nBase URL for API: ").strip()
        if base_url:
            config["planner"]["base_url"] = base_url

    # DeepSeek base URL (optional, has default)
    if provider == "deepseek":
        print("\nDeepSeek defaults to https://api.deepseek.com")
        base_url = input("Custom base URL (leave empty for default): ").strip()
        if base_url:
            config["planner"]["base_url"] = base_url

    # Agent model
    agent_model = input(f"\nAgent model [sonnet]: ").strip() or "sonnet"
    config["agent"]["model"] = agent_model

    # Max cost
    max_cost = input(f"Max cost in USD [50.0]: ").strip() or "50.0"
    try:
        config["planner"]["max_cost_usd"] = float(max_cost)
    except ValueError:
        config["planner"]["max_cost_usd"] = 50.0

    # Save config
    try:
        save_config(config)
        config_path = get_config_path()
        print(f"\n✓ Configuration saved to: {config_path}")
        print("\nYou can now run: chainsmith <project-dir>")
    except Exception as e:
        print(f"\n✗ Failed to save config: {e}")
        sys.exit(1)


def cmd_config_show(no_mask: bool = False):
    """Show current configuration."""
    config = load_config()
    config = merge_env_overrides(config)

    if not no_mask:
        config = mask_secrets(config)

    import json
    print("\nCurrent configuration:")
    print("="*60)
    print(json.dumps(config, indent=2))
    print("="*60)
    print(f"\nConfig file: {get_config_path()}")
    print("(Environment variables have been applied)")


def cmd_config_set(key: str, value: str):
    """Set a configuration value."""
    config = load_config()

    try:
        config = set_config_value(config, key, value)
        save_config(config)
        print(f"✓ Set {key} = {value}")
        print(f"Config saved to: {get_config_path()}")
    except Exception as e:
        print(f"✗ Failed to set config: {e}")
        sys.exit(1)


# ──────────────────────────────────────────────
# Entry point
# ──────────────────────────────────────────────


def main():
    """Main entry point for the CLI."""
    args = parse_args()

    # Handle config commands
    if args.command == "init":
        cmd_init()
        return

    if args.command == "config":
        if args.config_command == "show":
            cmd_config_show(no_mask=args.no_mask)
        elif args.config_command == "set":
            cmd_config_set(args.key, args.value)
        else:
            print("Usage: chainsmith config {show|set}")
            sys.exit(1)
        return

    # Handle run command (orchestration)
    orchestrator = Orchestrator(args)
    orchestrator.run()


if __name__ == "__main__":
    main()
