Trace utility network

/* Copyright 2025 Esri
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

package com.esri.arcgismaps.sample.traceutilitynetwork

import android.os.Bundle
import androidx.activity.ComponentActivity
import androidx.activity.compose.setContent
import androidx.activity.enableEdgeToEdge
import androidx.compose.material3.MaterialTheme
import androidx.compose.material3.Surface
import androidx.compose.runtime.Composable
import androidx.compose.runtime.DisposableEffect
import com.arcgismaps.ApiKey
import com.arcgismaps.ArcGISEnvironment
import com.esri.arcgismaps.sample.sampleslib.theme.SampleAppTheme
import com.esri.arcgismaps.sample.traceutilitynetwork.screens.TraceUtilityNetworkScreen

class MainActivity : ComponentActivity() {

    override fun onCreate(savedInstanceState: Bundle?) {
        super.onCreate(savedInstanceState)
        // authentication with an API key or named user is
        // required to access basemaps and other location services
        ArcGISEnvironment.apiKey = ApiKey.create(BuildConfig.ACCESS_TOKEN)

        enableEdgeToEdge()
        setContent {
            SampleAppTheme {
                TraceUtilityNetworkApp()
            }
        }
    }

    @Composable
    private fun TraceUtilityNetworkApp() {
        Surface(color = MaterialTheme.colorScheme.background) {
            TraceUtilityNetworkScreen(
                sampleName = getString(R.string.trace_utility_network_app_name)
            )

            // remove credentials on screen dispose
            DisposableEffect(Unit) {
                onDispose {
                    ArcGISEnvironment.authenticationManager.arcGISCredentialStore.removeAll()
                }
            }
        }
    }
}

/* Copyright 2025 Esri
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

package com.esri.arcgismaps.sample.traceutilitynetwork.components

import android.app.Application
import android.widget.Toast
import androidx.compose.ui.unit.dp
import androidx.lifecycle.AndroidViewModel
import androidx.lifecycle.viewModelScope
import com.arcgismaps.ArcGISEnvironment
import com.arcgismaps.Color
import com.arcgismaps.data.ArcGISFeature
import com.arcgismaps.data.QueryParameters
import com.arcgismaps.geometry.Geometry
import com.arcgismaps.geometry.GeometryEngine
import com.arcgismaps.geometry.Point
import com.arcgismaps.geometry.Polyline
import com.arcgismaps.httpcore.authentication.ArcGISAuthenticationChallengeHandler
import com.arcgismaps.httpcore.authentication.ArcGISAuthenticationChallengeResponse
import com.arcgismaps.httpcore.authentication.TokenCredential
import com.arcgismaps.mapping.ArcGISMap
import com.arcgismaps.mapping.Viewpoint
import com.arcgismaps.mapping.layers.FeatureLayer
import com.arcgismaps.mapping.layers.SelectionMode
import com.arcgismaps.mapping.symbology.SimpleLineSymbol
import com.arcgismaps.mapping.symbology.SimpleLineSymbolStyle
import com.arcgismaps.mapping.symbology.SimpleMarkerSymbol
import com.arcgismaps.mapping.symbology.SimpleMarkerSymbolStyle
import com.arcgismaps.mapping.symbology.UniqueValue
import com.arcgismaps.mapping.symbology.UniqueValueRenderer
import com.arcgismaps.mapping.view.Graphic
import com.arcgismaps.mapping.view.GraphicsOverlay
import com.arcgismaps.mapping.view.IdentifyLayerResult
import com.arcgismaps.mapping.view.ScreenCoordinate
import com.arcgismaps.toolkit.geoviewcompose.MapViewProxy
import com.arcgismaps.utilitynetworks.UtilityElement
import com.arcgismaps.utilitynetworks.UtilityElementTraceResult
import com.arcgismaps.utilitynetworks.UtilityNetwork
import com.arcgismaps.utilitynetworks.UtilityNetworkSource
import com.arcgismaps.utilitynetworks.UtilityNetworkSourceType
import com.arcgismaps.utilitynetworks.UtilityTerminal
import com.arcgismaps.utilitynetworks.UtilityTier
import com.arcgismaps.utilitynetworks.UtilityTraceParameters
import com.arcgismaps.utilitynetworks.UtilityTraceType
import com.esri.arcgismaps.sample.sampleslib.components.MessageDialogViewModel
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.asStateFlow
import kotlinx.coroutines.launch
import kotlinx.coroutines.runBlocking
import kotlin.math.roundToLong

class TraceUtilityNetworkViewModel(application: Application) : AndroidViewModel(application) {

    // The textual hint shown to the user
    private val _hint = MutableStateFlow<String?>(null)
    val hint = _hint.asStateFlow()

    // Is trace utility network enabled
    private val _canTrace = MutableStateFlow(false)
    val canTrace = _canTrace.asStateFlow()

    // The trace state used for the sample
    private val _traceState = MutableStateFlow(TraceState.ADD_STARTING_POINT)
    val traceState = _traceState.asStateFlow()

    // Currently selected utility trace type
    private val _selectedTraceType = MutableStateFlow<UtilityTraceType>(UtilityTraceType.Connected)
    val selectedTraceType = _selectedTraceType.asStateFlow()

    // Currently selected point type (start/barrier)
    private val _selectedPointType = MutableStateFlow(PointType.Start)
    val selectedPointType = _selectedPointType.asStateFlow()

    // Terminal configuration options (high/low)
    private val _terminalConfigurationOptions = MutableStateFlow<List<UtilityTerminal>>(listOf())
    val terminalConfigurationOptions = _terminalConfigurationOptions.asStateFlow()

    // Currently selected terminal configuration
    private var _selectedTerminalConfigurationIndex = MutableStateFlow<Int?>(null)

    // ArcGISMap holding the UtilityNetwork and operational layers
    val arcGISMap =
        ArcGISMap("https://sampleserver7.arcgisonline.com/portal/home/item.html?id=be0e4637620a453584118107931f718b")

    // Used to handle map view animations
    val mapViewProxy = MapViewProxy()

    // The utility network used for tracing
    private var utilityNetwork: UtilityNetwork? = null

    // The medium voltage tier used for the electric distribution domain network
    private var mediumVoltageTier: UtilityTier? = null

    // Create lists for starting locations and barriers
    private val utilityElementStartingLocations: MutableList<UtilityElement> = mutableListOf()
    private val utilityElementBarriers: MutableList<UtilityElement> = mutableListOf()

    // Graphics overlay for the starting locations and barrier graphics
    val graphicsOverlay = GraphicsOverlay()

    // Create symbols for the starting point and barriers
    private val startingPointSymbol = SimpleMarkerSymbol(
        style = SimpleMarkerSymbolStyle.Cross,
        color = Color.green,
        size = 25f
    )
    private val barrierPointSymbol = SimpleMarkerSymbol(
        style = SimpleMarkerSymbolStyle.X,
        color = Color.red,
        size = 25f
    )

    // Add custom unique renderer values for the electrical distribution layer
    private val electricalDistributionUniqueValueRenderer = UniqueValueRenderer(
        fieldNames = listOf("ASSETGROUP"),
        uniqueValues = listOf(
            UniqueValue(
                description = "Low voltage",
                label = "",
                symbol = SimpleLineSymbol(
                    style = SimpleLineSymbolStyle.Dash,
                    color = Color.green,
                    width = 3f
                ),
                values = listOf(3)
            ),
            UniqueValue(
                description = "Medium voltage",
                label = "",
                symbol = SimpleLineSymbol(
                    style = SimpleLineSymbolStyle.Solid,
                    color = Color.green,
                    width = 3f
                ),
                values = listOf(5)
            )
        )
    )

    /**
     * Returns a [ArcGISAuthenticationChallengeHandler] to access the utility network URL.
     */
    private fun getAuthenticationChallengeHandler(): ArcGISAuthenticationChallengeHandler {
        return ArcGISAuthenticationChallengeHandler { challenge ->
            val result: Result<TokenCredential> = runBlocking {
                TokenCredential.create(challenge.requestUrl, "viewer01", "I68VGU^nMurF", 0)
            }
            if (result.getOrNull() != null) {
                val credential = result.getOrNull()
                return@ArcGISAuthenticationChallengeHandler ArcGISAuthenticationChallengeResponse.ContinueWithCredential(
                    credential!!
                )
            } else {
                val ex = result.exceptionOrNull()
                return@ArcGISAuthenticationChallengeHandler ArcGISAuthenticationChallengeResponse.ContinueAndFailWithError(
                    ex!!
                )
            }
        }
    }

    /**
     * Initializes view model by adding credentials, loading map and utility network,
     * and electrical device and distribution feature layers.
     */
    suspend fun initializeTraceViewModel() {

        ArcGISEnvironment.authenticationManager.arcGISAuthenticationChallengeHandler =
            getAuthenticationChallengeHandler()

        // Load the map
        arcGISMap.load().onSuccess {

            // The utility network used for tracing
            utilityNetwork = arcGISMap.utilityNetworks.first()
            utilityNetwork?.let { utilityNetwork ->

                // Load the utility network
                utilityNetwork.load().onSuccess {

                    // Get the service geodatabase from the utility network
                    val serviceGeodatabase = utilityNetwork.serviceGeodatabase

                    // Use the ElectricDistribution domain network
                    val electricDistribution = utilityNetwork.definition?.getDomainNetwork("ElectricDistribution")

                    // Use the Medium Voltage Tier
                    mediumVoltageTier = electricDistribution?.getTier("Medium Voltage Radial")

                    (serviceGeodatabase?.getTable(3)?.layer as FeatureLayer).apply {
                        // Customize rendering for the layer
                        renderer = electricalDistributionUniqueValueRenderer
                    }

                    // Update hint values to reflect trace stage changes
                    viewModelScope.launch {
                        _traceState.collect { updateHint(it) }
                    }
                }
            }
        }
    }


    /**
     * Performs an identify operation to obtain the [ArcGISFeature] nearest to the
     * tapped [screenCoordinate]. The selected feature is then used to [identifyUtilityElement].
     */
    fun identifyNearestArcGISFeature(
        mapPoint: Point,
        screenCoordinate: ScreenCoordinate
    ) {
        viewModelScope.launch {
            // Identify the feature on the tapped location
            val identifyResults: List<IdentifyLayerResult> =
                mapViewProxy.identifyLayers(
                    screenCoordinate = screenCoordinate,
                    tolerance = 4.dp,
                    returnPopupsOnly = false,
                    maximumResults = 1
                ).getOrElse {
                    return@launch messageDialogVM.showMessageDialog(
                        title = it.message.toString(),
                        description = it.cause.toString()
                    )
                }
            // If the identify returns a result, retrieve the geoelement as an ArcGISFeature
            identifyResults.firstOrNull()?.geoElements?.firstOrNull()?.let { identifiedFeature ->
                (identifiedFeature as? ArcGISFeature)?.let { arcGISFeature ->
                    // Identify the utility element associated with the selected feature
                    identifyUtilityElement(
                        identifiedFeature = arcGISFeature,
                        mapPoint = mapPoint
                    )
                }
            }
        }
    }

    /**
     * Uses the [mapPoint] to identify any utility elements in the utility network.
     * Based on the [UtilityNetworkSourceType] create an element for a junction or an edge.
     */
    private fun identifyUtilityElement(
        identifiedFeature: ArcGISFeature,
        mapPoint: Point
    ) {
        // Get the network source of the identified feature
        val utilityNetworkSource = utilityNetwork?.definition?.networkSources?.value?.firstOrNull {
            it.featureTable.tableName == identifiedFeature.featureTable?.tableName
        } ?: return handleError("Selected feature does not contain a Utility Network Source.")

        // Check if the network source is a junction or an edge
        when (utilityNetworkSource.sourceType) {
            UtilityNetworkSourceType.Junction -> {
                // Create a junction element with the identified feature
                createJunctionUtilityElement(
                    identifiedFeature = identifiedFeature,
                    utilityNetworkSource = utilityNetworkSource
                )
            }

            UtilityNetworkSourceType.Edge -> {
                // Create an edge element with the identified feature
                createEdgeUtilityElement(
                    identifiedFeature = identifiedFeature,
                    mapPoint = mapPoint
                )
            }
        }
    }

    /**
     * Create a [UtilityElement] of the [identifiedFeature].
     */
    private fun createJunctionUtilityElement(
        identifiedFeature: ArcGISFeature,
        utilityNetworkSource: UtilityNetworkSource
    ) {
        // Find the code matching the asset group name in the feature's attributes
        val assetGroupCode = identifiedFeature.attributes["assetgroup"] as Int
        // Find the network source's asset group with the matching code
        utilityNetworkSource.assetGroups.first { it.code == assetGroupCode }.assetTypes
            // Find the asset group type code matching the feature's asset type code
            .first { it.code == identifiedFeature.attributes["assettype"].toString().toInt() }.let { utilityAssetType ->
                // Get the list of terminals for the feature
                val terminals = utilityAssetType.terminalConfiguration?.terminals
                    ?: return handleError("Error retrieving terminal configuration")

                // If there is only one terminal, use it to create a utility element
                when (terminals.size) {
                    1 -> {
                        // Create a utility element
                        utilityNetwork?.createElementOrNull(
                            arcGISFeature = identifiedFeature, terminal = terminals.first()
                        )?.let { utilityElement ->
                            // Add the utility element to the map
                            addUtilityElementToMap(
                                identifiedFeature = identifiedFeature,
                                mapPoint = identifiedFeature.geometry as Point,
                                utilityElement = utilityElement
                            )
                        }
                    }
                    // If there is more than one terminal, prompt the user to select one
                    else -> {
                        // Reset the index, as the user would need to make a choice
                        _selectedTerminalConfigurationIndex.value = null
                        // Get a list of terminal names from the terminal configuration
                        val terminalConfiguration = utilityAssetType.terminalConfiguration ?: return
                        // Update the list of available terminal options
                        _terminalConfigurationOptions.value = terminalConfiguration.terminals
                        // Show the dialog to choose a terminal configuration
                        _traceState.value = TraceState.TERMINAL_CONFIGURATION_REQUIRED

                        viewModelScope.launch {
                            _selectedTerminalConfigurationIndex.collect { selectedIndex ->
                                if (selectedIndex != null) {
                                    // Create a utility element
                                    val element = utilityNetwork?.createElementOrNull(
                                        arcGISFeature = identifiedFeature, terminal = terminals[selectedIndex]
                                    ) ?: return@collect handleError(
                                        "Error creating utility element"
                                    )
                                    // Add the utility element graphic to the map
                                    addUtilityElementToMap(
                                        identifiedFeature = identifiedFeature,
                                        mapPoint = identifiedFeature.geometry as Point,
                                        utilityElement = element
                                    )
                                    // Dismiss the dialog to choose another point
                                    _traceState.value = TraceState.ADD_STARTING_POINT
                                }
                            }
                        }
                    }
                }
            }
    }

    /**
     * Create a [UtilityElement] of the [identifiedFeature].
     */
    private fun createEdgeUtilityElement(
        identifiedFeature: ArcGISFeature,
        mapPoint: Point
    ) {
        // Create a utility element with the identified feature
        val element = (utilityNetwork?.createElementOrNull(
            arcGISFeature = identifiedFeature, terminal = null
        ) ?: return handleError("Error creating element"))
        // Calculate the fraction along these the map point is located
        element.fractionAlongEdge = GeometryEngine.fractionAlong(
            line = GeometryEngine.createWithZ(
                geometry = identifiedFeature.geometry!!,
                z = null // Remove the z-coordinate value from the identified geometry
            ) as Polyline,
            point = mapPoint,
            tolerance = -1.0
        ).roundToThreeDecimals()

        // Add the utility element graphic to the map
        addUtilityElementToMap(
            identifiedFeature = identifiedFeature,
            mapPoint = mapPoint,
            utilityElement = element
        )
        // Update the hint text
        updateHint("Fraction along the edge: ${element.fractionAlongEdge}")
    }

    /**
     * Add [utilityElement] to either the starting locations or barriers list
     * and add a graphic representing it to the [graphicsOverlay].
     */
    private fun addUtilityElementToMap(
        identifiedFeature: ArcGISFeature,
        mapPoint: Point,
        utilityElement: UtilityElement
    ) {
        graphicsOverlay.graphics.add(
            Graphic(
                geometry = GeometryEngine.nearestCoordinate(
                    geometry = identifiedFeature.geometry!!,
                    point = mapPoint
                )?.coordinate
            ).apply {
                // Add the element to the appropriate list (starting locations or barriers),
                // and add the appropriate symbol to the graphic
                when (_selectedPointType.value) {
                    PointType.Start -> {
                        utilityElementStartingLocations.add(utilityElement)
                        symbol = startingPointSymbol
                        _canTrace.value = true
                    }

                    PointType.Barrier -> {
                        utilityElementBarriers.add(utilityElement)
                        symbol = barrierPointSymbol
                    }
                }
            })
    }

    /**
     * Uses the elements selected as starting locations and (optionally) barriers
     * to perform a connected trace, then selects all connected elements
     * found in the trace to highlight them.
     */
    fun traceUtilityNetwork() {
        // Check that the utility trace parameters are valid
        if (utilityElementStartingLocations.isEmpty()) {
            return handleError("No starting locations provided for trace.")
        }

        val traceType = _selectedTraceType.value

        // Create utility trace parameters for the given trace type
        val traceParameters = UtilityTraceParameters(
            traceType = traceType,
            startingLocations = utilityElementStartingLocations
        ).apply {
            // If any barriers have been created, add them to the parameters
            barriers.addAll(utilityElementBarriers)
            // Set the trace configuration using the tier from the utility domain network
            traceConfiguration = mediumVoltageTier?.getDefaultTraceConfiguration()
        }

        // Run the utility trace and get the results
        viewModelScope.launch {
            // Update the trace state
            _traceState.value = TraceState.RUNNING_TRACE_UTILITY_NETWORK
            // Perform the trace with the above parameters, and obtain the results list
            val traceResults = utilityNetwork?.trace(traceParameters)?.getOrElse {
                return@launch handleError(
                    title = "Error performing trace",
                    description = it.message.toString()
                )
            }

            // Get the utility trace result's first result as a utility element trace result
            (traceResults?.first() as? UtilityElementTraceResult)?.let { utilityElementTraceResult ->
                // Ensure the result is not empty
                if (utilityElementTraceResult.elements.isEmpty()) return@launch handleError("No elements found in the trace result")

                arcGISMap.operationalLayers.filterIsInstance<FeatureLayer>().forEach { featureLayer ->
                    // Clear previous selection
                    featureLayer.clearSelection()
                    val params = QueryParameters().apply {
                        returnGeometry = true // Used to calculate the viewpoint result
                    }
                    // Create query parameters to find features who's network source name matches the layer's feature table name
                    utilityElementTraceResult.elements.filter {
                        it.networkSource.name == featureLayer.featureTable?.tableName
                    }.forEach { utilityElement ->
                        params.objectIds.add(utilityElement.objectId)
                    }

                    // Check if any trace results were added from the above filter
                    if (params.objectIds.isNotEmpty()) {
                        // Select features that match the query
                        val featureQueryResult = featureLayer.selectFeatures(
                            parameters = params, mode = SelectionMode.New
                        ).getOrElse {
                            return@launch handleError(
                                title = it.message.toString(), description = it.cause.toString()
                            )
                        }

                        // Create list of all the feature result geometries
                        val resultGeometryList = mutableListOf<Geometry>()
                        featureQueryResult.iterator().forEach { feature ->
                            feature.geometry?.let {
                                resultGeometryList.add(it)
                            }
                        }

                        // Obtain the union geometry of all the feature geometries
                        GeometryEngine.unionOrNull(resultGeometryList)?.let { unionGeometry ->
                            // Set the map's viewpoint to the union result geometry
                            mapViewProxy.setViewpointAnimated(Viewpoint(boundingGeometry = unionGeometry))
                        }
                    } else {
                        Toast.makeText(
                            getApplication(), "Trace result found 0 elements", Toast.LENGTH_SHORT
                        ).show()
                    }
                }

                // Update the trace state
                _traceState.value = TraceState.TRACE_COMPLETED
            }
        }
    }

    /**
     * Resets the trace, removing graphics and clearing selections.
     */
    fun reset() {
        arcGISMap.operationalLayers.filterIsInstance<FeatureLayer>().forEach { it.clearSelection() }
        utilityElementBarriers.clear()
        utilityElementStartingLocations.clear()
        graphicsOverlay.graphics.clear()
        _traceState.value = TraceState.ADD_STARTING_POINT
        _canTrace.value = false
        _selectedTraceType.value = UtilityTraceType.Connected
        _selectedTerminalConfigurationIndex.value = null
        _selectedPointType.value = PointType.Start
        _terminalConfigurationOptions.value = listOf()
    }

    /**
     * Update the [utilityTraceType] selected by the user
     */
    fun updateTraceType(utilityTraceType: UtilityTraceType) {
        _selectedTraceType.value = utilityTraceType
        _traceState.value = TraceState.ADD_STARTING_POINT
    }

    /**
     * Switch from adding .start points to adding .barrier, or vice versa.
     */
    fun updatePointType(pointType: PointType) {
        _selectedPointType.value = pointType
        when (pointType) {
            PointType.Start -> {
                _traceState.value = TraceState.ADD_STARTING_POINT
            }

            PointType.Barrier -> {
                _traceState.value = TraceState.ADD_BARRIER_POINT
            }
        }
    }

    /**
     * Update the index used to select the [terminalConfigurationOptions]
     */
    fun updateTerminalConfigurationOption(index: Int) {
        _selectedTerminalConfigurationIndex.value = index
    }

    /**
     * Update the hint flow to display new [message].
     */
    private fun updateHint(message: String) {
        _hint.value = message
    }

    // Create a message dialog view model for handling error messages
    val messageDialogVM = MessageDialogViewModel()

    private fun handleError(title: String, description: String = "") {
        reset()
        _traceState.value = TraceState.TRACE_FAILED
        messageDialogVM.showMessageDialog(title, description)
    }

    private fun Double.roundToThreeDecimals(): Double {
        return (this * 1000.0).roundToLong() / 1000.0
    }
}

