"""
ThinkingSession module for managing sequential thoughts through MCP.
"""

import uuid
from typing import Dict, List, Optional, Union, Any

from mcp import ClientSession
import os


class ThinkingSession:
    """
    Manages a sequential thinking session that records thoughts and allows for retrieval.

    Uses MCP to store and retrieve thoughts in a ChromaDB collection.
    """

    def __init__(self, client: ClientSession = None, session_id: str = None, mcp_server_url: Optional[str] = None):
        """
        Initialize a thinking session.

        Args:
            client: ClientSession instance for communicating with MCP server
            session_id: Optional session ID. If not provided, a new UUID will be generated.
            mcp_server_url: Optional MCP server URL. If not provided, it will be read from the environment variable MCP_SERVER_URL.
        """
        if client:
            self.client = client
        else:
            server_url = mcp_server_url or os.getenv("MCP_SERVER_URL", "http://localhost:8000")
            if not server_url:
                raise ValueError("MCP_SERVER_URL environment variable not set and no client or URL provided.")
            self.client = ClientSession(server_url)

        self.session_id = session_id or str(uuid.uuid4())

    def record_thought(
        self,
        thought: str,
        thought_number: int,
        total_thoughts: int,
        branch_id: str = "",
        branch_from_thought: int = 0,
        next_thought_needed: bool = False,
    ) -> Dict[str, Any]:
        """
        Record a single thought in the thinking session.

        Args:
            thought: Content of the thought
            thought_number: Sequential number of this thought
            total_thoughts: Total anticipated number of thoughts in this sequence
            branch_id: Optional identifier for a branch within the session
            branch_from_thought: If creating a new branch, the parent thought number it originates from
            next_thought_needed: Flag indicating if a subsequent thought is expected

        Returns:
            Response from MCP server including the session ID
        """
        response = self.client.mcp_chroma_dev_chroma_sequential_thinking(
            thought=thought,
            thought_number=thought_number,
            total_thoughts=total_thoughts,
            session_id=self.session_id,
            branch_id=branch_id,
            branch_from_thought=branch_from_thought,
            next_thought_needed=next_thought_needed,
        )

        # Update session_id if it was auto-generated by the server
        if not self.session_id and "session_id" in response:
            self.session_id = response["session_id"]

        return response

    def find_similar_thoughts(
        self, query: str, n_results: int = 5, threshold: float = -1.0, include_branches: bool = True
    ) -> List[Dict[str, Any]]:
        """
        Find thoughts semantically similar to the query.

        Args:
            query: Text to search for similar thoughts
            n_results: Maximum number of similar thoughts to return
            threshold: Similarity score threshold (0.0 to 1.0). Lower distance is more similar.
                      -1.0 to use default.
            include_branches: Whether to include thoughts from branches in the search

        Returns:
            List of similar thoughts with metadata
        """
        response = self.client.mcp_chroma_dev_chroma_find_similar_thoughts(
            query=query,
            session_id=self.session_id,
            n_results=n_results,
            threshold=threshold,
            include_branches=include_branches,
        )

        return response.get("similar_thoughts", [])

    def get_session_summary(self, include_branches: bool = True) -> Dict[str, Any]:
        """
        Get a summary of all thoughts in this thinking session.

        Args:
            include_branches: Whether to include thoughts from branches in the summary

        Returns:
            Dictionary containing session information and all thoughts
        """
        response = self.client.mcp_chroma_dev_chroma_get_session_summary(
            session_id=self.session_id, include_branches=include_branches
        )

        return response

    @classmethod
    def find_similar_sessions(
        cls, query: str, client: ClientSession = None, n_results: int = 5, threshold: float = -1.0
    ) -> List[Dict[str, Any]]:
        """
        Find thinking sessions similar to the query text.

        Args:
            query: Text to search for similar sessions
            client: ClientSession instance for communicating with MCP server
            n_results: Maximum number of similar sessions to return
            threshold: Similarity score threshold (0.0 to 1.0). Lower distance is more similar.
                      -1.0 to use default.

        Returns:
            List of similar sessions with metadata
        """
        client = client or ClientSession(os.getenv("MCP_SERVER_URL", "http://localhost:8000"))
        response = client.mcp_chroma_dev_chroma_find_similar_sessions(
            query=query, n_results=n_results, threshold=threshold
        )

        return response.get("similar_sessions", [])
