1 /*
<lambda>null2  * Copyright 2024 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package androidx.pdf.viewer.fragment
18 
19 import android.net.Uri
20 import androidx.annotation.RestrictTo
21 import androidx.core.os.OperationCanceledException
22 import androidx.lifecycle.SavedStateHandle
23 import androidx.lifecycle.ViewModel
24 import androidx.lifecycle.ViewModelProvider
25 import androidx.lifecycle.ViewModelProvider.AndroidViewModelFactory.Companion.APPLICATION_KEY
26 import androidx.lifecycle.createSavedStateHandle
27 import androidx.lifecycle.viewModelScope
28 import androidx.lifecycle.viewmodel.CreationExtras
29 import androidx.pdf.PdfDocument
30 import androidx.pdf.PdfLoader
31 import androidx.pdf.SandboxedPdfLoader
32 import androidx.pdf.exceptions.PdfPasswordException
33 import androidx.pdf.search.SearchRepository
34 import androidx.pdf.search.model.NoQuery
35 import androidx.pdf.search.model.QueryResults
36 import androidx.pdf.search.model.SearchResultState
37 import androidx.pdf.viewer.fragment.model.HighlightData
38 import androidx.pdf.viewer.fragment.model.PdfFragmentUiState
39 import androidx.pdf.viewer.fragment.model.SearchViewUiState
40 import androidx.pdf.viewer.fragment.util.fetchCounterData
41 import androidx.pdf.viewer.fragment.util.getCenter
42 import androidx.pdf.viewer.fragment.util.toHighlightsData
43 import java.util.concurrent.Executors
44 import kotlinx.coroutines.Job
45 import kotlinx.coroutines.SupervisorJob
46 import kotlinx.coroutines.asCoroutineDispatcher
47 import kotlinx.coroutines.flow.MutableStateFlow
48 import kotlinx.coroutines.flow.StateFlow
49 import kotlinx.coroutines.flow.asStateFlow
50 import kotlinx.coroutines.flow.update
51 import kotlinx.coroutines.launch
52 
53 /**
54  * A ViewModel class responsible for managing the loading and state of a PDF document.
55  *
56  * This ViewModel uses a [PdfLoader] to asynchronously open a PDF document from a given Uri. The
57  * loading result, which can be either a success with a [PdfDocument] or a failure with an
58  * exception, is exposed through the `pdfDocumentStateFlow`.
59  *
60  * The `loadDocument` function initiates the loading process within the `viewModelScope`, ensuring
61  * that the operation is properly managed and not cancelled by configuration changes.
62  *
63  * @constructor Creates a new [PdfDocumentViewModel] instance.
64  * @property loader The [PdfLoader] used to open the PDF document.
65  */
66 @RestrictTo(RestrictTo.Scope.LIBRARY)
67 internal class PdfDocumentViewModel(
68     private val state: SavedStateHandle,
69     private val loader: PdfLoader
70 ) : ViewModel() {
71 
72     /** A Coroutine [Job] that manages the PDF loading task. */
73     private var documentLoadJob: Job? = null
74 
75     /**
76      * Parent [Job] for search query and result collectors. All children jobs will be cancelled upon
77      * disabling [PdfViewerFragment.isTextSearchActive].
78      */
79     private val searchCollector = SupervisorJob(viewModelScope.coroutineContext[Job])
80 
81     /**
82      * Parent [Job] for search operations triggered on [SearchRepository]. All children jobs will
83      * cancelled upon updating search query.
84      */
85     private var searchJob: Job = SupervisorJob(viewModelScope.coroutineContext[Job])
86 
87     private val _fragmentUiScreenState =
88         MutableStateFlow<PdfFragmentUiState>(PdfFragmentUiState.Loading)
89 
90     /**
91      * Represents the UI state of the fragment.
92      *
93      * Exposes the UI state as a StateFlow to enable reactive consumption and ensure that consumers
94      * always receive the latest state.
95      */
96     internal val fragmentUiScreenState: StateFlow<PdfFragmentUiState>
97         get() = _fragmentUiScreenState.asStateFlow()
98 
99     private val _searchViewUiState = MutableStateFlow<SearchViewUiState>(SearchViewUiState.Closed)
100 
101     /** Stream of UI states of the PdfSearchView. */
102     internal val searchViewUiState: StateFlow<SearchViewUiState>
103         get() = _searchViewUiState.asStateFlow()
104 
105     internal val immersiveModeFlow: StateFlow<Boolean>
106         get() = state.getStateFlow(IMMERSIVE_MODE_STATE_KEY, false)
107 
108     private val _highlightsFlow = MutableStateFlow<HighlightData>(EMPTY_HIGHLIGHTS)
109 
110     /** Stream of highlights to be added on PdfView. Also includes scroll to page data. */
111     internal val highlightsFlow: StateFlow<HighlightData>
112         get() = _highlightsFlow.asStateFlow()
113 
114     /**
115      * Indicates whether the user is entering their password for the first time or making a repeated
116      * attempt.
117      *
118      * This state is used to determine the appropriate error message to display in the password
119      * dialog.
120      */
121     private var passwordFailed = false
122 
123     /** DocumentUri as set in [state] */
124     val documentUriFromState: Uri?
125         get() = state[DOCUMENT_URI_KEY]
126 
127     /** isTextSearchActive as set in [state] */
128     val isTextSearchActiveFromState: Boolean
129         get() = state[TEXT_SEARCH_STATE_KEY] ?: false
130 
131     /** isImmersiveModeFromState as set in [state] */
132     val isImmersiveModeDesired: Boolean
133         get() = state[IMMERSIVE_MODE_STATE_KEY] ?: false
134 
135     /** Holds business logic for search feature. */
136     private lateinit var searchRepository: SearchRepository
137 
138     init {
139         /**
140          * Open PDF if documentUri was previously set in state. This will be required in events like
141          * process death
142          */
143         state.get<Uri>(DOCUMENT_URI_KEY)?.let { uri ->
144             documentLoadJob = viewModelScope.launch { openDocument(uri) }
145             /*
146             Trigger restoring search view once document is loaded.
147             This is required as [SearchRepository] depends on [PdfDocument] which is created in
148             [PdfFragmentUiState.DocumentLoaded] state.
149             */
150             documentLoadJob?.invokeOnCompletion { maybeRestoreSearchState() }
151             documentLoadJob?.invokeOnCompletion { maybeRestoreImmersiveModeState() }
152         }
153     }
154 
155     private fun maybeRestoreImmersiveModeState() {
156         setImmersiveModeDesired(enterImmersive = isImmersiveModeDesired)
157     }
158 
159     private fun maybeRestoreSearchState() {
160         // Return early if search is disabled, as there's no result to restore.
161         if (!isTextSearchActiveFromState) return
162 
163         // Restore search session from last state saved
164         updateSearchState(isTextSearchActive = isTextSearchActiveFromState)
165         val query = state.get<String>(SEARCH_QUERY_KEY)
166         val pageNum = state.get<Int>(QUERY_RESULT_PAGE_NUM_KEY) ?: 0
167         val resultIndex = state.get<Int>(QUERY_RESULT_INDEX_KEY) ?: 0
168         query?.let {
169             viewModelScope.launch(searchJob) {
170                 searchRepository.produceSearchResults(
171                     query = query,
172                     currentVisiblePage = pageNum,
173                     resultIndex = resultIndex
174                 )
175             }
176         }
177     }
178 
179     /**
180      * Initiates the loading of a PDF document from the provided Uri.
181      *
182      * This function uses the provided [PdfLoader] to asynchronously open the PDF document. The
183      * loading result is then posted to the `pdfDocumentStateFlow` as a [Result] object, indicating
184      * either success with a [PdfDocument] or failure with an exception.
185      *
186      * The loading operation is executed within the `viewModelScope` to ensure that it continues
187      * even if a configuration change occurs.
188      *
189      * @param uri The Uri of the PDF document to load.
190      * @param password The optional password to use if the document is encrypted.
191      */
192     fun loadDocument(uri: Uri?, password: String?) {
193         uri?.let {
194             /*
195             Triggers the document loading process only under the following conditions:
196             1. **New Document URI:** The URI of the document to be loaded is different
197             `from the URI of the previously loaded document.
198             2. **Previous Load Failure or No Previous Load:** This is required when a
199              reload of document is required like document loading failed previous time or opened
200              using an incorrect password.
201              */
202             if (
203                 (uri != state[DOCUMENT_URI_KEY] ||
204                     fragmentUiScreenState.value !is PdfFragmentUiState.DocumentLoaded)
205             ) {
206                 state[DOCUMENT_URI_KEY] = uri
207                 // Ensure we don't schedule duplicate loading by canceling previous one.
208                 if (documentLoadJob?.isActive == true) documentLoadJob?.cancel()
209 
210                 // Loading a new document should not persist a search session from previous
211                 // document.
212                 updateSearchState(isTextSearchActive = false)
213                 setImmersiveModeDesired(enterImmersive = true)
214 
215                 documentLoadJob = viewModelScope.launch { openDocument(uri, password) }
216             }
217         }
218     }
219 
220     /**
221      * Called when the user toggles the search view's active state
222      * [PdfViewerFragment.isTextSearchActive].
223      *
224      * This function updates the search state in the [SavedStateHandle] and performs actions related
225      * to enabling/disabling the search view.
226      */
227     internal fun updateSearchState(isTextSearchActive: Boolean) {
228         /**
229          * [SearchRepository] is initialized only after a document is successfully loaded. If user
230          * triggers search before document is loaded, it will be a No-Op.
231          */
232         if (fragmentUiScreenState.value !is PdfFragmentUiState.DocumentLoaded) return
233 
234         state[TEXT_SEARCH_STATE_KEY] = isTextSearchActive
235 
236         if (isTextSearchActive) {
237             _searchViewUiState.update { SearchViewUiState.Init }
238             collectSearchResults()
239         } else {
240             searchJob.children.forEach { it.cancel() }
241             searchCollector.children.forEach { it.cancel() }
242             searchRepository.clearSearchResults()
243 
244             _searchViewUiState.update { SearchViewUiState.Closed }
245             _highlightsFlow.update { EMPTY_HIGHLIGHTS }
246             // Remove search params set in state on disabling search.
247             state.apply {
248                 remove<String>(SEARCH_QUERY_KEY)
249                 remove<Int>(QUERY_RESULT_PAGE_NUM_KEY)
250                 remove<Int>(QUERY_RESULT_INDEX_KEY)
251             }
252         }
253     }
254 
255     private fun collectSearchResults() {
256         viewModelScope.launch(searchCollector) {
257             searchRepository.queryResults.collect { queryResults ->
258                 handleQueryResults(queryResults)
259             }
260         }
261     }
262 
263     private fun handleQueryResults(queryResults: SearchResultState) {
264         when (queryResults) {
265             is NoQuery -> {
266                 _searchViewUiState.update { SearchViewUiState.Init }
267                 _highlightsFlow.update { EMPTY_HIGHLIGHTS }
268             }
269             is QueryResults.NoMatch -> {
270                 state[SEARCH_QUERY_KEY] = queryResults.query
271 
272                 _searchViewUiState.update {
273                     SearchViewUiState.Active(
274                         query = queryResults.query,
275                         currentMatch = 0,
276                         totalMatches = 0
277                     )
278                 }
279                 _highlightsFlow.update { EMPTY_HIGHLIGHTS }
280             }
281             is QueryResults.Matched -> {
282                 with(queryResults) {
283                     state[SEARCH_QUERY_KEY] = query
284                     state[QUERY_RESULT_PAGE_NUM_KEY] = queryResultsIndex.pageNum
285                     state[QUERY_RESULT_INDEX_KEY] = queryResultsIndex.resultBoundsIndex
286                 }
287 
288                 val (currentIndex, totalMatches) = queryResults.fetchCounterData()
289                 _searchViewUiState.update {
290                     SearchViewUiState.Active(
291                         query = queryResults.query,
292                         // The UI displays the search result counter starting from 1,
293                         // so we add 1 to the current index.
294                         currentMatch = if (totalMatches > 0) currentIndex + 1 else 0,
295                         totalMatches = totalMatches
296                     )
297                 }
298                 _highlightsFlow.update {
299                     HighlightData(
300                         currentIndex = currentIndex,
301                         highlightBounds = queryResults.toHighlightsData()
302                     )
303                 }
304             }
305         }
306     }
307 
308     /**
309      * Handles user interaction related to enabling the immersive mode.
310      *
311      * This function ensures that the immersive mode is properly applied and ready for user input
312      * when triggered.
313      */
314     fun setImmersiveModeDesired(enterImmersive: Boolean) {
315         /**
316          * Immersive mode state should be updated only after document is loaded. else it will be a
317          * No-Op.
318          */
319         if (fragmentUiScreenState.value !is PdfFragmentUiState.DocumentLoaded) return
320         state[IMMERSIVE_MODE_STATE_KEY] = enterImmersive
321     }
322 
323     /**
324      * Toggles the immersive mode state.
325      *
326      * This function ensures that the immersive mode is properly applied and ready for user input
327      * when triggered.
328      */
329     fun toggleImmersiveModeState() {
330         /**
331          * Immersive mode state should be updated only after document is loaded. else it will be a
332          * No-Op.
333          */
334         if (fragmentUiScreenState.value !is PdfFragmentUiState.DocumentLoaded) return
335         state[IMMERSIVE_MODE_STATE_KEY] = !isImmersiveModeDesired
336     }
337 
338     private suspend fun openDocument(uri: Uri, password: String? = null) {
339         /**
340          * PdfDocument, if ever created, will be stored in DocumentLoaded state. This state could be
341          * transitioned to other only if a new uri is submitted.
342          */
343         releaseDocument()
344 
345         /** Move to [PdfFragmentUiState.Loading] state before we begin load operation. */
346         _fragmentUiScreenState.update { PdfFragmentUiState.Loading }
347 
348         try {
349 
350             // Try opening pdf with provided params
351             val document = loader.openDocument(uri, password)
352 
353             searchRepository = SearchRepository(document)
354 
355             /** Successful load, move to [PdfFragmentUiState.DocumentLoaded] state. */
356             _fragmentUiScreenState.update { PdfFragmentUiState.DocumentLoaded(document) }
357             setImmersiveModeDesired(enterImmersive = false)
358 
359             /** Resets the [passwordFailed] state after a document is successfully loaded. */
360             passwordFailed = false
361         } catch (passwordException: PdfPasswordException) {
362             /** Move to [PdfFragmentUiState.PasswordRequested] for password protected pdf. */
363             _fragmentUiScreenState.update { PdfFragmentUiState.PasswordRequested(passwordFailed) }
364 
365             /** Enable [passwordFailed] for subsequent password attempts. */
366             passwordFailed = true
367         } catch (exception: Exception) {
368             /** Generic exception handling, move to [PdfFragmentUiState.DocumentError] state. */
369             _fragmentUiScreenState.update { PdfFragmentUiState.DocumentError(exception) }
370 
371             /** Resets the [passwordFailed] state after a document failed to load. */
372             passwordFailed = false
373         }
374     }
375 
376     /** Intent triggered when user submits a search query. */
377     fun searchDocument(query: String, visiblePageRange: IntRange) {
378         /**
379          * Cannot start searching document before it's loaded, i.e. fragment is moved to
380          * [PdfFragmentUiState.DocumentLoaded] state.
381          */
382         if (fragmentUiScreenState.value !is PdfFragmentUiState.DocumentLoaded) return
383 
384         val queryResults = searchRepository.queryResults.value
385         // Return early if the query is unchanged from the previous search to avoid redundant
386         // operations.
387         if (queryResults is QueryResults && queryResults.query == query) return
388 
389         // Cancel any on-going search operation(s) as the results will not be valid anymore.
390         searchJob.children.forEach { it.cancel() }
391 
392         viewModelScope.launch(searchJob) {
393             searchRepository.produceSearchResults(
394                 query = query,
395                 currentVisiblePage = visiblePageRange.getCenter()
396             )
397         }
398     }
399 
400     /** Intent triggered when user clicks prev button. */
401     fun findPreviousMatch() {
402         viewModelScope.launch(searchJob) { searchRepository.producePreviousResult() }
403     }
404 
405     /** Intent triggered when user clicks next button. */
406     fun findNextMatch() {
407         viewModelScope.launch(searchJob) { searchRepository.produceNextResult() }
408     }
409 
410     private fun IntRange.getCenterPage(): Int {
411         val size = endInclusive - first + 1
412         return first + size / 2
413     }
414 
415     fun passwordDialogCancelled() {
416         /** Resets the [passwordFailed] state after a password dialog is cancelled. */
417         passwordFailed = false
418         _fragmentUiScreenState.update {
419             PdfFragmentUiState.DocumentError(
420                 OperationCanceledException("Password cancelled. Cannot open PDF.")
421             )
422         }
423     }
424 
425     /**
426      * Closes the currently loaded PDF document, if one exists. This is important to release
427      * resources and prevent leaks.
428      */
429     private fun releaseDocument() {
430         (_fragmentUiScreenState.value as? PdfFragmentUiState.DocumentLoaded)?.pdfDocument?.close()
431     }
432 
433     override fun onCleared() {
434         super.onCleared()
435         releaseDocument()
436     }
437 
438     @Suppress("UNCHECKED_CAST")
439     companion object {
440 
441         private const val DOCUMENT_URI_KEY = "documentUri"
442         private const val TEXT_SEARCH_STATE_KEY = "textSearchState"
443         private const val IMMERSIVE_MODE_STATE_KEY = "immersiveModeState"
444         private const val SEARCH_QUERY_KEY = "searchQuery"
445         private const val QUERY_RESULT_INDEX_KEY = "queryResultIndex"
446         private const val QUERY_RESULT_PAGE_NUM_KEY = "queryResultPageNum"
447         private val EMPTY_HIGHLIGHTS = HighlightData(currentIndex = -1, highlightBounds = listOf())
448 
449         val Factory: ViewModelProvider.Factory =
450             object : ViewModelProvider.Factory {
451                 override fun <T : ViewModel> create(
452                     modelClass: Class<T>,
453                     extras: CreationExtras
454                 ): T {
455                     // Get the Application object from extras
456                     val application = checkNotNull(extras[APPLICATION_KEY])
457                     // Create a SavedStateHandle for this ViewModel from extras
458                     val savedStateHandle = extras.createSavedStateHandle()
459 
460                     val dispatcher = Executors.newSingleThreadExecutor().asCoroutineDispatcher()
461                     return (PdfDocumentViewModel(
462                         savedStateHandle,
463                         SandboxedPdfLoader(application, dispatcher)
464                     ))
465                         as T
466                 }
467             }
468     }
469 }
470