"""
Advanced chat interface for AI interactions with context management and response handling
"""
from typing import Dict, List, Optional, Any, Union
from dataclasses import dataclass
from datetime import datetime
import json
import asyncio
from enum import Enum
import re

class MessageRole(Enum):
    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"
    FUNCTION = "function"

@dataclass
class Message:
    role: MessageRole
    content: str
    timestamp: datetime
    metadata: Optional[Dict] = None

@dataclass
class CodeAction:
    action_type: str  # 'create', 'modify', 'delete', 'rename'
    file_path: str
    content: Optional[str] = None
    old_content: Optional[str] = None
    line_number: Optional[int] = None
    description: Optional[str] = None

@dataclass
class CommandAction:
    command: str
    cwd: Optional[str] = None
    expected_output: Optional[List[str]] = None
    timeout: Optional[int] = None
    description: Optional[str] = None

class ResponseType(Enum):
    TEXT = "text"
    CODE = "code"
    COMMAND = "command"
    ERROR = "error"
    SUGGESTION = "suggestion"

@dataclass
class Response:
    type: ResponseType
    content: Union[str, CodeAction, CommandAction]
    context: Optional[Dict] = None
    requires_confirmation: bool = False
    alternative_solutions: List[Any] = None

class ChatSession:
    def __init__(self, session_id: str):
        self.session_id = session_id
        self.messages: List[Message] = []
        self.context: Dict[str, Any] = {}
        self.start_time = datetime.now()
        
    def add_message(self, role: MessageRole, content: str, metadata: Optional[Dict] = None):
        message = Message(
            role=role,
            content=content,
            timestamp=datetime.now(),
            metadata=metadata
        )
        self.messages.append(message)
        
    def get_context_window(self, window_size: int = 10) -> List[Dict[str, Any]]:
        """Get the recent conversation context"""
        messages = self.messages[-window_size:]
        return [
            {
                "role": msg.role.value,
                "content": msg.content,
                "timestamp": msg.timestamp.isoformat()
            }
            for msg in messages
        ]
        
    def clear_history(self):
        """Clear conversation history"""
        self.messages.clear()
        
    def save_session(self, file_path: str):
        """Save session to file"""
        data = {
            "session_id": self.session_id,
            "start_time": self.start_time.isoformat(),
            "context": self.context,
            "messages": [
                {
                    "role": msg.role.value,
                    "content": msg.content,
                    "timestamp": msg.timestamp.isoformat(),
                    "metadata": msg.metadata
                }
                for msg in self.messages
            ]
        }
        with open(file_path, 'w') as f:
            json.dump(data, f, indent=2)
            
    @classmethod
    def load_session(cls, file_path: str) -> 'ChatSession':
        """Load session from file"""
        with open(file_path, 'r') as f:
            data = json.load(f)
            
        session = cls(data["session_id"])
        session.start_time = datetime.fromisoformat(data["start_time"])
        session.context = data["context"]
        
        for msg_data in data["messages"]:
            session.add_message(
                role=MessageRole(msg_data["role"]),
                content=msg_data["content"],
                metadata=msg_data.get("metadata")
            )
            
        return session

