"""
Cursor CLI wrapper with stream-json parsing and progress tracking.
"""
import json
import subprocess
import sys
from pathlib import Path
from typing import Optional, Dict, Any, List, Callable
from dataclasses import dataclass
from enum import Enum


class CursorResultStatus(Enum):
    """Status of cursor CLI execution."""
    SUCCESS = "success"
    FAILURE = "failure"
    TIMEOUT = "timeout"
    ERROR = "error"


@dataclass
class CursorResult:
    """Result of cursor CLI execution."""
    status: CursorResultStatus
    output: str
    error: Optional[str]
    duration_ms: Optional[int]
    files_modified: List[str]
    files_created: List[str]
    files_read: List[str]
    tool_calls: int
    

class CursorCLI:
    """Wrapper for Cursor CLI with stream-json parsing."""
    
    def __init__(self, 
                 project_root: Path, 
                 api_key: str,
                 model: str = "claude-3-5-sonnet-20241022",
                 verbose: bool = True):
        self.project_root = project_root
        self.api_key = api_key
        self.model = model
        self.verbose = verbose
    
    def run_agent(self,
                  task_description: str,
                  force: bool = True,
                  max_iterations: int = 5,
                  progress_callback: Optional[Callable[[str, Any], None]] = None) -> CursorResult:
        """
        Run cursor agent on a task.
        
        Args:
            task_description: Description of the task for the agent
            force: Allow file modifications
            max_iterations: Maximum number of agent iterations
            progress_callback: Callback function(event_type, data) for progress updates
        
        Returns:
            CursorResult with execution details
        """
        # Build command
        cmd = [
            "cursor-agent",
            "-p",  # Print mode (non-interactive)
            "--output-format", "stream-json",
            "--stream-partial-output",
        ]
        
        if force:
            cmd.append("--force")
        
        # Add task description
        cmd.append(task_description)
        
        # Set environment
        env = {
            "CURSOR_API_KEY": self.api_key,
            "CURSOR_MODEL": self.model,
        }
        
        if self.verbose:
            print(f"🚀 Running cursor-agent in {self.project_root}")
            print(f"📝 Task: {task_description[:100]}...")
        
        # Track execution metrics
        accumulated_text = ""
        tool_count = 0
        files_modified = []
        files_created = []
        files_read = []
        duration_ms = None
        error_message = None
        
        try:
            # Run cursor-agent with streaming output
            process = subprocess.Popen(
                cmd,
                cwd=self.project_root,
                env={**subprocess.os.environ, **env},
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                bufsize=1
            )
            
            # Process streaming output line by line
            for line in process.stdout:
                line = line.strip()
                if not line:
                    continue
                
                try:
                    event = json.loads(line)
                    event_type = event.get("type")
                    subtype = event.get("subtype")
                    
                    # Call progress callback
                    if progress_callback:
                        progress_callback(event_type, event)
                    
                    # Process different event types
                    if event_type == "system":
                        self._handle_system_event(event, subtype)
                    
                    elif event_type == "assistant":
                        accumulated_text += self._handle_assistant_event(event)
                    
                    elif event_type == "tool_call":
                        tool_result = self._handle_tool_call_event(
                            event, subtype, files_modified, files_created, files_read
                        )
                        if tool_result:
                            tool_count += 1
                    
                    elif event_type == "result":
                        duration_ms = event.get("duration_ms")
                        if self.verbose:
                            self._print_summary(
                                duration_ms, tool_count, 
                                len(accumulated_text),
                                files_created, files_modified
                            )
                
                except json.JSONDecodeError:
                    # Not JSON, might be error message
                    if self.verbose:
                        print(f"⚠️  {line}")
                    error_message = line
            
            # Wait for process to complete
            return_code = process.wait()
            
            # Check stderr
            stderr = process.stderr.read()
            if stderr:
                error_message = stderr
                if self.verbose:
                    print(f"❌ Error: {stderr}")
            
            # Determine status
            if return_code == 0:
                status = CursorResultStatus.SUCCESS
            else:
                status = CursorResultStatus.FAILURE
            
            return CursorResult(
                status=status,
                output=accumulated_text,
                error=error_message,
                duration_ms=duration_ms,
                files_modified=files_modified,
                files_created=files_created,
                files_read=files_read,
                tool_calls=tool_count
            )
        
        except subprocess.TimeoutExpired:
            return CursorResult(
                status=CursorResultStatus.TIMEOUT,
                output=accumulated_text,
                error="Cursor CLI execution timeout",
                duration_ms=None,
                files_modified=files_modified,
                files_created=files_created,
                files_read=files_read,
                tool_calls=tool_count
            )
        
        except Exception as e:
            return CursorResult(
                status=CursorResultStatus.ERROR,
                output=accumulated_text,
                error=str(e),
                duration_ms=None,
                files_modified=files_modified,
                files_created=files_created,
                files_read=files_read,
                tool_calls=tool_count
            )
    
    def _handle_system_event(self, event: Dict[str, Any], subtype: Optional[str]) -> None:
        """Handle system event."""
        if subtype == "init" and self.verbose:
            model = event.get("model", "unknown")
            print(f"🤖 Using model: {model}")
    
    def _handle_assistant_event(self, event: Dict[str, Any]) -> str:
        """Handle assistant event and return text delta."""
        content = event.get("message", {}).get("content", [])
        if content:
            text = content[0].get("text", "")
            if text and self.verbose:
                # Show progress indicator
                sys.stdout.write(".")
                sys.stdout.flush()
            return text
        return ""
    
    def _handle_tool_call_event(self,
                                event: Dict[str, Any],
                                subtype: Optional[str],
                                files_modified: List[str],
                                files_created: List[str],
                                files_read: List[str]) -> bool:
        """Handle tool call event. Returns True if tool completed."""
        tool_call = event.get("tool_call", {})
        
        if subtype == "started":
            # Tool started
            if "writeToolCall" in tool_call:
                path = tool_call["writeToolCall"].get("args", {}).get("path", "unknown")
                if self.verbose:
                    print(f"\n🔧 Writing: {path}")
            
            elif "readToolCall" in tool_call:
                path = tool_call["readToolCall"].get("args", {}).get("path", "unknown")
                files_read.append(path)
                if self.verbose:
                    print(f"\n📖 Reading: {path}")
            
            return False
        
        elif subtype == "completed":
            # Tool completed
            if "writeToolCall" in tool_call:
                result = tool_call["writeToolCall"].get("result", {})
                if "success" in result:
                    success = result["success"]
                    path = tool_call["writeToolCall"].get("args", {}).get("path", "unknown")
                    lines = success.get("linesCreated", 0)
                    
                    if lines > 0:
                        files_created.append(path)
                        if self.verbose:
                            print(f"   ✅ Created {lines} lines")
                    else:
                        files_modified.append(path)
                        if self.verbose:
                            lines_modified = success.get("linesModified", 0)
                            print(f"   ✅ Modified {lines_modified} lines")
            
            elif "readToolCall" in tool_call:
                result = tool_call["readToolCall"].get("result", {})
                if "success" in result and self.verbose:
                    lines = result["success"].get("totalLines", 0)
                    print(f"   ✅ Read {lines} lines")
            
            return True
        
        return False
    
    def _print_summary(self, 
                      duration_ms: Optional[int],
                      tool_count: int,
                      chars_generated: int,
                      files_created: List[str],
                      files_modified: List[str]) -> None:
        """Print execution summary."""
        print(f"\n\n🎯 Completed in {duration_ms}ms")
        print(f"📊 Stats: {tool_count} tool calls, {chars_generated} chars generated")
        if files_created:
            print(f"📝 Created {len(files_created)} file(s)")
        if files_modified:
            print(f"✏️  Modified {len(files_modified)} file(s)")


def test_cursor_cli():
    """Test cursor CLI wrapper."""
    import os
    from pathlib import Path
    
    api_key = os.getenv("CURSOR_API_KEY")
    if not api_key:
        print("❌ CURSOR_API_KEY not set")
        return
    
    project_root = Path.cwd()
    
    cli = CursorCLI(project_root, api_key)
    
    task = "List all Python files in tools/agent directory"
    
    def progress_callback(event_type: str, event: Dict[str, Any]):
        """Example progress callback."""
        if event_type == "tool_call" and event.get("subtype") == "started":
            print(f"🔄 Tool started...")
    
    result = cli.run_agent(
        task_description=task,
        force=False,
        progress_callback=progress_callback
    )
    
    print(f"\n📋 Result: {result.status}")
    print(f"📄 Output: {result.output[:200]}...")


if __name__ == "__main__":
    test_cursor_cli()