#!/usr/bin/env python3
"""Task management utilities for the guardrail workflow."""

import argparse
import os
import sys
from datetime import datetime
from pathlib import Path

# Try to import yaml, fall back to basic parsing if not available
try:
    import yaml
    HAS_YAML = True
except ImportError:
    HAS_YAML = False


def parse_yaml_simple(content: str) -> dict:
    """Simple YAML parser for basic task files."""
    result = {}
    current_key = None
    current_list = None

    for line in content.split('\n'):
        line = line.rstrip()
        if not line or line.startswith('#'):
            continue

        # Handle list items
        if line.startswith('  - '):
            if current_list is not None:
                current_list.append(line[4:].strip())
            continue

        # Handle key-value pairs
        if ':' in line and not line.startswith(' '):
            key, _, value = line.partition(':')
            key = key.strip()
            value = value.strip()

            if value:
                result[key] = value
                current_list = None
            else:
                result[key] = []
                current_list = result[key]
            current_key = key

    return result


def load_yaml(filepath: str) -> dict:
    """Load YAML file."""
    with open(filepath, 'r') as f:
        content = f.read()

    if HAS_YAML:
        return yaml.safe_load(content) or {}
    return parse_yaml_simple(content)


def save_yaml(filepath: str, data: dict):
    """Save data to YAML file."""
    if HAS_YAML:
        with open(filepath, 'w') as f:
            yaml.dump(data, f, default_flow_style=False, sort_keys=False)
    else:
        # Simple YAML writer
        lines = []
        for key, value in data.items():
            if isinstance(value, list):
                lines.append(f"{key}:")
                for item in value:
                    lines.append(f"  - {item}")
            elif isinstance(value, str) and '\n' in value:
                lines.append(f"{key}: |")
                for line in value.split('\n'):
                    lines.append(f"  {line}")
            else:
                lines.append(f"{key}: {value}")

        with open(filepath, 'w') as f:
            f.write('\n'.join(lines))


# ============================================================================
# Version-Aware Task Directory
# ============================================================================

def get_workflow_dir() -> Path:
    """Get the .workflow directory path."""
    return Path('.workflow')


def get_current_tasks_dir() -> str:
    """Get the tasks directory for the currently active workflow version.
    
    Returns the version-specific tasks directory if a workflow is active,
    otherwise falls back to 'tasks' for backward compatibility.
    """
    current_path = get_workflow_dir() / 'current.yml'
    if not current_path.exists():
        return 'tasks'  # Fallback for no active workflow
    
    current = load_yaml(str(current_path))
    version = current.get('active_version')
    if not version:
        return 'tasks'  # Fallback
    
    tasks_dir = get_workflow_dir() / 'versions' / version / 'tasks'
    tasks_dir.mkdir(parents=True, exist_ok=True)
    return str(tasks_dir)


# ============================================================================
# Task Operations
# ============================================================================

def find_tasks(tasks_dir: str, filters: dict = None) -> list:
    """Find all task files matching filters."""
    tasks = []
    tasks_path = Path(tasks_dir)

    if not tasks_path.exists():
        return tasks

    for filepath in tasks_path.glob('**/*.yml'):
        try:
            task = load_yaml(str(filepath))
            task['_filepath'] = str(filepath)

            # Apply filters
            if filters:
                match = True
                for key, value in filters.items():
                    if task.get(key) != value:
                        match = False
                        break
                if not match:
                    continue

            tasks.append(task)
        except Exception as e:
            print(f"Warning: Could not parse {filepath}: {e}", file=sys.stderr)

    return tasks


