package com.rugvedsarika.docscan.viewmodel

import android.app.Application
import android.net.Uri
import androidx.lifecycle.AndroidViewModel
import androidx.lifecycle.viewModelScope
import com.google.mlkit.vision.common.InputImage
import com.google.mlkit.vision.text.TextRecognition
import com.google.mlkit.vision.text.latin.TextRecognizerOptions
import com.rugvedsarika.docscan.data.AppDatabase
import com.rugvedsarika.docscan.data.entity.Document
import com.rugvedsarika.docscan.data.entity.Page
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.flatMapLatest
import kotlinx.coroutines.launch
import kotlinx.coroutines.tasks.await
import java.io.File

class MainViewModel(private val application: Application) : AndroidViewModel(application) {
    private val dao = AppDatabase.getDatabase(application).documentDao()
    private val recognizer = TextRecognition.getClient(TextRecognizerOptions.DEFAULT_OPTIONS)

    private val searchQuery = MutableStateFlow("")
    
    @OptIn(ExperimentalCoroutinesApi::class)
    val allDocuments: Flow<List<Document>> = searchQuery.flatMapLatest { query ->
        if (query.isEmpty()) dao.getAllDocuments()
        else dao.searchDocuments("%$query%")
    }

    fun setSearchQuery(query: String) {
        searchQuery.value = query
    }

    fun insertDocument(name: String, imagePaths: List<String>) {
        viewModelScope.launch {
            val thumbnail = if (imagePaths.isNotEmpty()) imagePaths[0] else null
            val document = Document(name = name, thumbnailPath = thumbnail)
            val docId = dao.insertDocument(document)
            
            imagePaths.forEachIndexed { index, path ->
                val extractedText = try {
                    val image = InputImage.fromFilePath(application, Uri.fromFile(File(path)))
                    val result = recognizer.process(image).await()
                    result.text
                } catch (e: Exception) {
                    null
                }
                
                dao.insertPage(Page(
                    documentId = docId, 
                    imagePath = path, 
                    orderIndex = index,
                    extractedText = extractedText
                ))
            }
        }
    }

    fun deleteDocument(document: Document) {
        deleteDocuments(listOf(document))
    }
    
    fun deleteDocuments(documents: List<Document>) {
        viewModelScope.launch {
            documents.forEach { doc ->
                // Delete image files
                val pages = dao.getPagesListForDocument(doc.id)
                pages.forEach { page ->
                    try {
                        File(page.imagePath).delete()
                    } catch (e: Exception) {
                        e.printStackTrace()
                    }
                }
                // Delete records (cascades to pages)
                dao.deleteDocument(doc)
            }
        }
    }
    
    fun toggleLock(document: Document) {
        viewModelScope.launch {
            dao.updateDocument(document.copy(isLocked = !document.isLocked))
        }
    }

    fun mergeDocuments(name: String, documents: List<Document>) {
        viewModelScope.launch {
            val thumbnail = if (documents.isNotEmpty()) documents[0].thumbnailPath else null
            val newDocId = dao.insertDocument(Document(name = name, thumbnailPath = thumbnail))
            
            var orderIndex = 0
            documents.forEach { doc ->
                val pages = dao.getPagesListForDocument(doc.id)
                pages.forEach { page ->
                    dao.insertPage(page.copy(id = 0, documentId = newDocId, orderIndex = orderIndex++))
                }
            }
        }
    }

    fun getPagesForDocument(documentId: Long): Flow<List<Page>> {
        return dao.getPagesForDocument(documentId)
    }

    suspend fun getPagesListForDocument(documentId: Long): List<Page> {
        return dao.getPagesListForDocument(documentId)
    }

    suspend fun getFullTextForDocuments(documents: List<Document>): String {
        val sb = StringBuilder()
        for (doc in documents) {
            val pages = dao.getPagesListForDocument(doc.id)
            sb.append("--- ${doc.name} ---\n")
            for (page in pages) {
                sb.append(page.extractedText ?: "").append("\n")
            }
            sb.append("\n")
        }
        return sb.toString()
    }

    suspend fun generatePdfFile(documents: List<Document>, outputFile: File) {
        val pdfDocument = android.graphics.pdf.PdfDocument()
        var pageNumber = 1
        
        for (doc in documents) {
            val pages = dao.getPagesListForDocument(doc.id)
            for (page in pages) {
                val bitmap = android.graphics.BitmapFactory.decodeFile(page.imagePath)
                if (bitmap != null) {
                    val pageInfo = android.graphics.pdf.PdfDocument.PageInfo.Builder(bitmap.width, bitmap.height, pageNumber++).create()
                    val pdfPage = pdfDocument.startPage(pageInfo)
                    val canvas = pdfPage.canvas
                    canvas.drawBitmap(bitmap, 0f, 0f, null)
                    pdfDocument.finishPage(pdfPage)
                    bitmap.recycle()
                }
            }
        }
        
        outputFile.outputStream().use { 
            pdfDocument.writeTo(it)
        }
        pdfDocument.close()
    }
}
