#!/usr/bin/env python3
"""
Task Runner for Allegro-Primus
Executes tasks in isolated worktrees with result capture
"""

import json
import os
import subprocess
import sys
import time
import traceback
from dataclasses import dataclass, asdict
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Callable, Any
from contextlib import contextmanager


@dataclass
class TaskResult:
    """Result of a task execution"""
    task_id: str
    issue_number: str
    success: bool
    return_code: int
    stdout: str
    stderr: str
    execution_time: float
    timestamp: str
    artifacts: List[str]
    commit_hash: Optional[str] = None
    error: Optional[str] = None
    
    def to_dict(self) -> Dict:
        return asdict(self)
    
    def to_json(self, indent: int = 2) -> str:
        return json.dumps(self.to_dict(), indent=indent)


class TaskRunner:
    """Runs tasks in isolated worktree environments"""
    
    def __init__(self, worktree_path: str, logs_dir: str = None):
        self.worktree_path = Path(worktree_path)
        self.logs_dir = Path(logs_dir) if logs_dir else self.worktree_path / ".ap-logs"
        self.logs_dir.mkdir(exist_ok=True)
        self.task_history: List[TaskResult] = []
    
    def _run_git(self, args: List[str], cwd: str = None) -> subprocess.CompletedProcess:
        """Run git command in worktree"""
        return subprocess.run(
            ["git"] + args,
            capture_output=True,
            text=True,
            cwd=cwd or str(self.worktree_path)
        )
    
    def _write_log(self, task_id: str, content: str):
        """Write content to task log file"""
        log_file = self.logs_dir / f"{task_id}.log"
        log_file.write_text(content)
    
    @contextmanager
    def _capture_context(self):
        """Context manager for capturing output"""
        yield
    
    def run_command(self, 
                   task_id: str,
                   command: List[str],
                   issue_number: str,
                   env: Dict[str, str] = None,
                   timeout: int = 300,
                   capture_output: bool = True) -> TaskResult:
        """
        Run a shell command in the worktree
        
        Args:
            task_id: Unique identifier for this task
            command: Command and arguments as list
            issue_number: Associated issue number
            env: Additional environment variables
            timeout: Command timeout in seconds
            capture_output: Whether to capture stdout/stderr
            
        Returns:
            TaskResult with execution details
        """
        start_time = time.time()
        timestamp = datetime.now().isoformat()
        
        # Prepare environment
        run_env = os.environ.copy()
        run_env["AP_WORKTREE"] = str(self.worktree_path)
        run_env["AP_ISSUE"] = issue_number
        run_env["AP_TASK_ID"] = task_id
        if env:
            run_env.update(env)
        
        try:
            # Run command
            result = subprocess.run(
                command,
                capture_output=capture_output,
                text=True,
                cwd=str(self.worktree_path),
                env=run_env,
                timeout=timeout
            )
            
            execution_time = time.time() - start_time
            
            # Create result
            task_result = TaskResult(
                task_id=task_id,
                issue_number=issue_number,
                success=result.returncode == 0,
                return_code=result.returncode,
                stdout=result.stdout,
                stderr=result.stderr,
                execution_time=execution_time,
                timestamp=timestamp,
                artifacts=[]
            )
            
            # Write log
            log_content = f"""Task: {task_id}
Issue: {issue_number}
Command: {' '.join(command)}
Timestamp: {timestamp}
Duration: {execution_time:.2f}s
Return Code: {result.returncode}

=== STDOUT ===
{result.stdout}

=== STDERR ===
{result.stderr}
"""
            self._write_log(task_id, log_content)
            
        except subprocess.TimeoutExpired as e:
            execution_time = time.time() - start_time
            task_result = TaskResult(
                task_id=task_id,
                issue_number=issue_number,
                success=False,
                return_code=-1,
                stdout=e.stdout if e.stdout else "",
                stderr=f"Timeout after {timeout}s",
                execution_time=execution_time,
                timestamp=timestamp,
                artifacts=[],
                error="Timeout"
            )
        except Exception as e:
            execution_time = time.time() - start_time
            task_result = TaskResult(
                task_id=task_id,
                issue_number=issue_number,
                success=False,
                return_code=-1,
                stdout="",
                stderr=str(e),
                execution_time=execution_time,
                timestamp=timestamp,
                artifacts=[],
                error=traceback.format_exc()
            )
        
        self.task_history.append(task_result)
        return task_result
    
    def run_script(self,
                  task_id: str,
                  script_content: str,
                  issue_number: str,
                  interpreter: str = "python3",
                  timeout: int = 300) -> TaskResult:
        """
        Run a script in the worktree
        
        Args:
            task_id: Task identifier
            script_content: Script to execute
            issue_number: Associated issue
            interpreter: Script interpreter
            timeout: Execution timeout
        """
        # Write script to temporary file
        script_path = self.worktree_path / f".ap-script-{task_id}.py"
        script_path.write_text(script_content)
        
        try:
            result = self.run_command(
                task_id=task_id,
                command=[interpreter, str(script_path)],
                issue_number=issue_number,
                timeout=timeout
            )
        finally:
            # Cleanup script
            script_path.unlink(missing_ok=True)
        
        return result
    
    def run_python_function(self,
                           task_id: str,
                           func: Callable,
                           issue_number: str,
                           args: tuple = (),
                           kwargs: Dict = None) -> TaskResult:
        """
        Execute a Python function in the worktree context
        
        Args:
            task_id: Task identifier
            func: Function to execute
            issue_number: Associated issue
            args: Positional arguments
            kwargs: Keyword arguments
        """
        start_time = time.time()
        timestamp = datetime.now().isoformat()
        kwargs = kwargs or {}
        
        # Store original cwd
        original_cwd = os.getcwd()
        
        try:
            # Change to worktree
            os.chdir(self.worktree_path)
            
            # Set environment markers
            os.environ["AP_WORKTREE"] = str(self.worktree_path)
            os.environ["AP_ISSUE"] = issue_number
            os.environ["AP_TASK_ID"] = task_id
            
            # Execute function
            result = func(*args, **kwargs)
            
            execution_time = time.time() - start_time
            
            task_result = TaskResult(
                task_id=task_id,
                issue_number=issue_number,
                success=True,
                return_code=0,
                stdout=str(result) if result else "",
                stderr="",
                execution_time=execution_time,
                timestamp=timestamp,
                artifacts=[]
            )
            
        except Exception as e:
            execution_time = time.time() - start_time
            task_result = TaskResult(
                task_id=task_id,
                issue_number=issue_number,
                success=False,
                return_code=-1,
                stdout="",
                stderr=str(e),
                execution_time=execution_time,
                timestamp=timestamp,
                artifacts=[],
                error=traceback.format_exc()
            )
        finally:
            os.chdir(original_cwd)
        
        self.task_history.append(task_result)
        return task_result
    
    def auto_commit(self, 
                   message: str,
                   files: List[str] = None,
                   allow_empty: bool = False) -> Optional[str]:
        """
        Automatically commit changes in the worktree
        
        Args:
            message: Commit message
            files: Specific files to commit (None = all changes)
            allow_empty: Allow empty commits
            
        Returns:
            Commit hash if successful
        """
        # Add files
        if files:
            for f in files:
                self._run_git(["add", f])
        else:
            self._run_git(["add", "-A"])
        
        # Check if there are changes
        status = self._run_git(["status", "--porcelain"])
        if not status.stdout.strip() and not allow_empty:
            return None
        
        # Commit
        cmd = ["commit", "-m", f"[AP] {message}"]
        if allow_empty:
            cmd.append("--allow-empty")
        
        result = self._run_git(cmd)
        
        if result.returncode == 0:
            # Get commit hash
            log_result = self._run_git(["log", "-1", "--format=%H"])
            return log_result.stdout.strip()
        
        return None
    
    def get_logs(self, task_id: str = None) -> str:
        """Get logs for a task or all tasks"""
        if task_id:
            log_file = self.logs_dir / f"{task_id}.log"
            if log_file.exists():
                return log_file.read_text()
            return ""
        
        # Return all logs
        all_logs = []
        for log_file in sorted(self.logs_dir.glob("*.log")):
            all_logs.append(f"=== {log_file.stem} ===")
            all_logs.append(log_file.read_text())
            all_logs.append("")
        
        return "\n".join(all_logs)
    
    def generate_report(self, output_path: str = None) -> str:
        """Generate a report of all task executions"""
        report = {
            "worktree": str(self.worktree_path),
            "total_tasks": len(self.task_history),
            "successful": sum(1 for t in self.task_history if t.success),
            "failed": sum(1 for t in self.task_history if not t.success),
            "tasks": [t.to_dict() for t in self.task_history]
        }
        
        report_json = json.dumps(report, indent=2)
        
        if output_path:
            Path(output_path).write_text(report_json)
        
        return report_json


