@file:OptIn(InternalSerializationApi::class)

package com.example.notesai.data.local

import android.content.Context
import androidx.datastore.core.DataStore
import androidx.datastore.preferences.core.Preferences
import androidx.datastore.preferences.core.edit
import androidx.datastore.preferences.core.emptyPreferences
import androidx.datastore.preferences.core.stringPreferencesKey
import androidx.datastore.preferences.preferencesDataStore
import com.example.notesai.data.model.Note
import com.example.notesai.data.model.Category
import com.example.notesai.data.model.ChatHistory
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.catch
import kotlinx.coroutines.flow.map
import kotlinx.serialization.InternalSerializationApi
import kotlinx.serialization.encodeToString
import kotlinx.serialization.json.Json
import kotlinx.serialization.Serializable
import java.io.IOException

val Context.dataStore: DataStore<Preferences> by preferencesDataStore(name = "notes_prefs")

@Serializable
data class SerializableCategory(
    val id: String,
    val name: String,
    val gradientStart: Long,
    val gradientEnd: Long,
    val timestamp: Long,
    val isDeleted: Boolean = false,
    val isPinned: Boolean = false // NEW
)

@Serializable
data class SerializableNote(
    val id: String,
    val categoryId: String,
    val title: String,
    val description: String = "",
    val content: String = "",
    val timestamp: Long,
    val isArchived: Boolean = false,
    val isDeleted: Boolean = false,
    val isPinned: Boolean = false
)

class DataStoreManager(private val context: Context) {
    companion object {
        val CATEGORIES_KEY = stringPreferencesKey("categories")
        val NOTES_KEY = stringPreferencesKey("notes")
        val CHAT_HISTORY_KEY = stringPreferencesKey("chat_history")
        val THEME_KEY = stringPreferencesKey("theme") // "dark" or "light"
    }

    private val json = Json {
        ignoreUnknownKeys = true
        encodeDefaults = true
    }

    val categoriesFlow: Flow<List<Category>> = context.dataStore.data
        .catch { exception ->
            if (exception is IOException) {
                emit(emptyPreferences())
            } else {
                throw exception
            }
        }
        .map { preferences ->
            val jsonString = preferences[CATEGORIES_KEY] ?: "[]"
            try {
                json.decodeFromString<List<SerializableCategory>>(jsonString).map {
                    Category(
                        id = it.id,
                        name = it.name,
                        gradientStart = it.gradientStart,
                        gradientEnd = it.gradientEnd,
                        timestamp = it.timestamp,
                        isDeleted = it.isDeleted,
                        isPinned = it.isPinned // NEW
                    )
                }
            } catch (e: Exception) {
                e.printStackTrace()
                emptyList()
            }
        }

    val notesFlow: Flow<List<Note>> = context.dataStore.data
        .catch { exception ->
            if (exception is IOException) {
                emit(emptyPreferences())
            } else {
                throw exception
            }
        }
        .map { preferences ->
            val jsonString = preferences[NOTES_KEY] ?: "[]"
            try {
                json.decodeFromString<List<SerializableNote>>(jsonString).map {
                    Note(
                        id = it.id,
                        categoryId = it.categoryId,
                        title = it.title,
                        description = it.description,
                        content = it.content,
                        timestamp = it.timestamp,
                        isPinned = it.isPinned,
                        isArchived = it.isArchived,
                        isDeleted = it.isDeleted
                    )
                }
            } catch (e: Exception) {
                e.printStackTrace()
                emptyList()
            }
        }

    // Chat History Flow
    val chatHistoryFlow: Flow<List<ChatHistory>> = context.dataStore.data
        .catch { exception ->
            if (exception is IOException) {
                emit(emptyPreferences())
            } else {
                throw exception
            }
        }
        .map { preferences ->
            val jsonString = preferences[CHAT_HISTORY_KEY] ?: "[]"
            try {
                json.decodeFromString<List<ChatHistory>>(jsonString)
                    .filter { !it.isDeleted }
                    .sortedByDescending { it.timestamp }
            } catch (e: Exception) {
                e.printStackTrace()
                emptyList()
            }
        }