def list_tasks(tasks_dir: str, status: str = None, agent: str = None):
    """List tasks with optional filtering."""
    filters = {}
    if status:
        filters['status'] = status
    if agent:
        filters['agent'] = agent

    tasks = find_tasks(tasks_dir, filters)

    if not tasks:
        print("No tasks found.")
        return

    # Group by status
    by_status = {}
    for task in tasks:
        s = task.get('status', 'unknown')
        if s not in by_status:
            by_status[s] = []
        by_status[s].append(task)

    print("\n" + "=" * 60)
    print("TASK LIST")
    print("=" * 60)

    status_order = ['pending', 'in_progress', 'review', 'approved', 'completed', 'blocked']
    for s in status_order:
        if s in by_status:
            print(f"\n{s.upper()} ({len(by_status[s])})")
            print("-" * 40)
            for task in by_status[s]:
                agent = task.get('agent', '?')
                priority = task.get('priority', 'medium')
                print(f"  [{agent}] {task.get('id', 'unknown')} ({priority})")
                print(f"         {task.get('title', 'No title')}")


def get_next_task(tasks_dir: str, agent: str) -> dict:
    """Get next available task for an agent."""
    tasks = find_tasks(tasks_dir, {'agent': agent, 'status': 'pending'})

    if not tasks:
        return None

    # Sort by priority (high > medium > low)
    priority_order = {'high': 0, 'medium': 1, 'low': 2}
    tasks.sort(key=lambda t: priority_order.get(t.get('priority', 'medium'), 1))

    # Check dependencies
    for task in tasks:
        deps = task.get('dependencies', [])
        if not deps:
            return task

        # Check if all dependencies are completed
        all_deps_done = True
        for dep_id in deps:
            dep_tasks = find_tasks(tasks_dir, {'id': dep_id})
            if dep_tasks and dep_tasks[0].get('status') != 'completed':
                all_deps_done = False
                break

        if all_deps_done:
            return task

    return None


def update_task_status(tasks_dir: str, task_id: str, new_status: str, notes: str = None):
    """Update task status."""
    tasks = find_tasks(tasks_dir, {'id': task_id})

    if not tasks:
        print(f"Error: Task {task_id} not found")
        return False

    task = tasks[0]
    filepath = task['_filepath']

    # Remove internal field
    del task['_filepath']

    # Update status
    task['status'] = new_status

    if new_status == 'completed':
        task['completed_at'] = datetime.now().isoformat()

    if notes:
        task['review_notes'] = notes

    save_yaml(filepath, task)
    print(f"Updated {task_id} to {new_status}")
    return True


def complete_all_tasks(tasks_dir: str):
    """Mark all non-completed tasks as completed."""
    tasks = find_tasks(tasks_dir)
    completed_count = 0

    for task in tasks:
        if task.get('status') != 'completed':
            filepath = task['_filepath']
            del task['_filepath']
            task['status'] = 'completed'
            task['completed_at'] = datetime.now().isoformat()
            save_yaml(filepath, task)
            completed_count += 1
            print(f"  Completed: {task.get('id', 'unknown')}")

    print(f"\nMarked {completed_count} task(s) as completed.")
    return completed_count


