import json
import logging
import os
import threading
from typing import List, Optional, Dict, Any
from pathlib import Path
from datetime import datetime

import portalocker

from .models import ThoughtData, ThoughtStage
from .logging_conf import configure_logging
from .storage_utils import prepare_thoughts_for_serialization, save_thoughts_to_file, load_thoughts_from_file

logger = configure_logging("sequential-thinking.storage")


class ThoughtStorage:
    """Storage manager for thought data."""

    def __init__(self, storage_dir: Optional[str] = None):
        """Initialize the storage manager.

        Args:
            storage_dir: Directory to store thought data files. If None, uses a default directory.
        """
        if storage_dir is None:
            # Use user's home directory by default
            home_dir = Path.home()
            self.storage_dir = home_dir / ".mcp_coaia_sequential_thinking"
        else:
            self.storage_dir = Path(storage_dir)

        # Create storage directory if it doesn't exist
        self.storage_dir.mkdir(parents=True, exist_ok=True)

        # Default session file
        self.current_session_file = self.storage_dir / "current_session.json"
        self.lock_file = self.storage_dir / "current_session.lock"

        # Thread safety
        self._lock = threading.RLock()
        self.thought_history: List[ThoughtData] = []

        # Load existing session if available
        self._load_session()

    def _load_session(self) -> None:
        """Load thought history from the current session file if it exists."""
        with self._lock:
            # Use the utility function to handle loading with proper error handling
            self.thought_history = load_thoughts_from_file(self.current_session_file, self.lock_file)

    def _save_session(self) -> None:
        """Save the current thought history to the session file."""
        # Use thread lock to ensure consistent data
        with self._lock:
            # Use utility functions to prepare and save thoughts
            thoughts_with_ids = prepare_thoughts_for_serialization(self.thought_history)
        
        # Save to file with proper locking
        save_thoughts_to_file(self.current_session_file, thoughts_with_ids, self.lock_file)

    def add_thought(self, thought: ThoughtData) -> None:
        """Add a thought to the history and save the session.

        Args:
            thought: The thought data to add
        """
        with self._lock:
            self.thought_history.append(thought)
        self._save_session()

    def get_all_thoughts(self) -> List[ThoughtData]:
        """Get all thoughts in the current session.

        Returns:
            List[ThoughtData]: All thoughts in the current session
        """
        with self._lock:
            # Return a copy to avoid external modification
            return list(self.thought_history)

    def get_thoughts_by_stage(self, stage: ThoughtStage) -> List[ThoughtData]:
        """Get all thoughts in a specific stage.

        Args:
            stage: The thinking stage to filter by

        Returns:
            List[ThoughtData]: Thoughts in the specified stage
        """
        with self._lock:
            return [t for t in self.thought_history if t.stage == stage]

    def clear_history(self) -> None:
        """Clear the thought history and save the empty session."""
        with self._lock:
            self.thought_history.clear()
        self._save_session()

    def export_session(self, file_path: str) -> None:
        """Export the current session to a file.

        Args:
            file_path: Path to save the exported session
        """
        with self._lock:
            # Use utility function to prepare thoughts for serialization
            thoughts_with_ids = prepare_thoughts_for_serialization(self.thought_history)
            
            # Create export-specific metadata
            metadata = {
                "exportedAt": datetime.now().isoformat(),
                "metadata": {
                    "totalThoughts": len(self.thought_history),
                    "stages": {
                        stage.value: len([t for t in self.thought_history if t.stage == stage])
                        for stage in ThoughtStage
                    }
                }
            }
        
        # Convert string path to Path object for compatibility with utility
        file_path_obj = Path(file_path)
        lock_file = file_path_obj.with_suffix('.lock')
        
        # Use utility function to save with proper locking
        save_thoughts_to_file(file_path_obj, thoughts_with_ids, lock_file, metadata)

    def import_session(self, file_path: str) -> None:
        """Import a session from a file.

        Args:
            file_path: Path to the file to import

        Raises:
            FileNotFoundError: If the file doesn't exist
            json.JSONDecodeError: If the file is not valid JSON
            KeyError: If the file doesn't contain valid thought data
        """
        # Convert string path to Path object for compatibility with utility
        file_path_obj = Path(file_path)
        lock_file = file_path_obj.with_suffix('.lock')
        
        # Use utility function to load thoughts with proper error handling
        thoughts = load_thoughts_from_file(file_path_obj, lock_file)
        
        with self._lock:
            self.thought_history = thoughts

        self._save_session()