class TaskBatch:
    """Batch multiple tasks for sequential execution"""
    
    def __init__(self, runner: TaskRunner, issue_number: str):
        self.runner = runner
        self.issue_number = issue_number
        self.tasks: List[Dict] = []
        self.results: List[TaskResult] = []
    
    def add_command(self, task_id: str, command: List[str], **kwargs):
        """Add a command task to the batch"""
        self.tasks.append({
            "type": "command",
            "task_id": task_id,
            "command": command,
            "kwargs": kwargs
        })
    
    def add_script(self, task_id: str, script: str, **kwargs):
        """Add a script task to the batch"""
        self.tasks.append({
            "type": "script",
            "task_id": task_id,
            "script": script,
            "kwargs": kwargs
        })
    
    def execute(self, stop_on_error: bool = True) -> List[TaskResult]:
        """Execute all tasks in the batch"""
        self.results = []
        
        for task in self.tasks:
            if task["type"] == "command":
                result = self.runner.run_command(
                    task_id=task["task_id"],
                    command=task["command"],
                    issue_number=self.issue_number,
                    **task["kwargs"]
                )
            elif task["type"] == "script":
                result = self.runner.run_script(
                    task_id=task["task_id"],
                    script_content=task["script"],
                    issue_number=self.issue_number,
                    **task["kwargs"]
                )
            
            self.results.append(result)
            
            if stop_on_error and not result.success:
                break
        
        return self.results


if __name__ == "__main__":
    # Test
    import tempfile
    
    with tempfile.TemporaryDirectory() as tmpdir:
        runner = TaskRunner(tmpdir)
        
        # Test command execution
        result = runner.run_command(
            task_id="test-1",
            command=["echo", "Hello from AP"],
            issue_number="123"
        )
        
        print(f"Success: {result.success}")
        print(f"Output: {result.stdout}")