package com.example.notesai

import android.os.Bundle
import androidx.activity.ComponentActivity
import androidx.activity.compose.setContent
import androidx.compose.animation.*
import androidx.compose.animation.core.Spring
import androidx.compose.animation.core.spring
import androidx.compose.foundation.layout.*
import androidx.compose.material.icons.Icons
import androidx.compose.material.icons.filled.*
import androidx.compose.material3.*
import androidx.compose.runtime.*
import androidx.compose.ui.Modifier
import androidx.compose.ui.graphics.Color
import androidx.compose.ui.platform.LocalContext
import androidx.compose.ui.text.font.FontWeight
import androidx.compose.ui.unit.dp
import androidx.compose.ui.unit.sp
import androidx.compose.ui.zIndex
import com.example.notesai.data.local.DataStoreManager
import com.example.notesai.presentation.components.DrawerMenu
import com.example.notesai.presentation.components.ModernBottomBar
import com.example.notesai.presentation.components.ModernTopBar
import com.example.notesai.presentation.dialogs.CategoryDialog
import com.example.notesai.presentation.dialogs.NoteDialog
import com.example.notesai.presentation.screens.ai.AIHelperScreen
import com.example.notesai.presentation.screens.archive.ArchiveScreen
import com.example.notesai.presentation.screens.main.MainScreen
import com.example.notesai.presentation.screens.note.EditableFullScreenNoteView
import com.example.notesai.presentation.screens.starred.StarredNotesScreen
import com.example.notesai.presentation.screens.trash.TrashScreen
import com.example.notesai.data.model.Note
import com.example.notesai.data.model.Category
import com.example.notesai.util.AppColors
import com.example.notesai.util.Constants
import kotlinx.coroutines.launch

class MainActivity : ComponentActivity() {
    override fun onCreate(savedInstanceState: Bundle?) {
        super.onCreate(savedInstanceState)
        setContent {
            NotesAppTheme()
        }
    }
}

@Composable
fun NotesAppTheme(content: @Composable () -> Unit = { NotesApp() }) {
    val context = LocalContext.current
    val dataStoreManager = remember { DataStoreManager(context) }
    var isDarkTheme by remember { mutableStateOf(true) }

    // Load theme preference
    LaunchedEffect(Unit) {
        dataStoreManager.themeFlow.collect { theme ->
            isDarkTheme = theme == "dark"
            AppColors.setTheme(isDarkTheme)
        }
    }

    // Create dynamic color scheme based on theme
    val colorScheme = if (isDarkTheme) {
        darkColorScheme(
            primary = AppColors.Primary,
            onPrimary = Color.White,
            primaryContainer = AppColors.Primary.copy(alpha = 0.3f),
            onPrimaryContainer = Color.White,
            secondary = AppColors.Secondary,
            onSecondary = Color.White,
            secondaryContainer = AppColors.Secondary.copy(alpha = 0.3f),
            onSecondaryContainer = Color.White,
            background = AppColors.Background,
            onBackground = AppColors.OnBackground,
            surface = AppColors.Surface,
            onSurface = AppColors.OnSurface,
            surfaceVariant = AppColors.SurfaceVariant,
            onSurfaceVariant = AppColors.OnSurfaceVariant,
            error = AppColors.Error,
            onError = Color.White,
            outline = AppColors.Border,
            outlineVariant = AppColors.Divider
        )
    } else {
        lightColorScheme(
            primary = AppColors.Primary,
            onPrimary = Color.White,
            primaryContainer = AppColors.Primary.copy(alpha = 0.1f),
            onPrimaryContainer = AppColors.Primary,
            secondary = AppColors.Secondary,
            onSecondary = Color.White,
            secondaryContainer = AppColors.Secondary.copy(alpha = 0.1f),
            onSecondaryContainer = AppColors.Secondary,
            background = AppColors.Background,
            onBackground = AppColors.OnBackground,
            surface = AppColors.Surface,
            onSurface = AppColors.OnSurface,
            surfaceVariant = AppColors.SurfaceVariant,
            onSurfaceVariant = AppColors.OnSurfaceVariant,
            error = AppColors.Error,
            onError = Color.White,
            outline = AppColors.Border,
            outlineVariant = AppColors.Divider
        )
    }

    MaterialTheme(
        colorScheme = colorScheme,
        typography = Typography(
            displayLarge = MaterialTheme.typography.displayLarge.copy(
                fontWeight = FontWeight.Bold
            ),
            headlineLarge = MaterialTheme.typography.headlineLarge.copy(
                fontWeight = FontWeight.Bold
            ),
            titleLarge = MaterialTheme.typography.titleLarge.copy(
                fontWeight = FontWeight.SemiBold
            ),
            bodyLarge = MaterialTheme.typography.bodyLarge.copy(
                lineHeight = 24.sp
            ),
            bodyMedium = MaterialTheme.typography.bodyMedium.copy(
                lineHeight = 20.sp
            )
        )
    ) {
        Surface(
            modifier = Modifier.fillMaxSize(),
            color = MaterialTheme.colorScheme.background
        ) {
            content()
        }
    }
}