enum class PointType {
    Start,
    Barrier
}

object TraceState {
    const val ADD_STARTING_POINT = "Tap on map to add a stating location point(s)"
    const val ADD_BARRIER_POINT = "Tap on map to add a barrier point(s)"
    const val TERMINAL_CONFIGURATION_REQUIRED = "Select Terminal Configuration"
    const val RUNNING_TRACE_UTILITY_NETWORK = "Evaluating trace utility network"
    const val TRACE_COMPLETED = "Trace completed"
    const val TRACE_FAILED = "Fail to run trace"
}

/* Copyright 2025 Esri
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

@file:OptIn(ExperimentalMaterial3Api::class)

package com.esri.arcgismaps.sample.traceutilitynetwork.screens

import android.content.res.Configuration
import androidx.compose.foundation.background
import androidx.compose.foundation.basicMarquee
import androidx.compose.foundation.layout.Arrangement
import androidx.compose.foundation.layout.Box
import androidx.compose.foundation.layout.Column
import androidx.compose.foundation.layout.Row
import androidx.compose.foundation.layout.fillMaxSize
import androidx.compose.foundation.layout.fillMaxWidth
import androidx.compose.foundation.layout.padding
import androidx.compose.foundation.layout.wrapContentSize
import androidx.compose.foundation.rememberScrollState
import androidx.compose.foundation.verticalScroll
import androidx.compose.material.icons.Icons
import androidx.compose.material.icons.filled.Settings
import androidx.compose.material3.BasicAlertDialog
import androidx.compose.material3.Button
import androidx.compose.material3.DropdownMenuItem
import androidx.compose.material3.ExperimentalMaterial3Api
import androidx.compose.material3.ExposedDropdownMenuBox
import androidx.compose.material3.ExposedDropdownMenuDefaults
import androidx.compose.material3.FloatingActionButton
import androidx.compose.material3.HorizontalDivider
import androidx.compose.material3.Icon
import androidx.compose.material3.MaterialTheme
import androidx.compose.material3.MenuAnchorType
import androidx.compose.material3.OutlinedButton
import androidx.compose.material3.Scaffold
import androidx.compose.material3.SegmentedButton
import androidx.compose.material3.SegmentedButtonDefaults
import androidx.compose.material3.SingleChoiceSegmentedButtonRow
import androidx.compose.material3.Surface
import androidx.compose.material3.Text
import androidx.compose.material3.TextField
import androidx.compose.runtime.Composable
import androidx.compose.runtime.LaunchedEffect
import androidx.compose.runtime.getValue
import androidx.compose.runtime.mutableStateOf
import androidx.compose.runtime.remember
import androidx.compose.runtime.setValue
import androidx.compose.ui.Alignment
import androidx.compose.ui.Modifier
import androidx.compose.ui.tooling.preview.Preview
import androidx.compose.ui.unit.dp
import androidx.compose.ui.util.fastForEachIndexed
import androidx.compose.ui.window.DialogProperties
import androidx.lifecycle.compose.collectAsStateWithLifecycle
import androidx.lifecycle.viewmodel.compose.viewModel
import com.arcgismaps.Color
import com.arcgismaps.LoadStatus
import com.arcgismaps.mapping.view.SelectionProperties
import com.arcgismaps.toolkit.geoviewcompose.MapView
import com.arcgismaps.utilitynetworks.UtilityTerminal
import com.arcgismaps.utilitynetworks.UtilityTraceType
import com.esri.arcgismaps.sample.sampleslib.components.BottomSheet
import com.esri.arcgismaps.sample.sampleslib.components.LoadingDialog
import com.esri.arcgismaps.sample.sampleslib.components.MessageDialog
import com.esri.arcgismaps.sample.sampleslib.components.SampleTopAppBar
import com.esri.arcgismaps.sample.sampleslib.theme.SampleAppTheme
import com.esri.arcgismaps.sample.traceutilitynetwork.components.PointType
import com.esri.arcgismaps.sample.traceutilitynetwork.components.TraceState
import com.esri.arcgismaps.sample.traceutilitynetwork.components.TraceUtilityNetworkViewModel

/**
 * Main screen containing utility network map view and trace controls.
 */
