import os
import json
import datetime
from typing import Dict, List, Optional, Any

# Use absolute import for api_models
from api_service.api_models import (
    Conversation,
    UserSettings,
    Message,
    RoleCategoryPreset,
    GuildRoleCategoryConfig,
    UserCustomColorRole,
)

# ============= Database Class =============


class Database:
    def __init__(self, data_dir="data"):
        self.data_dir = data_dir
        self.conversations_file = os.path.join(data_dir, "conversations.json")
        self.settings_file = os.path.join(data_dir, "user_settings.json")
        self.tokens_file = os.path.join(data_dir, "user_tokens.json")
        self.role_presets_file = os.path.join(data_dir, "role_category_presets.json")
        self.guild_role_configs_file = os.path.join(
            data_dir, "guild_role_category_configs.json"
        )
        self.user_color_roles_file = os.path.join(
            data_dir, "user_custom_color_roles.json"
        )

        # Create data directory if it doesn't exist
        os.makedirs(data_dir, exist_ok=True)

        # In-memory storage
        self.conversations: Dict[str, Dict[str, Conversation]] = (
            {}
        )  # user_id -> conversation_id -> Conversation
        self.user_settings: Dict[str, UserSettings] = {}  # user_id -> UserSettings
        self.user_tokens: Dict[str, Dict[str, Any]] = {}  # user_id -> token_data
        self.role_category_presets: Dict[str, RoleCategoryPreset] = (
            {}
        )  # preset_id -> RoleCategoryPreset
        self.guild_role_category_configs: Dict[str, List[GuildRoleCategoryConfig]] = (
            {}
        )  # guild_id -> List[GuildRoleCategoryConfig]
        self.user_custom_color_roles: Dict[str, Dict[str, UserCustomColorRole]] = (
            {}
        )  # guild_id -> user_id -> UserCustomColorRole

        # Load data from files
        self.load_data()

    def load_data(self):
        """Load all data from files"""
        self.load_conversations()
        self.load_user_settings()
        self.load_user_tokens()
        self.load_role_category_presets()
        self.load_guild_role_category_configs()
        self.load_user_custom_color_roles()

    def save_data(self):
        """Save all data to files"""
        self.save_conversations()
        self.save_all_user_settings()
        self.save_user_tokens()
        self.save_role_category_presets()
        self.save_guild_role_category_configs()
        self.save_user_custom_color_roles()

    # ============= Role Selector Data Load/Save Methods =============

    def load_role_category_presets(self):
        """Load role category presets from file"""
        if os.path.exists(self.role_presets_file):
            try:
                with open(self.role_presets_file, "r", encoding="utf-8") as f:
                    data = json.load(f)
                    self.role_category_presets = {
                        preset_id: RoleCategoryPreset.model_validate(preset_data)
                        for preset_id, preset_data in data.items()
                    }
                print(
                    f"Loaded {len(self.role_category_presets)} role category presets."
                )
            except Exception as e:
                print(f"Error loading role category presets: {e}")
                self.role_category_presets = {}
        else:
            self.role_category_presets = {}  # Initialize if file doesn't exist

    def save_role_category_presets(self):
        """Save role category presets to file"""
        try:
            serializable_data = {
                preset_id: preset.model_dump()
                for preset_id, preset in self.role_category_presets.items()
            }
            with open(self.role_presets_file, "w", encoding="utf-8") as f:
                json.dump(
                    serializable_data, f, indent=2, default=str, ensure_ascii=False
                )
        except Exception as e:
            print(f"Error saving role category presets: {e}")

    def load_guild_role_category_configs(self):
        """Load guild role category configs from file"""
        if os.path.exists(self.guild_role_configs_file):
            try:
                with open(self.guild_role_configs_file, "r", encoding="utf-8") as f:
                    data = json.load(f)
                    self.guild_role_category_configs = {
                        guild_id: [
                            GuildRoleCategoryConfig.model_validate(config_data)
                            for config_data in configs_list
                        ]
                        for guild_id, configs_list in data.items()
                    }
                print(
                    f"Loaded guild role category configs for {len(self.guild_role_category_configs)} guilds."
                )
            except Exception as e:
                print(f"Error loading guild role category configs: {e}")
                self.guild_role_category_configs = {}
        else:
            self.guild_role_category_configs = {}

    def save_guild_role_category_configs(self):
        """Save guild role category configs to file"""
        try:
            serializable_data = {
                guild_id: [config.model_dump() for config in configs_list]
                for guild_id, configs_list in self.guild_role_category_configs.items()
            }
            with open(self.guild_role_configs_file, "w", encoding="utf-8") as f:
                json.dump(
                    serializable_data, f, indent=2, default=str, ensure_ascii=False
                )
        except Exception as e:
            print(f"Error saving guild role category configs: {e}")

    def load_user_custom_color_roles(self):
        """Load user custom color roles from file"""
        if os.path.exists(self.user_color_roles_file):
            try:
                with open(self.user_color_roles_file, "r", encoding="utf-8") as f:
                    data = json.load(f)
                    self.user_custom_color_roles = {
                        guild_id: {
                            user_id: UserCustomColorRole.model_validate(role_data)
                            for user_id, role_data in user_roles.items()
                        }
                        for guild_id, user_roles in data.items()
                    }
                print(
                    f"Loaded user custom color roles for {len(self.user_custom_color_roles)} guilds."
                )
            except Exception as e:
                print(f"Error loading user custom color roles: {e}")
                self.user_custom_color_roles = {}
        else:
            self.user_custom_color_roles = {}

    def save_user_custom_color_roles(self):
        """Save user custom color roles to file"""
        try:
            serializable_data = {
                guild_id: {
                    user_id: role.model_dump() for user_id, role in user_roles.items()
                }
                for guild_id, user_roles in self.user_custom_color_roles.items()
            }
            with open(self.user_color_roles_file, "w", encoding="utf-8") as f:
                json.dump(
                    serializable_data, f, indent=2, default=str, ensure_ascii=False
                )
        except Exception as e:
            print(f"Error saving user custom color roles: {e}")

    # ============= Existing Data Load/Save Methods =============

    def load_conversations(self):
        """Load conversations from file"""
        if os.path.exists(self.conversations_file):
            try:
                with open(self.conversations_file, "r", encoding="utf-8") as f:
                    data = json.load(f)
                    # Convert to Conversation objects
                    self.conversations = {
                        user_id: {
                            conv_id: Conversation.model_validate(conv_data)
                            for conv_id, conv_data in user_convs.items()
                        }
                        for user_id, user_convs in data.items()
                    }
                print(f"Loaded conversations for {len(self.conversations)} users")
            except Exception as e:
                print(f"Error loading conversations: {e}")
                self.conversations = {}

    def save_conversations(self):
        """Save conversations to file"""
        try:
            # Convert to JSON-serializable format
            serializable_data = {
                user_id: {
                    conv_id: conv.model_dump() for conv_id, conv in user_convs.items()
                }
                for user_id, user_convs in self.conversations.items()
            }
            with open(self.conversations_file, "w", encoding="utf-8") as f:
                json.dump(
                    serializable_data, f, indent=2, default=str, ensure_ascii=False
                )
        except Exception as e:
            print(f"Error saving conversations: {e}")

    def load_user_settings(self):
        """Load user settings from file"""
        if os.path.exists(self.settings_file):
            try:
                with open(self.settings_file, "r", encoding="utf-8") as f:
                    data = json.load(f)
                    # Convert to UserSettings objects
                    self.user_settings = {
                        user_id: UserSettings.model_validate(settings_data)
                        for user_id, settings_data in data.items()
                    }
                print(f"Loaded settings for {len(self.user_settings)} users")
            except Exception as e:
                print(f"Error loading user settings: {e}")
                self.user_settings = {}

    def save_all_user_settings(self):
        """Save all user settings to file"""
        try:
            # Convert to JSON-serializable format
            serializable_data = {
                user_id: settings.model_dump()
                for user_id, settings in self.user_settings.items()
            }
            with open(self.settings_file, "w", encoding="utf-8") as f:
                json.dump(
                    serializable_data, f, indent=2, default=str, ensure_ascii=False
                )
        except Exception as e:
            print(f"Error saving user settings: {e}")

    # ============= Conversation Methods =============

    def get_user_conversations(self, user_id: str) -> List[Conversation]:
        """Get all conversations for a user"""
        return list(self.conversations.get(user_id, {}).values())

    def get_conversation(
        self, user_id: str, conversation_id: str
    ) -> Optional[Conversation]:
        """Get a specific conversation for a user"""
        return self.conversations.get(user_id, {}).get(conversation_id)

    def save_conversation(
        self, user_id: str, conversation: Conversation
    ) -> Conversation:
        """Save a conversation for a user"""
        # Update the timestamp
        conversation.updated_at = datetime.datetime.now()

        # Initialize user's conversations dict if it doesn't exist
        if user_id not in self.conversations:
            self.conversations[user_id] = {}

        # Save the conversation
        self.conversations[user_id][conversation.id] = conversation

        # Save to disk
        self.save_conversations()

        return conversation

    def delete_conversation(self, user_id: str, conversation_id: str) -> bool:
        """Delete a conversation for a user"""
        if (
            user_id in self.conversations
            and conversation_id in self.conversations[user_id]
        ):
            del self.conversations[user_id][conversation_id]
            self.save_conversations()
            return True
        return False

    # ============= Role Category Preset Methods =============

    def get_role_category_preset(self, preset_id: str) -> Optional[RoleCategoryPreset]:
        """Get a specific role category preset by ID."""
        return self.role_category_presets.get(preset_id)

    def get_all_role_category_presets(self) -> List[RoleCategoryPreset]:
        """Get all role category presets."""
        return list(self.role_category_presets.values())

    def save_role_category_preset(
        self, preset: RoleCategoryPreset
    ) -> RoleCategoryPreset:
        """Save a role category preset."""
        self.role_category_presets[preset.id] = preset
        self.save_role_category_presets()
        return preset

    def delete_role_category_preset(self, preset_id: str) -> bool:
        """Delete a role category preset."""
        if preset_id in self.role_category_presets:
            del self.role_category_presets[preset_id]
            self.save_role_category_presets()
            return True
        return False

    # ============= Guild Role Category Config Methods =============

    def get_guild_role_category_configs(
        self, guild_id: str
    ) -> List[GuildRoleCategoryConfig]:
        """Get all role category configurations for a specific guild."""
        return self.guild_role_category_configs.get(guild_id, [])

    def get_all_guild_role_category_configs(
        self,
    ) -> Dict[str, List[GuildRoleCategoryConfig]]:
        """Get all role category configurations for all guilds."""
        return self.guild_role_category_configs

    def get_guild_role_category_config(
        self, guild_id: str, category_id: str
    ) -> Optional[GuildRoleCategoryConfig]:
        """Get a specific role category configuration for a guild."""
        for config in self.get_guild_role_category_configs(guild_id):
            if config.category_id == category_id:
                return config
        return None

    def save_guild_role_category_config(
        self, config: GuildRoleCategoryConfig
    ) -> GuildRoleCategoryConfig:
        """Save a guild's role category configuration."""
        guild_id = config.guild_id
        if guild_id not in self.guild_role_category_configs:
            self.guild_role_category_configs[guild_id] = []

        # Remove existing config with the same category_id if it exists, then add the new/updated one
        self.guild_role_category_configs[guild_id] = [
            c
            for c in self.guild_role_category_configs[guild_id]
            if c.category_id != config.category_id
        ]
        self.guild_role_category_configs[guild_id].append(config)
        self.save_guild_role_category_configs()
        return config

    def delete_guild_role_category_config(
        self, guild_id: str, category_id: str
    ) -> bool:
        """Delete a specific role category configuration for a guild."""
        if guild_id in self.guild_role_category_configs:
            initial_len = len(self.guild_role_category_configs[guild_id])
            self.guild_role_category_configs[guild_id] = [
                c
                for c in self.guild_role_category_configs[guild_id]
                if c.category_id != category_id
            ]
            if len(self.guild_role_category_configs[guild_id]) < initial_len:
                self.save_guild_role_category_configs()
                return True
        return False

    # ============= User Custom Color Role Methods =============

    def get_user_custom_color_role(
        self, guild_id: str, user_id: str
    ) -> Optional[UserCustomColorRole]:
        """Get a user's custom color role in a specific guild."""
        return self.user_custom_color_roles.get(guild_id, {}).get(user_id)

    def save_user_custom_color_role(
        self, color_role: UserCustomColorRole
    ) -> UserCustomColorRole:
        """Save a user's custom color role."""
        guild_id = color_role.guild_id
        user_id = color_role.user_id

        if guild_id not in self.user_custom_color_roles:
            self.user_custom_color_roles[guild_id] = {}

        color_role.last_updated = datetime.datetime.now()
        self.user_custom_color_roles[guild_id][user_id] = color_role
        self.save_user_custom_color_roles()
        return color_role

    def delete_user_custom_color_role(self, guild_id: str, user_id: str) -> bool:
        """Delete a user's custom color role in a specific guild."""
        if (
            guild_id in self.user_custom_color_roles
            and user_id in self.user_custom_color_roles[guild_id]
        ):
            del self.user_custom_color_roles[guild_id][user_id]
            self.save_user_custom_color_roles()
            return True
        return False

    # ============= User Settings Methods =============

    def get_user_settings(self, user_id: str) -> UserSettings:
        """Get settings for a user, creating default settings if they don't exist"""
        if user_id not in self.user_settings:
            self.user_settings[user_id] = UserSettings()

        return self.user_settings[user_id]

    def save_user_settings(self, user_id: str, settings: UserSettings) -> UserSettings:
        """Save settings for a user"""
        # Update the timestamp
        settings.last_updated = datetime.datetime.now()

        # Save the settings
        self.user_settings[user_id] = settings

        # Save to disk
        self.save_all_user_settings()

        return settings

    # ============= User Tokens Methods =============

    def load_user_tokens(self):
        """Load user tokens from file"""
        if os.path.exists(self.tokens_file):
            try:
                with open(self.tokens_file, "r", encoding="utf-8") as f:
                    self.user_tokens = json.load(f)
                print(f"Loaded tokens for {len(self.user_tokens)} users")
            except Exception as e:
                print(f"Error loading user tokens: {e}")
                self.user_tokens = {}

    def save_user_tokens(self):
        """Save user tokens to file"""
        try:
            with open(self.tokens_file, "w", encoding="utf-8") as f:
                json.dump(
                    self.user_tokens, f, indent=2, default=str, ensure_ascii=False
                )
        except Exception as e:
            print(f"Error saving user tokens: {e}")

    def get_user_token(self, user_id: str) -> Optional[Dict[str, Any]]:
        """Get token data for a user"""
        return self.user_tokens.get(user_id)

    def save_user_token(
        self, user_id: str, token_data: Dict[str, Any]
    ) -> Dict[str, Any]:
        """Save token data for a user"""
        # Add the time when the token was saved
        token_data["saved_at"] = datetime.datetime.now().isoformat()

        # Save the token data
        self.user_tokens[user_id] = token_data

        # Save to disk
        self.save_user_tokens()

        return token_data

    def delete_user_token(self, user_id: str) -> bool:
        """Delete token data for a user"""
        if user_id in self.user_tokens:
            del self.user_tokens[user_id]
            self.save_user_tokens()
            return True
        return False