#!/usr/bin/env python3
"""Session State Manager - Persistent session state and continuity"""

import json
import uuid
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Any, Set

try:
    from .models import ToolExecution
except ImportError:
    from models import ToolExecution


class SessionStateManager:
    """Manages persistent session state across Claude interactions"""
    
    def __init__(self, state_dir: str = ".claude_hooks"):
        self.state_dir = Path(state_dir)
        self.state_dir.mkdir(parents=True, exist_ok=True)
        
        self.state_file = self.state_dir / "session_state.json"
        self.todos_file = Path("ACTIVE_TODOS.md")
        self.last_session_file = Path("LAST_SESSION.md")
        
        # Initialize session
        self.session_id = str(uuid.uuid4())[:8]
        self.current_state = self._load_or_create_state()
        
    def _load_or_create_state(self) -> Dict[str, Any]:
        """Load existing state or create new session state"""
        try:
            if self.state_file.exists():
                with open(self.state_file, 'r') as f:
                    state = json.load(f)
                
                # Check if this is a continuation of recent session
                last_activity = datetime.fromisoformat(state.get("last_activity", "1970-01-01"))
                if (datetime.now() - last_activity).total_seconds() < 3600:  # Within 1 hour
                    # Continue existing session
                    return state
                    
            # Create new session
            return self._create_new_session()
            
        except Exception:
            # If loading fails, create new session
            return self._create_new_session()
    
    def _create_new_session(self) -> Dict[str, Any]:
        """Create new session state"""
        return {
            "session_id": self.session_id,
            "start_time": datetime.now().isoformat(),
            "last_activity": datetime.now().isoformat(),
            "modified_files": [],
            "commands_executed": [],
            "tool_usage": {},
            "backup_history": [],
            "todos": [],
            "context_snapshots": []
        }
    
    def update_from_tool_use(self, tool_data: Dict[str, Any]):
        """Update session state from tool usage"""
        try:
            tool = tool_data.get("tool", "")
            params = tool_data.get("parameters", {})
            timestamp = datetime.now().isoformat()
            
            # Track file modifications
            if tool in ["Edit", "Write", "MultiEdit"]:
                file_path = params.get("file_path", "")
                if file_path and file_path not in self.current_state["modified_files"]:
                    self.current_state["modified_files"].append(file_path)
                    
            # Track commands executed
            if tool == "Bash":
                command = params.get("command", "")
                if command:
                    self.current_state["commands_executed"].append({
                        "command": command,
                        "timestamp": timestamp
                    })
                    
                    # Keep only last 50 commands
                    if len(self.current_state["commands_executed"]) > 50:
                        self.current_state["commands_executed"] = self.current_state["commands_executed"][-50:]
                    
            # Track tool usage statistics
            self.current_state["tool_usage"][tool] = self.current_state["tool_usage"].get(tool, 0) + 1
            self.current_state["last_activity"] = timestamp
            
            # Save state periodically
            self._save_state()
            
        except Exception:
            pass  # Don't let state tracking errors break the system
    
    def add_backup(self, backup_id: str, backup_info: Dict[str, Any]):
        """Record backup in session history"""
        try:
            backup_record = {
                "backup_id": backup_id,
                "timestamp": datetime.now().isoformat(),
                "reason": backup_info.get("reason", "unknown"),
                "success": backup_info.get("success", False)
            }
            
            self.current_state["backup_history"].append(backup_record)
            
            # Keep only last 10 backups
            if len(self.current_state["backup_history"]) > 10:
                self.current_state["backup_history"] = self.current_state["backup_history"][-10:]
                
            self._save_state()
            
        except Exception:
            pass
    
    def add_context_snapshot(self, context_data: Dict[str, Any]):
        """Add context snapshot for recovery"""
        try:
            snapshot = {
                "timestamp": datetime.now().isoformat(),
                "context_ratio": context_data.get("usage_ratio", 0.0),
                "prompt_count": context_data.get("prompt_count", 0),
                "tool_count": context_data.get("tool_executions", 0)
            }
            
            self.current_state["context_snapshots"].append(snapshot)
            
            # Keep only last 20 snapshots
            if len(self.current_state["context_snapshots"]) > 20:
                self.current_state["context_snapshots"] = self.current_state["context_snapshots"][-20:]
                
        except Exception:
            pass
    
    def update_todos(self, todos: List[Dict[str, Any]]):
        """Update active todos list"""
        try:
            self.current_state["todos"] = todos
            self._save_state()
            self._update_todos_file()
            
        except Exception:
            pass
    
    def get_session_summary(self) -> Dict[str, Any]:
        """Generate comprehensive session summary"""
        try:
            return {
                "session_id": self.current_state.get("session_id", "unknown"),
                "start_time": self.current_state.get("start_time", "unknown"),
                "last_activity": self.current_state.get("last_activity", "unknown"),
                "modified_files": self.current_state.get("modified_files", []),
                "tool_usage": self.current_state.get("tool_usage", {}),
                "commands_executed": self.current_state.get("commands_executed", []),
                "backup_history": self.current_state.get("backup_history", []),
                "todos": self.current_state.get("todos", []),
                "session_stats": self._calculate_session_stats()
            }
        except Exception:
            return {"error": "Failed to generate session summary"}
    
    def _calculate_session_stats(self) -> Dict[str, Any]:
        """Calculate session statistics"""
        try:
            total_tools = sum(self.current_state.get("tool_usage", {}).values())
            total_commands = len(self.current_state.get("commands_executed", []))
            total_files = len(self.current_state.get("modified_files", []))
            
            start_time = datetime.fromisoformat(self.current_state.get("start_time", datetime.now().isoformat()))
            duration = datetime.now() - start_time
            
            return {
                "duration_minutes": round(duration.total_seconds() / 60, 1),
                "total_tool_calls": total_tools,
                "total_commands": total_commands,
                "total_files_modified": total_files,
                "most_used_tools": self._get_top_tools(3)
            }
        except Exception:
            return {}
    
    def _get_top_tools(self, count: int) -> List[Dict[str, Any]]:
        """Get most frequently used tools"""
        try:
            tool_usage = self.current_state.get("tool_usage", {})
            sorted_tools = sorted(tool_usage.items(), key=lambda x: x[1], reverse=True)
            return [{"tool": tool, "count": usage} for tool, usage in sorted_tools[:count]]
        except Exception:
            return []
    
    def create_continuation_docs(self):
        """Create LAST_SESSION.md and ACTIVE_TODOS.md"""
        try:
            self._create_last_session_doc()
            self._update_todos_file()
        except Exception:
            pass  # Don't let doc creation errors break the system
    
    def _create_last_session_doc(self):
        """Create LAST_SESSION.md with session summary"""
        try:
            summary = self.get_session_summary()
            
            content = f"""# Last Claude Session Summary

**Session ID**: {summary['session_id']}
**Duration**: {summary['start_time']} → {summary['last_activity']}
**Session Length**: {summary.get('session_stats', {}).get('duration_minutes', 0)} minutes

## Files Modified ({len(summary['modified_files'])})
"""
            
            for file_path in summary['modified_files']:
                content += f"- {file_path}\n"
            
            content += f"\n## Tools Used ({summary.get('session_stats', {}).get('total_tool_calls', 0)} total)\n"
            
            for tool, count in summary['tool_usage'].items():
                content += f"- {tool}: {count} times\n"
            
            content += f"\n## Recent Commands ({len(summary['commands_executed'])})\n"
            
            # Show last 10 commands
            recent_commands = summary['commands_executed'][-10:]
            for cmd_info in recent_commands:
                timestamp = cmd_info['timestamp'][:19]  # Remove microseconds
                content += f"- `{cmd_info['command']}` ({timestamp})\n"
            
            content += f"\n## Backup History\n"
            
            for backup in summary['backup_history']:
                status = "✅" if backup['success'] else "❌"
                content += f"- {status} {backup['backup_id']} - {backup['reason']} ({backup['timestamp'][:19]})\n"
            
            content += f"""
## To Continue This Session

1. **Review Modified Files**: Check the files listed above for your recent changes
2. **Check Active Tasks**: Review `ACTIVE_TODOS.md` for pending work
3. **Restore Context**: Reference the commands and tools used above
4. **Use Backups**: If needed, restore from backup using `claude-hooks restore {summary['backup_history'][-1]['backup_id'] if summary['backup_history'] else 'latest'}`

## Quick Commands

```bash
# View current project status
git status

# Check for any uncommitted changes
git diff

# List available backups
claude-hooks list-backups

# Continue with active todos
cat ACTIVE_TODOS.md
```
"""
            
            with open(self.last_session_file, 'w') as f:
                f.write(content)
                
        except Exception as e:
            # Create minimal doc on error
            try:
                with open(self.last_session_file, 'w') as f:
                    f.write(f"# Last Session\n\nSession ended at {datetime.now().isoformat()}\n\nError creating summary: {e}\n")
            except Exception:
                pass
    
    def _update_todos_file(self):
        """Update ACTIVE_TODOS.md file"""
        try:
            todos = self.current_state.get("todos", [])
            
            if not todos:
                content = """# Active TODOs

*No active todos. Add some to track your progress!*

## How to Add TODOs

Use Claude's TodoWrite tool to manage your task list:
- Track progress across sessions
- Break down complex tasks
- Never lose track of what you're working on
"""
            else:
                content = f"""# Active TODOs

*Updated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}*

"""
                
                # Group by status
                pending_todos = [t for t in todos if t.get('status') == 'pending']
                in_progress_todos = [t for t in todos if t.get('status') == 'in_progress']
                completed_todos = [t for t in todos if t.get('status') == 'completed']
                
                if in_progress_todos:
                    content += "## 🚀 In Progress\n\n"
                    for todo in in_progress_todos:
                        priority = todo.get('priority', 'medium')
                        priority_emoji = {'high': '🔥', 'medium': '⭐', 'low': '📝'}.get(priority, '⭐')
                        content += f"- {priority_emoji} {todo.get('content', 'Unknown task')}\n"
                    content += "\n"
                
                if pending_todos:
                    content += "## 📋 Pending\n\n"
                    for todo in pending_todos:
                        priority = todo.get('priority', 'medium')
                        priority_emoji = {'high': '🔥', 'medium': '⭐', 'low': '📝'}.get(priority, '⭐')
                        content += f"- {priority_emoji} {todo.get('content', 'Unknown task')}\n"
                    content += "\n"
                
                if completed_todos:
                    content += "## ✅ Completed\n\n"
                    for todo in completed_todos[-5:]:  # Show last 5 completed
                        content += f"- ✅ {todo.get('content', 'Unknown task')}\n"
                    content += "\n"
            
            with open(self.todos_file, 'w') as f:
                f.write(content)
                
        except Exception:
            pass  # Don't let todo file creation break the system
    
    def _save_state(self):
        """Save current state to disk"""
        try:
            with open(self.state_file, 'w') as f:
                json.dump(self.current_state, f, indent=2)
        except Exception:
            pass  # Don't let state saving errors break the system
    
    def cleanup_session(self):
        """Clean up session and create final documentation"""
        try:
            self.create_continuation_docs()
            self._save_state()
        except Exception:
            pass