@Composable
fun TraceUtilityNetworkScreen(sampleName: String) {
    // Create the view model used by the sample
    val mapViewModel: TraceUtilityNetworkViewModel = viewModel()

    // Load map, utility network, and adds layers.
    LaunchedEffect(Unit) {
        mapViewModel.arcGISMap.apply {
            // Check if the network is not loaded
            if (utilityNetworks.isEmpty() || loadStatus.value == LoadStatus.NotLoaded) {
                mapViewModel.initializeTraceViewModel()
            }
        }
    }

    // Observe relevant states
    val hintText by mapViewModel.hint
        .collectAsStateWithLifecycle(null)
    val selectedTraceType by mapViewModel.selectedTraceType
        .collectAsStateWithLifecycle()
    val selectedPointType by mapViewModel.selectedPointType
        .collectAsStateWithLifecycle(PointType.Start)
    val canPerformTrace by mapViewModel.canTrace
        .collectAsStateWithLifecycle(false)
    val traceState by mapViewModel.traceState
        .collectAsStateWithLifecycle(TraceState.ADD_STARTING_POINT)
    val terminalConfigurationOptions by mapViewModel.terminalConfigurationOptions
        .collectAsStateWithLifecycle(listOf())

    // Set up the bottom sheet controls
    var isBottomSheetVisible by remember { mutableStateOf(false) }

    Scaffold(
        topBar = { SampleTopAppBar(title = sampleName) },
        content = { padding ->
            Box(modifier = Modifier.padding(padding)) {
                MapView(
                    modifier = Modifier
                        .fillMaxSize(),
                    arcGISMap = mapViewModel.arcGISMap,
                    mapViewProxy = mapViewModel.mapViewProxy,
                    graphicsOverlays = listOf(mapViewModel.graphicsOverlay),
                    selectionProperties = SelectionProperties(color = Color.yellow),
                    onSingleTapConfirmed = { tapEvent ->
                        // Identify tapped location if current state is valid.
                        tapEvent.mapPoint?.let {
                            mapViewModel.identifyNearestArcGISFeature(
                                mapPoint = it,
                                screenCoordinate = tapEvent.screenCoordinate
                            )
                        }
                    },
                    onDown = {
                        isBottomSheetVisible = false
                    }
                )

                BottomSheet(isVisible = isBottomSheetVisible) {
                    TraceOptions(
                        hintText = hintText ?: "Trace Options",
                        isTraceButtonEnabled = canPerformTrace,
                        utilityTraceType = selectedTraceType,
                        pointType = selectedPointType,
                        traceState = traceState,
                        onTraceSelected = mapViewModel::traceUtilityNetwork,
                        onPointTypeChanged = mapViewModel::updatePointType,
                        onTraceTypeSelected = mapViewModel::updateTraceType,
                        onResetSelected = mapViewModel::reset
                    )
                }
            }

            // Displays dialog to select a terminal configuration when required
            if (traceState == TraceState.TERMINAL_CONFIGURATION_REQUIRED) {
                TerminalConfigurationDialog(
                    terminalConfigurationOptions = terminalConfigurationOptions,
                    onTerminalConfigurationSelected = mapViewModel::updateTerminalConfigurationOption
                )
            }

            // Displays a loading dialog when trace is running
            if (traceState == TraceState.RUNNING_TRACE_UTILITY_NETWORK) {
                RunningTraceDialog(
                    traceName = selectedTraceType.javaClass.simpleName.toString()
                )
            }

            // Displays a dialog when sample encounters an error
            mapViewModel.messageDialogVM.apply {
                if (dialogStatus) {
                    MessageDialog(
                        title = messageTitle,
                        description = messageDescription,
                        onDismissRequest = ::dismissDialog
                    )
                }
            }
        },
        floatingActionButton = {
            if (!isBottomSheetVisible) {
                FloatingActionButton(
                    modifier = Modifier.padding(bottom = 36.dp, end = 12.dp),
                    onClick = { isBottomSheetVisible = true }
                ) { Icon(Icons.Filled.Settings, contentDescription = "Show Trace Options") }
            }
        }
    )
}

