"""
Session Management for XoneAI Agents

A simple wrapper around existing stateful capabilities to provide a unified
session API for developers building stateful agent applications.
"""
from __future__ import annotations

import os
import uuid
import requests
import json
import time
from typing import Any, Dict, List, Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from .memory.memory import Memory
    from .knowledge.knowledge import Knowledge

from .agent import Agent


class Session:
    """
    A simple wrapper around XoneAI's existing stateful capabilities.
    
    Provides a unified API for:
    - Session management with persistent state
    - Memory operations (short-term, long-term, user-specific)
    - Knowledge base operations
    - Agent state management
    - Remote agent connectivity
    
    Examples:
        # Local session with agent
        session = Session(session_id="chat_123", user_id="user_456")
        agent = session.Agent(name="Assistant", role="Helpful AI")
        
        # Remote agent session (similar to Google ADK)
        session = Session(agent_url="192.168.1.10:8000/agent")
        response = session.chat("Hello from remote client!")
        
        # Save session state
        session.save_state({"conversation_topic": "AI research"})
    """

    def __init__(
        self, 
        session_id: Optional[str] = None,
        user_id: Optional[str] = None,
        agent_url: Optional[str] = None,
        memory_config: Optional[Dict[str, Any]] = None,
        knowledge_config: Optional[Dict[str, Any]] = None,
        timeout: int = 30
    ):
        """
        Initialize a new session with optional persistence or remote agent connectivity.
        
        Args:
            session_id: Unique session identifier. Auto-generated if None.
            user_id: User identifier for user-specific memory operations.
            agent_url: URL of remote agent for direct connectivity (e.g., "192.168.1.10:8000/agent")
            memory_config: Configuration for memory system (defaults to RAG)
            knowledge_config: Configuration for knowledge base system  
            timeout: HTTP timeout for remote agent calls (default: 30 seconds)
        """
        self.session_id = session_id or str(uuid.uuid4())[:8]
        self.user_id = user_id or "default_user"
        self.agent_url = agent_url
        self.timeout = timeout
        self.is_remote = agent_url is not None

        # Validate agent_url format
        if self.is_remote:
            if not self.agent_url.startswith(('http://', 'https://')):
                # Assume http if no protocol specified
                self.agent_url = f"http://{self.agent_url}"
            # Test connectivity to remote agent
            self._test_remote_connection()

        # Initialize memory with sensible defaults (only for local sessions)
        if not self.is_remote:
            default_memory_config = {
                "provider": "rag",
                "use_embedding": False,  # Disable embeddings to avoid OpenAI API key requirement
                "rag_db_path": f".xoneai/sessions/{self.session_id}/chroma_db"
            }
            if memory_config:
                default_memory_config.update(memory_config)
            self.memory_config = default_memory_config

            # Initialize knowledge with session-specific config
            default_knowledge_config = knowledge_config or {}
            self.knowledge_config = default_knowledge_config

            # Create session directory
            os.makedirs(f".xoneai/sessions/{self.session_id}", exist_ok=True)

            # Initialize components lazily
            self._memory = None
            self._knowledge = None
            self._agents_instance = None
            self._agents = {}  # Track agents and their chat histories
        else:
            # For remote sessions, disable local memory/knowledge
            self.memory_config = {}
            self.knowledge_config = {}
            self._memory = None
            self._knowledge = None
            self._agents_instance = None
            self._agents = {}  # Track agents and their chat histories

    @property
    def memory(self) -> Memory:
        """Lazy-loaded memory instance"""
        if self.is_remote:
            raise ValueError("Memory operations are not available for remote agent sessions")
        if self._memory is None:
            from .memory.memory import Memory
            self._memory = Memory(config=self.memory_config)
        return self._memory

    @property
    def knowledge(self) -> Knowledge:
        """Lazy-loaded knowledge instance"""
        if self.is_remote:
            raise ValueError("Knowledge operations are not available for remote agent sessions")
        if self._knowledge is None:
            from .knowledge.knowledge import Knowledge
            self._knowledge = Knowledge(config=self.knowledge_config)
        return self._knowledge

    def Agent(
        self,
        name: str,
        role: str = "Assistant", 
        instructions: Optional[str] = None,
        tools: Optional[List[Any]] = None,
        memory: bool = True,
        knowledge: Optional[List[str]] = None,
        **kwargs
    ) -> Agent:
        """
        Create an agent with session context.
        
        Args:
            name: Agent name
            role: Agent role
            instructions: Agent instructions
            tools: List of tools for the agent
            memory: Enable memory for the agent
            knowledge: Knowledge sources for the agent
            **kwargs: Additional agent parameters
        
        Returns:
            Configured Agent instance
            
        Raises:
            ValueError: If this is a remote session (use chat() instead)
        """
        if self.is_remote:
            raise ValueError("Cannot create local agents in remote sessions. Use chat() to communicate with the remote agent.")

        agent_kwargs = {
            "name": name,
            "role": role,
            "user_id": self.user_id,
            **kwargs
        }

        if instructions:
            agent_kwargs["instructions"] = instructions
        if tools:
            agent_kwargs["tools"] = tools
        if memory:
            agent_kwargs["memory"] = self.memory
        if knowledge:
            agent_kwargs["knowledge"] = knowledge
            agent_kwargs["knowledge_config"] = self.knowledge_config

        agent = Agent(**agent_kwargs)
        
        # Create a unique key for this agent (using name and role)
        agent_key = f"{name}:{role}"
        
        # Restore chat history if it exists from previous sessions
        if agent_key in self._agents:
            agent.chat_history = self._agents[agent_key].get("chat_history", [])
        else:
            # Try to restore from memory for backward compatibility
            restored_history = self._restore_agent_chat_history(agent_key)
            agent.chat_history = restored_history
        
        # Track the agent
        self._agents[agent_key] = {"agent": agent, "chat_history": agent.chat_history}
        
        return agent

    # Keep create_agent for backward compatibility
    def create_agent(self, *args, **kwargs) -> Agent:
        """Backward compatibility wrapper for Agent method"""
        return self.Agent(*args, **kwargs)

    def save_state(self, state_data: Dict[str, Any]) -> None:
        """
        Save session state data to memory.
        
        Args:
            state_data: Dictionary of state data to save
            
        Raises:
            ValueError: If this is a remote session
        """
        if self.is_remote:
            raise ValueError("State operations are not available for remote agent sessions")
            
        # Save agent chat histories first
        self._save_agent_chat_histories()
        
        # Save session state
        state_text = f"Session state: {state_data}"
        self.memory.store_short_term(
            text=state_text,
            metadata={
                "type": "session_state",
                "session_id": self.session_id,
                "user_id": self.user_id,
                **state_data
            }
        )

    def restore_state(self) -> Dict[str, Any]:
        """
        Restore session state from memory.
        
        Returns:
            Dictionary of restored state data
            
        Raises:
            ValueError: If this is a remote session
        """
        if self.is_remote:
            raise ValueError("State operations are not available for remote agent sessions")
        # Use content-based search for better SQLite compatibility
        results = self.memory.search_short_term(
            query="Session state:",
            limit=10  # Get more results to filter by session_id
        )

        # Restore agent chat histories first
        self._restore_agent_chat_histories()
        
        # Filter results by session_id in metadata
        for result in results:
            metadata = result.get("metadata", {})
            if (metadata.get("type") == "session_state" and 
                metadata.get("session_id") == self.session_id):
                # Extract state data from metadata (excluding system fields)
                state_data = {k: v for k, v in metadata.items() 
                             if k not in ["type", "session_id", "user_id"]}
                return state_data

        return {}

    def _restore_agent_chat_history(self, agent_key: str) -> List[Dict[str, Any]]:
        """
        Restore agent chat history from memory.
        
        Args:
            agent_key: Unique identifier for the agent
            
        Returns:
            List of chat history messages or empty list if not found
        """
        if self.is_remote:
            return []
        
        # Search for agent chat history in memory
        results = self.memory.search_short_term(
            query="Agent chat history for",
            limit=10
        )
        
        # Filter results by session_id and agent_key
        for result in results:
            metadata = result.get("metadata", {})
            if (metadata.get("type") == "agent_chat_history" and 
                metadata.get("session_id") == self.session_id and
                metadata.get("agent_key") == agent_key):
                # Extract chat history from metadata
                chat_history = metadata.get("chat_history", [])
                return chat_history
        
        return []

    def _restore_agent_chat_histories(self) -> None:
        """
        Restore all agent chat histories from memory.
        """
        if self.is_remote:
            return
        
        # Search for all agent chat histories in memory
        results = self.memory.search_short_term(
            query="Agent chat history for",
            limit=50  # Get many results to find all agents
        )
        
        # Filter and restore chat histories for this session
        for result in results:
            metadata = result.get("metadata", {})
            if (metadata.get("type") == "agent_chat_history" and 
                metadata.get("session_id") == self.session_id):
                agent_key = metadata.get("agent_key")
                chat_history = metadata.get("chat_history", [])
                
                if agent_key and chat_history:
                    # Store in _agents dict for later retrieval
                    self._agents[agent_key] = {
                        "agent": None,  # Will be populated when Agent is created
                        "chat_history": chat_history
                    }

    def _save_agent_chat_histories(self) -> None:
        """
        Save all agent chat histories to memory.
        """
        if self.is_remote:
            return
        
        for agent_key, agent_data in self._agents.items():
            agent = agent_data.get("agent")
            chat_history = None
            
            # Prioritize history from the live agent object, but fall back to restored history
            if agent and hasattr(agent, 'chat_history'):
                chat_history = agent.chat_history
                agent_data["chat_history"] = chat_history  # Ensure tracked history is up-to-date
            else:
                chat_history = agent_data.get("chat_history")
            
            if chat_history is not None:
                # Save to memory
                history_text = f"Agent chat history for {agent_key}"
                self.memory.store_short_term(
                    text=history_text,
                    metadata={
                        "type": "agent_chat_history",
                        "session_id": self.session_id,
                        "user_id": self.user_id,
                        "agent_key": agent_key,
                        "chat_history": chat_history
                    }
                )

    def get_state(self, key: str, default: Any = None) -> Any:
        """Get a specific state value"""
        state = self.restore_state()
        return state.get(key, default)

    def set_state(self, key: str, value: Any) -> None:
        """Set a specific state value"""
        current_state = self.restore_state()
        current_state[key] = value
        self.save_state(current_state)

    def increment_state(self, key: str, increment: int = 1, default: int = 0) -> None:
        """Increment a numeric state value"""
        current_value = self.get_state(key, default)
        self.set_state(key, current_value + increment)

    def add_memory(self, text: str, memory_type: str = "long", **metadata) -> None:
        """
        Add information to session memory.
        
        Args:
            text: Text to store
            memory_type: "short" or "long" term memory
            **metadata: Additional metadata
        """
        metadata.update({
            "session_id": self.session_id,
            "user_id": self.user_id
        })

        if memory_type == "short":
            self.memory.store_short_term(text, metadata=metadata)
        else:
            self.memory.store_long_term(text, metadata=metadata)

    def search_memory(self, query: str, memory_type: str = "long", limit: int = 5) -> List[Dict[str, Any]]:
        """
        Search session memory.
        
        Args:
            query: Search query
            memory_type: "short" or "long" term memory
            limit: Maximum results to return

        Returns:
            List of memory results
        """
        if memory_type == "short":
            return self.memory.search_short_term(query, limit=limit)
        return self.memory.search_long_term(query, limit=limit)

    def add_knowledge(self, source: str) -> None:
        """
        Add knowledge source to session.
        
        Args:
            source: File path, URL, or text content
        """
        self.knowledge.add(source, user_id=self.user_id, agent_id=self.session_id)

    def search_knowledge(self, query: str, limit: int = 5) -> List[Dict[str, Any]]:
        """
        Search session knowledge base.
        
        Args:
            query: Search query
            limit: Maximum results to return

        Returns:
            List of knowledge results
        """
        return self.knowledge.search(query, agent_id=self.session_id, limit=limit)

    def clear_memory(self, memory_type: str = "all") -> None:
        """
        Clear session memory.
        
        Args:
            memory_type: "short", "long", or "all"
        """
        if memory_type in ["short", "all"]:
            self.memory.reset_short_term()
        if memory_type in ["long", "all"]:
            self.memory.reset_long_term()

    def get_context(self, query: str, max_items: int = 3) -> str:
        """
        Build context from session memory and knowledge.
        
        Args:
            query: Query to build context for
            max_items: Maximum items per section

        Returns:
            Formatted context string
        """
        return self.memory.build_context_for_task(
            task_descr=query,
            user_id=self.user_id,
            max_items=max_items
        )

    def _test_remote_connection(self) -> None:
        """
        Test connectivity to the remote agent.
        
        Raises:
            ConnectionError: If unable to connect to the remote agent
        """
        try:
            # Try a simple GET request to check if the server is responding
            test_url = self.agent_url.rstrip('/') + '/health' if '/health' not in self.agent_url else self.agent_url
            response = requests.get(test_url, timeout=self.timeout)
            if response.status_code != 200:
                # If health endpoint fails, try the main endpoint
                response = requests.head(self.agent_url, timeout=self.timeout)
                if response.status_code not in [200, 405]:  # 405 = Method Not Allowed is OK
                    raise ConnectionError(f"Remote agent returned status code: {response.status_code}")
            print(f"✅ Successfully connected to remote agent at {self.agent_url}")
        except requests.exceptions.Timeout:
            raise ConnectionError(f"Timeout connecting to remote agent at {self.agent_url}")
        except requests.exceptions.ConnectionError:
            raise ConnectionError(f"Failed to connect to remote agent at {self.agent_url}")
        except Exception as e:
            raise ConnectionError(f"Error connecting to remote agent: {str(e)}")

    def chat(self, message: str, **kwargs) -> str:
        """
        Send a message to the remote agent or handle local session.
        
        Args:
            message: The message to send to the agent
            **kwargs: Additional parameters for the request
            
        Returns:
            The agent's response
            
        Raises:
            ValueError: If this is not a remote session
            ConnectionError: If unable to communicate with remote agent
        """
        if not self.is_remote:
            raise ValueError("chat() method is only available for remote agent sessions. Use Agent.chat() for local agents.")
        
        try:
            # Prepare the request payload
            payload = {
                "query": message,
                "session_id": self.session_id,
                "user_id": self.user_id,
                **kwargs
            }
            
            # Send POST request to the remote agent
            response = requests.post(
                self.agent_url,
                json=payload,
                headers={"Content-Type": "application/json"},
                timeout=self.timeout
            )
            
            # Check if request was successful
            response.raise_for_status()
            
            # Parse the response
            result = response.json()
            
            # Extract the agent's response
            if isinstance(result, dict):
                return result.get("response", str(result))
            else:
                return str(result)
                
        except requests.exceptions.Timeout:
            raise ConnectionError(f"Timeout communicating with remote agent at {self.agent_url}")
        except requests.exceptions.ConnectionError:
            raise ConnectionError(f"Failed to communicate with remote agent at {self.agent_url}")
        except requests.exceptions.HTTPError as e:
            raise ConnectionError(f"HTTP error from remote agent: {e}")
        except json.JSONDecodeError:
            # If response is not JSON, return the raw text
            return response.text
        except Exception as e:
            raise ConnectionError(f"Error communicating with remote agent: {str(e)}")

    def send_message(self, message: str, **kwargs) -> str:
        """
        Alias for chat() method to match Google ADK pattern.
        
        Args:
            message: The message to send to the agent
            **kwargs: Additional parameters for the request
            
        Returns:
            The agent's response
        """
        return self.chat(message, **kwargs)

    def __str__(self) -> str:
        if self.is_remote:
            return f"Session(id='{self.session_id}', user='{self.user_id}', remote_agent='{self.agent_url}')"
        return f"Session(id='{self.session_id}', user='{self.user_id}')"

    def __repr__(self) -> str:
        return self.__str__()