def show_status(tasks_dir: str, manifest_path: str):
    """Show overall workflow status."""
    tasks = find_tasks(tasks_dir)

    # Count by status
    status_counts = {}
    agent_counts = {'frontend': {'pending': 0, 'completed': 0},
                    'backend': {'pending': 0, 'completed': 0},
                    'reviewer': {'pending': 0}}

    for task in tasks:
        s = task.get('status', 'unknown')
        status_counts[s] = status_counts.get(s, 0) + 1

        agent = task.get('agent', 'unknown')
        if agent in agent_counts:
            if s == 'pending':
                agent_counts[agent]['pending'] += 1
            elif s == 'completed':
                if 'completed' in agent_counts[agent]:
                    agent_counts[agent]['completed'] += 1

    print("\n" + "╔" + "═" * 58 + "╗")
    print("║" + "WORKFLOW STATUS".center(58) + "║")
    print("╠" + "═" * 58 + "╣")
    print("║" + " TASKS BY STATUS".ljust(58) + "║")
    print("║" + f"   ⏳ Pending:     {status_counts.get('pending', 0)}".ljust(58) + "║")
    print("║" + f"   🔄 In Progress: {status_counts.get('in_progress', 0)}".ljust(58) + "║")
    print("║" + f"   🔍 Review:      {status_counts.get('review', 0)}".ljust(58) + "║")
    print("║" + f"   ✅ Approved:    {status_counts.get('approved', 0)}".ljust(58) + "║")
    print("║" + f"   ✓  Completed:   {status_counts.get('completed', 0)}".ljust(58) + "║")
    print("║" + f"   🚫 Blocked:     {status_counts.get('blocked', 0)}".ljust(58) + "║")
    print("╠" + "═" * 58 + "╣")
    print("║" + " TASKS BY AGENT".ljust(58) + "║")
    print("║" + f"   🎨 Frontend: {agent_counts['frontend']['pending']} pending, {agent_counts['frontend']['completed']} completed".ljust(58) + "║")
    print("║" + f"   ⚙️  Backend:  {agent_counts['backend']['pending']} pending, {agent_counts['backend']['completed']} completed".ljust(58) + "║")
    print("║" + f"   🔍 Reviewer: {agent_counts['reviewer']['pending']} pending".ljust(58) + "║")
    print("╚" + "═" * 58 + "╝")


def main():
    parser = argparse.ArgumentParser(description="Task management for guardrail workflow")
    subparsers = parser.add_subparsers(dest='command', help='Commands')

    # list command
    list_parser = subparsers.add_parser('list', help='List tasks')
    list_parser.add_argument('--status', help='Filter by status')
    list_parser.add_argument('--agent', help='Filter by agent')
    list_parser.add_argument('--tasks-dir', default=None, help='Tasks directory (defaults to current version)')

    # next command
    next_parser = subparsers.add_parser('next', help='Get next task for agent')
    next_parser.add_argument('agent', choices=['frontend', 'backend', 'reviewer'])
    next_parser.add_argument('--tasks-dir', default=None, help='Tasks directory (defaults to current version)')

    # update command
    update_parser = subparsers.add_parser('update', help='Update task status')
    update_parser.add_argument('task_id', help='Task ID')
    update_parser.add_argument('status', choices=['pending', 'in_progress', 'review', 'approved', 'completed', 'blocked'])
    update_parser.add_argument('--notes', help='Review notes')
    update_parser.add_argument('--tasks-dir', default=None, help='Tasks directory (defaults to current version)')

    # status command
    status_parser = subparsers.add_parser('status', help='Show workflow status')
    status_parser.add_argument('--tasks-dir', default=None, help='Tasks directory (defaults to current version)')
    status_parser.add_argument('--manifest', default='project_manifest.json', help='Manifest path')

    # complete-all command
    complete_all_parser = subparsers.add_parser('complete-all', help='Mark all tasks as completed')
    complete_all_parser.add_argument('--tasks-dir', default=None, help='Tasks directory (defaults to current version)')

    args = parser.parse_args()
    
    # Resolve tasks_dir to version-specific directory if not explicitly provided
    if hasattr(args, 'tasks_dir') and args.tasks_dir is None:
        args.tasks_dir = get_current_tasks_dir()

    if args.command == 'list':
        list_tasks(args.tasks_dir, args.status, args.agent)
    elif args.command == 'next':
        task = get_next_task(args.tasks_dir, args.agent)
        if task:
            print(f"Next task for {args.agent}: {task.get('id')}")
            print(f"  Title: {task.get('title')}")
            print(f"  Files: {task.get('file_paths', [])}")
        else:
            print(f"No pending tasks for {args.agent}")
    elif args.command == 'update':
        update_task_status(args.tasks_dir, args.task_id, args.status, args.notes)
    elif args.command == 'status':
        show_status(args.tasks_dir, args.manifest)
    elif args.command == 'complete-all':
        complete_all_tasks(args.tasks_dir)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()