/**
 * Displays dialog to select a terminal configuration when required
 */
@Composable
fun TerminalConfigurationDialog(
    terminalConfigurationOptions: List<UtilityTerminal>,
    onTerminalConfigurationSelected: (Int) -> Unit,
) {
    BasicAlertDialog(
        onDismissRequest = {},
        properties = DialogProperties()
    ) {
        Surface {
            Column(
                modifier = Modifier
                    .verticalScroll(rememberScrollState())
                    .padding(24.dp),
                horizontalAlignment = Alignment.CenterHorizontally,
                verticalArrangement = Arrangement.spacedBy(12.dp)
            ) {
                Text(
                    text = "Select Terminal",
                    style = MaterialTheme.typography.titleMedium
                )
                terminalConfigurationOptions.fastForEachIndexed { index, utilityTerminal ->
                    OutlinedButton(onClick = { onTerminalConfigurationSelected(index) }) {
                        Text(utilityTerminal.name)
                    }
                }
            }
        }
    }
}


/**
 * Trace options layout with options for the trace type, starting vs barrier locations,
 * reset and tracing buttons.
 */
@Composable
fun TraceOptions(
    isTraceButtonEnabled: Boolean,
    hintText: String,
    utilityTraceType: UtilityTraceType,
    pointType: PointType,
    traceState: String,
    onTraceTypeSelected: (UtilityTraceType) -> Unit,
    onPointTypeChanged: (PointType) -> Unit,
    onResetSelected: () -> Unit,
    onTraceSelected: () -> Unit
) {
    Column(
        modifier = Modifier
            .wrapContentSize()
            .background(MaterialTheme.colorScheme.background)
            .padding(12.dp)
            .verticalScroll(rememberScrollState()),
        verticalArrangement = Arrangement.spacedBy(12.dp),
        horizontalAlignment = Alignment.CenterHorizontally
    ) {
        // Displays contextual hints
        Text(text = hintText, style = MaterialTheme.typography.labelLarge)

        // Show a dropdown menu to pick a new trace type
        ExposedDropdownMenuBoxWithTraceTypes(
            selectedTraceType = utilityTraceType,
            onTraceTypeSelected = onTraceTypeSelected
        )

        // Display segmented button for starting point type or barrier point type
        SegmentedButtonTracePointTypes(
            currentPointType = pointType,
            onPointTypeChanged = onPointTypeChanged,
        )

        // Display a row with reset and trace controls
        Row(
            modifier = Modifier.fillMaxWidth(),
            horizontalArrangement = Arrangement.SpaceAround
        ) {
            OutlinedButton(
                onClick = { onResetSelected() },
                enabled = traceState != TraceState.RUNNING_TRACE_UTILITY_NETWORK
            ) {
                Text("Reset")
            }
            Button(
                onClick = { onTraceSelected() },
                enabled = isTraceButtonEnabled
            ) {
                Text("Trace")
            }
        }
    }
}