@OptIn(ExperimentalMaterial3Api::class)
@Composable
fun NotesApp() {
    val context = LocalContext.current
    val dataStoreManager = remember { DataStoreManager(context) }
    val scope = rememberCoroutineScope()
    val lifecycleOwner = androidx.lifecycle.compose.LocalLifecycleOwner.current

    var categories by remember { mutableStateOf(listOf<Category>()) }
    var notes by remember { mutableStateOf(listOf<Note>()) }
    var selectedCategory by remember { mutableStateOf<Category?>(null) }
    var currentScreen by remember { mutableStateOf("main") }
    var drawerState by remember { mutableStateOf(false) }
    var showCategoryDialog by remember { mutableStateOf(false) }
    var showNoteDialog by remember { mutableStateOf(false) }
    var editingNote by remember { mutableStateOf<Note?>(null) }
    var searchQuery by remember { mutableStateOf("") }
    var showSearch by remember { mutableStateOf(false) }
    var showFullScreenNote by remember { mutableStateOf(false) }
    var fullScreenNote by remember { mutableStateOf<Note?>(null) }
    var isDarkTheme by remember { mutableStateOf(true) }

    // STATE UNTUK AI
    var showAIDrawer by remember { mutableStateOf(false) }
    var aiSelectedCategory by remember { mutableStateOf<Category?>(null) }
    var currentChatId by remember { mutableStateOf<String?>(null) }

    var isDataLoaded by remember { mutableStateOf(false) }

    // Load chat histories dari DataStore
    val chatHistories by dataStoreManager.chatHistoryFlow.collectAsState(initial = emptyList())

    fun sortCategories(categories: List<Category>): List<Category> {
        return categories
            .filter { !it.isDeleted }
            .sortedWith(
                compareByDescending<Category> { it.isPinned }
                    .thenByDescending { it.timestamp }
            )
    }

    LaunchedEffect(Unit) {
        dataStoreManager.themeFlow.collect { theme ->
            isDarkTheme = theme == "dark"
            AppColors.setTheme(isDarkTheme)
        }
    }

    LaunchedEffect(Unit) {
        dataStoreManager.categoriesFlow.collect { loadedCategories ->
            if (!isDataLoaded) {
                android.util.Log.d("NotesApp", "Loading ${loadedCategories.size} categories")
                categories = loadedCategories
            }
        }
    }

    LaunchedEffect(Unit) {
        dataStoreManager.notesFlow.collect { loadedNotes ->
            if (!isDataLoaded) {
                android.util.Log.d("NotesApp", "Loading ${loadedNotes.size} notes")
                notes = loadedNotes
                isDataLoaded = true
            }
        }
    }

    LaunchedEffect(categories) {
        if (isDataLoaded && categories.isNotEmpty()) {
            android.util.Log.d("NotesApp", "Saving ${categories.size} categories")
            scope.launch {
                dataStoreManager.saveCategories(categories)
            }
        }
    }

    LaunchedEffect(notes) {
        if (isDataLoaded && notes.isNotEmpty()) {
            android.util.Log.d("NotesApp", "Saving ${notes.size} notes")
            scope.launch {
                dataStoreManager.saveNotes(notes)
            }
        }
    }

    DisposableEffect(lifecycleOwner) {
        val observer = androidx.lifecycle.LifecycleEventObserver { _, event ->
            if (event == androidx.lifecycle.Lifecycle.Event.ON_PAUSE ||
                event == androidx.lifecycle.Lifecycle.Event.ON_STOP) {
                android.util.Log.d("NotesApp", "Lifecycle ${event.name}: Saving data")
                scope.launch {
                    if (categories.isNotEmpty()) {
                        dataStoreManager.saveCategories(categories)
                    }
                    if (notes.isNotEmpty()) {
                        dataStoreManager.saveNotes(notes)
                    }
                }
            }
        }
        lifecycleOwner.lifecycle.addObserver(observer)
        onDispose {
            lifecycleOwner.lifecycle.removeObserver(observer)
        }
    }

    Box(modifier = Modifier.fillMaxSize()) {
        // LAYER 1: Main Content (Scaffold)
        Scaffold(
            containerColor = AppColors.Background,
            topBar = {
                if (!showFullScreenNote && currentScreen != "ai") {
                    ModernTopBar(
                        title = when(currentScreen) {
                            "main" -> if (selectedCategory != null) selectedCategory!!.name else "NotesAI"
                            "starred" -> "Berbintang"
                            "archive" -> "Arsip"
                            "trash" -> "Sampah"
                            else -> "NotesAI"
                        },
                        showBackButton = (selectedCategory != null && currentScreen == "main"),
                        onBackClick = {
                            selectedCategory = null
                        },
                        onMenuClick = { drawerState = !drawerState },
                        onSearchClick = {
                            showSearch = !showSearch
                            if (!showSearch) searchQuery = "" // Reset search saat close
                        },
                        searchQuery = searchQuery,
                        onSearchQueryChange = { searchQuery = it },
                        showSearch = showSearch // AKTIFKAN UNTUK SEMUA SCREEN
                    )
                }
            },
            floatingActionButton = {
                AnimatedVisibility(
                    visible = currentScreen == "main" && !showFullScreenNote,
                    enter = scaleIn(
                        animationSpec = spring(
                            dampingRatio = Spring.DampingRatioMediumBouncy,
                            stiffness = Spring.StiffnessLow
                        )
                    ) + fadeIn(),
                    exit = scaleOut() + fadeOut()
                ) {
                    FloatingActionButton(
                        onClick = {
                            if (selectedCategory != null) {
                                editingNote = null
                                showNoteDialog = true
                            } else {
                                showCategoryDialog = true
                            }
                        },
                        containerColor = AppColors.Primary,
                        contentColor = Color.White,
                        elevation = FloatingActionButtonDefaults.elevation(
                            defaultElevation = 8.dp,
                            pressedElevation = 12.dp
                        ),
                        modifier = Modifier.size(64.dp)
                    ) {
                        Icon(
                            Icons.Default.Add,
                            contentDescription = if (selectedCategory != null) "Tambah Catatan" else "Tambah Kategori",
                            modifier = Modifier.size(28.dp)
                        )
                    }
                }
            },
            bottomBar = {
                if (!showFullScreenNote) {
                    ModernBottomBar(
                        currentScreen = currentScreen,
                        onHomeClick = {
                            currentScreen = "main"
                            selectedCategory = null
                        },
                        onAIClick = { currentScreen = "ai" }
                    )
                }
            }
        ) { padding ->
            Box(modifier = Modifier.fillMaxSize()) {
                if (showFullScreenNote && fullScreenNote != null) {
                    EditableFullScreenNoteView(
                        note = fullScreenNote!!,
                        onBack = {
                            showFullScreenNote = false
                            fullScreenNote = null
                        },
                        onSave = { title, content ->
                            notes = notes.map {
                                if (it.id == fullScreenNote!!.id) it.copy(
                                    title = title,
                                    content = content,
                                    timestamp = System.currentTimeMillis()
                                )
                                else it
                            }
                            fullScreenNote = fullScreenNote!!.copy(title = title, content = content)
                        },
                        onArchive = {
                            notes = notes.map {
                                if (it.id == fullScreenNote!!.id) it.copy(isArchived = true)
                                else it
                            }
                            showFullScreenNote = false
                            fullScreenNote = null
                        },
                        onDelete = {
                            notes = notes.map {
                                if (it.id == fullScreenNote!!.id) it.copy(isDeleted = true)
                                else it
                            }
                            showFullScreenNote = false
                            fullScreenNote = null
                        },
                        onPinToggle = {
                            notes = notes.map {
                                if (it.id == fullScreenNote!!.id) it.copy(isPinned = !it.isPinned)
                                else it
                            }
                            fullScreenNote = fullScreenNote!!.copy(isPinned = !fullScreenNote!!.isPinned)
                        }
                    )
                } else {
                    Column(
                        modifier = Modifier
                            .padding(padding)
                            .fillMaxSize()
                    ) {
                        when (currentScreen) {
                            "main" -> MainScreen(
                                categories = sortCategories(categories),
                                notes = notes,
                                selectedCategory = selectedCategory,
                                searchQuery = searchQuery,
                                onCategoryClick = { selectedCategory = it },
                                onNoteClick = { note ->
                                    fullScreenNote = note
                                    showFullScreenNote = true
                                },
                                onPinToggle = { note ->
                                    notes = notes.map {
                                        if (it.id == note.id) it.copy(isPinned = !it.isPinned)
                                        else it
                                    }
                                },
                                onCategoryDelete = { category ->
                                    categories = categories.map {
                                        if (it.id == category.id) it.copy(isDeleted = true)
                                        else it
                                    }
                                    notes = notes.map {
                                        if (it.categoryId == category.id) it.copy(isDeleted = true)
                                        else it
                                    }
                                    selectedCategory = null
                                },
                                onCategoryEdit = { category, newName, newGradientStart, newGradientEnd ->
                                    categories = categories.map {
                                        if (it.id == category.id) {
                                            it.copy(
                                                name = newName,
                                                gradientStart = newGradientStart,
                                                gradientEnd = newGradientEnd,
                                                timestamp = System.currentTimeMillis()
                                            )
                                        } else {
                                            it
                                        }
                                    }
                                },
                                onCategoryPin = { category ->
                                    categories = categories.map {
                                        if (it.id == category.id) it.copy(isPinned = !it.isPinned)
                                        else it
                                    }
                                },
                                onNoteEdit = { note ->
                                    editingNote = note
                                    showNoteDialog = true
                                },
                                onNoteDelete = { note ->
                                    notes = notes.map {
                                        if (it.id == note.id) it.copy(isDeleted = true)
                                        else it
                                    }
                                }
                            )

                            "trash" -> TrashScreen(
                                notes = notes.filter { it.isDeleted },
                                categories = categories,
                                searchQuery = searchQuery, // TAMBAHKAN INI
                                onRestoreNote = { note ->
                                    notes = notes.map {
                                        if (it.id == note.id) it.copy(isDeleted = false, isArchived = false)
                                        else it
                                    }
                                },
                                onDeleteNotePermanent = { note ->
                                    notes = notes.filter { it.id != note.id }
                                },
                                onRestoreCategory = { category ->
                                    categories = categories.map {
                                        if (it.id == category.id) it.copy(isDeleted = false)
                                        else it
                                    }
                                    notes = notes.map {
                                        if (it.categoryId == category.id) it.copy(isDeleted = false, isArchived = false)
                                        else it
                                    }
                                },
                                onDeleteCategoryPermanent = { category ->
                                    categories = categories.filter { it.id != category.id }
                                    notes = notes.filter { it.categoryId != category.id }
                                }
                            )

                            "starred" -> StarredNotesScreen(
                                notes = notes,
                                categories = categories.filter { !it.isDeleted },
                                searchQuery = searchQuery, // TAMBAHKAN INI
                                onNoteClick = { note ->
                                    fullScreenNote = note
                                    showFullScreenNote = true
                                },
                                onUnpin = { note ->
                                    notes = notes.map {
                                        if (it.id == note.id) it.copy(isPinned = false)
                                        else it
                                    }
                                }
                            )

                            "archive" -> ArchiveScreen(
                                notes = notes.filter { it.isArchived && !it.isDeleted },
                                categories = categories.filter { !it.isDeleted },
                                searchQuery = searchQuery, 
                                onRestore = { note ->
                                    notes = notes.map {
                                        if (it.id == note.id) it.copy(isArchived = false)
                                        else it
                                    }
                                },
                                onDelete = { note ->
                                    notes = notes.map {
                                        if (it.id == note.id) it.copy(isDeleted = true)
                                        else it
                                    }
                                }
                            )

                            "ai" -> AIHelperScreen(
                                categories = categories.filter { !it.isDeleted },
                                notes = notes.filter { !it.isDeleted },
                                onShowDrawer = { showAIDrawer = true }
                            )
                        }
                    }
                }

                if (showCategoryDialog) {
                    CategoryDialog(
                        onDismiss = { showCategoryDialog = false },
                        onSave = { name, gradientStart, gradientEnd ->
                            categories = categories + Category(
                                name = name,
                                gradientStart = gradientStart,
                                gradientEnd = gradientEnd
                            )
                            showCategoryDialog = false
                        }
                    )
                }

                if (showNoteDialog && selectedCategory != null) {
                    NoteDialog(
                        note = editingNote,
                        categoryId = selectedCategory!!.id,
                        onDismiss = {
                            showNoteDialog = false
                            editingNote = null
                        },
                        onSave = { title, description ->
                            if (editingNote != null) {
                                notes = notes.map {
                                    if (it.id == editingNote!!.id)
                                        it.copy(
                                            title = title,
                                            description = description,
                                            timestamp = System.currentTimeMillis()
                                        )
                                    else it
                                }
                            } else {
                                notes = notes + Note(
                                    categoryId = selectedCategory!!.id,
                                    title = title,
                                    description = description,
                                    content = ""
                                )
                            }
                            showNoteDialog = false
                            editingNote = null
                        },
                        onDelete = if (editingNote != null) {
                            {
                                notes = notes.map {
                                    if (it.id == editingNote!!.id) it.copy(isDeleted = true)
                                    else it
                                }
                                showNoteDialog = false
                                editingNote = null
                            }
                        } else null
                    )
                }
            }
        }

        // LAYER 2: Main Drawer (z-index 150)
        AnimatedVisibility(
            visible = drawerState,
            enter = fadeIn() + slideInHorizontally(initialOffsetX = { -it }),
            exit = fadeOut() + slideOutHorizontally(targetOffsetX = { -it }),
            modifier = Modifier.zIndex(150f)
        ) {
            DrawerMenu(
                currentScreen = currentScreen,
                isDarkTheme = isDarkTheme,
                onDismiss = { drawerState = false },
                onItemClick = { screen ->
                    currentScreen = screen
                    selectedCategory = null
                    drawerState = false
                    showSearch = false // TUTUP SEARCH
                    searchQuery = ""   // RESET SEARCH QUERY
                },
                onThemeToggle = {
                    isDarkTheme = !isDarkTheme
                    AppColors.setTheme(isDarkTheme)
                    scope.launch {
                        dataStoreManager.saveTheme(if (isDarkTheme) "dark" else "light")
                    }
                }
            )
        }

        // LAYER 3: AI History Drawer (z-index 200 - PALING ATAS)
        AnimatedVisibility(
            visible = showAIDrawer,
            enter = fadeIn() + slideInHorizontally(initialOffsetX = { -it }),
            exit = fadeOut() + slideOutHorizontally(targetOffsetX = { -it }),
            modifier = Modifier.zIndex(200f)
        ) {
            com.example.notesai.presentation.screens.ai.components.ChatHistoryDrawer(
                chatHistories = chatHistories, // GUNAKAN chatHistories dari collectAsState
                categories = categories.filter { !it.isDeleted },
                notes = notes.filter { !it.isDeleted },
                selectedCategory = aiSelectedCategory,
                onDismiss = { showAIDrawer = false },
                onHistoryClick = { history ->
                    // Load chat history
                    aiSelectedCategory = categories.find { it.id == history.categoryId }
                    currentChatId = history.id
                    showAIDrawer = false
                    // Anda perlu cara untuk pass data ini ke AIHelperScreen
                },
                onDeleteHistory = { historyId ->
                    scope.launch {
                        dataStoreManager.deleteChatHistory(historyId)
                    }
                },
                onCategorySelected = { category ->
                    aiSelectedCategory = category
                },
                onNewChat = {
                    aiSelectedCategory = null
                    currentChatId = null
                    showAIDrawer = false
                },
                onEditHistoryTitle = { historyId, newTitle ->
                    scope.launch {
                        dataStoreManager.updateChatHistoryTitle(historyId, newTitle)
                    }
                }
            )
        }
    }
}