class ChatInterface:
    def __init__(self, api_key: str, model: str = "gpt-4"):
        self.api_key = api_key
        self.model = model
        self.sessions: Dict[str, ChatSession] = {}
        self.active_session: Optional[ChatSession] = None
        self._load_prompts()
        
    def _load_prompts(self):
        """Load system prompts and templates"""
        self.prompts = {
            "system_init": """You are an advanced AI coding assistant with expertise in:
1. Code analysis and modification
2. Project structure understanding
3. Task breakdown and planning
4. Terminal command execution
5. File operations and management

You can:
1. Read and modify code files
2. Execute terminal commands
3. Install dependencies
4. Run tests
5. Manage git operations

Please provide clear, specific responses and always validate actions before execution.""",
            
            "code_review": """Please review the following code changes:

{code_diff}

Consider:
1. Potential bugs or issues
2. Performance implications
3. Security concerns
4. Best practices
5. Testing requirements""",
            
            "error_analysis": """An error occurred during execution:

{error_message}

Please analyze:
1. Root cause
2. Potential fixes
3. Prevention measures""",
        }
        
    def create_session(self, session_id: Optional[str] = None) -> str:
        """Create a new chat session"""
        if not session_id:
            session_id = f"session_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
            
        session = ChatSession(session_id)
        self.sessions[session_id] = session
        self.active_session = session
        
        # Initialize with system prompt
        session.add_message(
            role=MessageRole.SYSTEM,
            content=self.prompts["system_init"]
        )
        
        return session_id
        
    def switch_session(self, session_id: str):
        """Switch to a different session"""
        if session_id not in self.sessions:
            raise ValueError(f"Session {session_id} not found")
        self.active_session = self.sessions[session_id]
        
    async def send_message(self, 
                        content: str, 
                        context: Optional[Dict] = None) -> Response:
        """Send a message and get response"""
        if not self.active_session:
            self.create_session()
            
        # Add user message
        self.active_session.add_message(
            role=MessageRole.USER,
            content=content,
            metadata={"context": context}
        )
        
        # Prepare conversation history
        messages = self.active_session.get_context_window()
        
        # Add current context if provided
        if context:
            messages.append({
                "role": "system",
                "content": f"Current context: {json.dumps(context, indent=2)}"
            })
            
        try:
            # Get AI response
            response = await self._get_ai_response(messages)
            
            # Parse and structure the response
            structured_response = self._parse_response(response)
            
            # Add assistant response to history
            self.active_session.add_message(
                role=MessageRole.ASSISTANT,
                content=response,
                metadata={"structured": structured_response}
            )
            
            return structured_response
            
        except Exception as e:
            return Response(
                type=ResponseType.ERROR,
                content=str(e)
            )
            
    async def _get_ai_response(self, messages: List[Dict]) -> str:
        """Get response from AI model"""
        # Implementation depends on the AI API being used
        # This is a placeholder for the actual API call
        return "AI response placeholder"
        
    def _parse_response(self, response: str) -> Response:
        """Parse AI response into structured format"""
        # Look for code blocks
        code_blocks = re.findall(r'```[\w]*\n(.*?)```', response, re.DOTALL)
        
        # Look for command blocks
        command_blocks = re.findall(r'\$\s*(.*?)(?:\n|$)', response)
        
        if code_blocks:
            # Response contains code
            return Response(
                type=ResponseType.CODE,
                content=CodeAction(
                    action_type='modify',
                    file_path='',  # Would be extracted from context
                    content=code_blocks[0]
                )
            )
        elif command_blocks:
            # Response contains commands
            return Response(
                type=ResponseType.COMMAND,
                content=CommandAction(
                    command=command_blocks[0]
                )
            )
        else:
            # Regular text response
            return Response(
                type=ResponseType.TEXT,
                content=response
            )
            
    def get_session_stats(self, session_id: Optional[str] = None) -> Dict[str, Any]:
        """Get statistics for a session"""
        session = self.sessions.get(session_id, self.active_session)
        if not session:
            raise ValueError("No active session")
            
        messages = session.messages
        return {
            "total_messages": len(messages),
            "user_messages": sum(1 for m in messages if m.role == MessageRole.USER),
            "assistant_messages": sum(1 for m in messages if m.role == MessageRole.ASSISTANT),
            "start_time": session.start_time,
            "duration": datetime.now() - session.start_time,
            "context_updates": len([m for m in messages if m.metadata and "context" in m.metadata])
        }
        
    def get_session_summary(self, session_id: Optional[str] = None) -> Dict[str, Any]:
        """Get summary of session activities"""
        session = self.sessions.get(session_id, self.active_session)
        if not session:
            raise ValueError("No active session")
            
        code_actions = []
        command_actions = []
        
        for msg in session.messages:
            if msg.role == MessageRole.ASSISTANT and msg.metadata:
                response = msg.metadata.get("structured")
                if response:
                    if response.type == ResponseType.CODE:
                        code_actions.append(response.content)
                    elif response.type == ResponseType.COMMAND:
                        command_actions.append(response.content)
                        
        return {
            "code_actions": len(code_actions),
            "command_actions": len(command_actions),
            "last_action": session.messages[-1].timestamp if session.messages else None,
            "session_id": session.session_id
        }