/**
 * A ExposedDropdownMenuBox with a list of supported [UtilityTraceType].
 */
@OptIn(ExperimentalMaterial3Api::class)
@Composable
fun ExposedDropdownMenuBoxWithTraceTypes(
    selectedTraceType: UtilityTraceType,
    onTraceTypeSelected: (UtilityTraceType) -> Unit
) {
    val traceOptions = listOf(
        UtilityTraceType.Downstream,
        UtilityTraceType.Upstream,
        UtilityTraceType.Subnetwork,
        UtilityTraceType.Connected
    )

    var selectedTraceName by remember { mutableStateOf("") }
    var expanded by remember { mutableStateOf(false) }

    LaunchedEffect(selectedTraceType) {
        selectedTraceName = selectedTraceType.javaClass.simpleName
    }

    ExposedDropdownMenuBox(
        expanded = expanded,
        onExpandedChange = { expanded = !expanded }
    ) {
        TextField(
            modifier = Modifier
                .menuAnchor(MenuAnchorType.PrimaryNotEditable)
                .fillMaxWidth(),
            label = { Text("Trace Type") },
            value = selectedTraceName,
            readOnly = true,
            onValueChange = { },
            trailingIcon = { ExposedDropdownMenuDefaults.TrailingIcon(expanded = expanded) }
        )
        ExposedDropdownMenu(
            expanded = expanded,
            onDismissRequest = { expanded = false }
        ) {
            traceOptions.forEachIndexed { index, traceType ->
                DropdownMenuItem(
                    text = { Text(traceType.javaClass.simpleName) },
                    onClick = {
                        expanded = false
                        onTraceTypeSelected(traceType)
                    }
                )
                if (index < traceOptions.lastIndex) {
                    HorizontalDivider()
                }
            }
        }
    }
}