    suspend fun saveCategories(categories: List<Category>) {
        try {
            context.dataStore.edit { preferences ->
                val serializable = categories.map {
                    SerializableCategory(
                        id = it.id,
                        name = it.name,
                        gradientStart = it.gradientStart,
                        gradientEnd = it.gradientEnd,
                        timestamp = it.timestamp,
                        isDeleted = it.isDeleted,
                        isPinned = it.isPinned // NEW
                    )
                }
                preferences[CATEGORIES_KEY] = json.encodeToString(serializable)
            }
        } catch (e: Exception) {
            e.printStackTrace()
        }
    }

    suspend fun saveNotes(notes: List<Note>) {
        try {
            context.dataStore.edit { preferences ->
                val serializable = notes.map {
                    SerializableNote(
                        id = it.id,
                        categoryId = it.categoryId,
                        title = it.title,
                        description = it.description,
                        content = it.content,
                        timestamp = it.timestamp,
                        isPinned = it.isPinned,
                        isArchived = it.isArchived,
                        isDeleted = it.isDeleted
                    )
                }
                preferences[NOTES_KEY] = json.encodeToString(serializable)
            }
        } catch (e: Exception) {
            e.printStackTrace()
        }
    }

    // Save Chat History
    suspend fun saveChatHistory(chatHistoryList: List<ChatHistory>) {
        try {
            context.dataStore.edit { preferences ->
                preferences[CHAT_HISTORY_KEY] = json.encodeToString(chatHistoryList)
            }
        } catch (e: Exception) {
            e.printStackTrace()
        }
    }

    // Add new chat history
    suspend fun addChatHistory(chatHistory: ChatHistory) {
        try {
            context.dataStore.edit { preferences ->
                val jsonString = preferences[CHAT_HISTORY_KEY] ?: "[]"
                val currentList = json.decodeFromString<List<ChatHistory>>(jsonString).toMutableList()

                // Check if already exists, update instead
                val existingIndex = currentList.indexOfFirst { it.id == chatHistory.id }
                if (existingIndex != -1) {
                    currentList[existingIndex] = chatHistory
                } else {
                    currentList.add(0, chatHistory) // Add to beginning
                }

                preferences[CHAT_HISTORY_KEY] = json.encodeToString(currentList)
            }
        } catch (e: Exception) {
            e.printStackTrace()
        }
    }

    // Delete chat history (soft delete)
    suspend fun deleteChatHistory(historyId: String) {
        try {
            context.dataStore.edit { preferences ->
                val jsonString = preferences[CHAT_HISTORY_KEY] ?: "[]"
                val currentList = json.decodeFromString<List<ChatHistory>>(jsonString).toMutableList()
                val updatedList = currentList.map {
                    if (it.id == historyId) it.copy(isDeleted = true) else it
                }
                preferences[CHAT_HISTORY_KEY] = json.encodeToString(updatedList)
            }
        } catch (e: Exception) {
            e.printStackTrace()
        }
    }

    // NEW: Update chat history title
    suspend fun updateChatHistoryTitle(historyId: String, newTitle: String) {
        try {
            context.dataStore.edit { preferences ->
                val jsonString = preferences[CHAT_HISTORY_KEY] ?: "[]"
                val currentList = json.decodeFromString<List<ChatHistory>>(jsonString).toMutableList()
                val updatedList = currentList.map {
                    if (it.id == historyId) {
                        it.copy(customTitle = newTitle)
                    } else {
                        it
                    }
                }
                preferences[CHAT_HISTORY_KEY] = json.encodeToString(updatedList)
            }
        } catch (e: Exception) {
            e.printStackTrace()
        }
    }

    // Theme Preference Flow
    val themeFlow: Flow<String> = context.dataStore.data
        .catch { exception ->
            if (exception is IOException) {
                emit(emptyPreferences())
            } else {
                throw exception
            }
        }
        .map { preferences ->
            preferences[THEME_KEY] ?: "dark" // Default dark theme
        }

    // Save Theme Preference
    suspend fun saveTheme(theme: String) {
        try {
            context.dataStore.edit { preferences ->
                preferences[THEME_KEY] = theme
            }
        } catch (e: Exception) {
            e.printStackTrace()
        }
    }
}