/**
 * A SingleChoiceSegmentedButtonRow with a choice between
 * starting or barrier point type.
 */
@Composable
fun SegmentedButtonTracePointTypes(
    currentPointType: PointType,
    onPointTypeChanged: (PointType) -> Unit
) {
    var selectedIndex = when (currentPointType) {
        PointType.Start -> 0
        PointType.Barrier -> 1
    }

    LaunchedEffect(currentPointType) {
        selectedIndex = if (currentPointType == PointType.Start) 0 else 1
    }

    SingleChoiceSegmentedButtonRow(modifier = Modifier.fillMaxWidth()) {
        val options = listOf("Add starting location(s)", "Add barrier(s)")
        options.forEachIndexed { index, label ->
            SegmentedButton(
                shape = SegmentedButtonDefaults.itemShape(index = index, count = options.size),
                onClick = {
                    selectedIndex = index
                    // Switch between Start/Barrier
                    onPointTypeChanged(if (index == 0) PointType.Start else PointType.Barrier)
                },
                selected = (index == selectedIndex)
            ) {
                Text(
                    modifier = Modifier.basicMarquee(),
                    text = label,
                    maxLines = 1,
                    style = MaterialTheme.typography.labelSmall
                )
            }
        }
    }
}


@Composable
fun RunningTraceDialog(traceName: String) {
    LoadingDialog(loadingMessage = "Running $traceName trace...")
}

@Preview(showBackground = true)
@Preview(uiMode = Configuration.UI_MODE_NIGHT_YES, showBackground = true)
@Composable
fun PreviewTraceUtilityNetworkScreen() {
    SampleAppTheme {
        Surface {
            TraceOptions(
                isTraceButtonEnabled = true,
                hintText = "Trace options",
                utilityTraceType = UtilityTraceType.Downstream,
                pointType = PointType.Start,
                traceState = TraceState.ADD_STARTING_POINT,
                onTraceTypeSelected = { },
                onPointTypeChanged = { },
                onResetSelected = { },
                onTraceSelected = { }
            )
        }
    }
}


@Preview(showBackground = true)
@Preview(uiMode = Configuration.UI_MODE_NIGHT_YES, showBackground = true)
@Composable
fun PreviewTerminalConfigurationDialog() {
    SampleAppTheme { Surface { TerminalConfigurationDialog(listOf()) {} } }
}


@Preview(showBackground = true)
@Preview(uiMode = Configuration.UI_MODE_NIGHT_YES, showBackground = true)
@Composable
fun PreviewRunningTraceDialog() {
    SampleAppTheme { Surface { RunningTraceDialog("Downstream") } }
}