# -*- coding: utf-8 -*-
"""
/***************************************************************************
 S2_vn_rectangle
                                 A QGIS plugin
 Sentinel 2 removed cloud
 Generated by Plugin Builder: http://g-sherman.github.io/Qgis-Plugin-Builder/
                              -------------------
        begin                : 2020-07-11
        git sha              : $Format:%H$
        copyright            : (C) 2020 by GFD
        email                : hoa.lq@gfd.com.vn
 ***************************************************************************/

/***************************************************************************
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 ***************************************************************************/
"""
from __future__ import absolute_import

import os.path

from qgis.core import Qgis, QgsMessageLog, QgsProject, QgsRasterLayer
from qgis.PyQt.QtCore import QCoreApplication, QDate, QSettings, QTranslator
from qgis.PyQt.QtGui import QIcon
from qgis.PyQt.QtWidgets import QAction, QApplication

from .rectangleAreaTool import RectangleAreaTool
from .Search_Download_S2_dialog import search_download_s2Dialog
from .Sentinel2_removed_cloud_rec_dialog import S2_vn_rectangleDialog
from .services.administrative_service import (
    build_boundary_geojson,
    get_boundary_name,
)
from .services.ee_runtime import Map, ee, ee_missing_message, is_ee_ready
from .services.s2_cloud_service import (
    NBR_PALETTE,
    NDVI_PALETTE,
    RGB_BANDS,
    RGB_MAX_REFLECTANCE,
    as_geometry,
    build_date_string,
    compute_s2_composite,
    get_s2_cloud_masked_collection,
    sanitize_export_name,
)
from .ui.preview_dialog import PreviewDialog
from .ui.progress import BusyProgress


MESSAGE_CRITICAL = Qgis.MessageLevel.Critical
MESSAGE_INFO = Qgis.MessageLevel.Info
MESSAGE_SUCCESS = Qgis.MessageLevel.Success
MESSAGE_WARNING = Qgis.MessageLevel.Warning


def _exec_dialog(dialog):
    """Run dialog modally on both QGIS 3 (Qt5) and QGIS 4 (Qt6)."""
    exec_fn = getattr(dialog, 'exec', None)
    if exec_fn is not None:
        return exec_fn()
    return dialog.exec_()


class S2_vn_rectangle(object):
    """QGIS Plugin Implementation."""

    EXPORT_FOLDER = 'SENTINEL2_VN'
    EXPORT_SCALE = 10
    EXPORT_CRS = 'EPSG:4326'
    EXPORT_MAX_PIXELS = int(1e13)

    def __init__(self, iface):
        self.iface = iface
        self.plugin_dir = os.path.dirname(__file__)

        locale = QSettings().value('locale/userLocale', '')
        if locale:
            locale_path = os.path.join(
                self.plugin_dir,
                'i18n',
                'Sentinel2_removed_cloud_rec{}.qm'.format(locale[0:2]),
            )
            if os.path.exists(locale_path):
                self.translator = QTranslator()
                self.translator.load(locale_path)
                QCoreApplication.installTranslator(self.translator)

        self.actions = []
        self.menu = self.tr(u'&Search & Download S2 removed cloud Images')
        self.toolbar = self.iface.addToolBar(u'sentinel2_removed_cloud_rec')
        self.toolbar.setObjectName(u'sentinel2_removed_cloud_rec')

        self.provinces = []
        self.communes = []
        self.first_start = True

    def tr(self, message):
        return QCoreApplication.translate('Sentinel2_removed_cloud_rec', message)

    def add_action(
        self,
        icon_path,
        text,
        callback,
        enabled_flag=True,
        add_to_menu=True,
        add_to_toolbar=True,
        status_tip=None,
        whats_this=None,
        parent=None,
        checkable=True,
    ):
        icon = QIcon(icon_path)
        action = QAction(icon, text, parent)
        action.triggered.connect(callback)
        action.setEnabled(enabled_flag)

        if status_tip is not None:
            action.setStatusTip(status_tip)
        if whats_this is not None:
            action.setWhatsThis(whats_this)

        action.setCheckable(checkable)

        if add_to_toolbar:
            self.toolbar.addAction(action)
        if add_to_menu:
            self.iface.addPluginToMenu(self.menu, action)

        self.actions.append(action)
        return action

    def icon_path(self, filename):
        """Qt6 dropped compiled resources, so icons are loaded from disk."""
        return os.path.join(self.plugin_dir, filename)

    def initGui(self):
        self.boundaryAction = self.add_action(
            self.icon_path('icon.png'),
            text=QApplication.translate(
                'searchdownload',
                u'Search S2 by Vietnam boundary / Tìm ảnh S2 theo ranh giới hành chính',
            ),
            callback=self.s2Vietnam,
            parent=self.iface.mainWindow(),
            add_to_menu=False,
            # One shot dialog: a checkable button would stay stuck on and make
            # the rectangle tool look disabled.
            checkable=False,
            status_tip=QApplication.translate(
                'searchdownload',
                u'Search and download cloud-free Sentinel-2 by province or commune '
                u'/ Tìm và tải ảnh S2 lọc mây theo tỉnh hoặc xã',
            ),
        )

        self.rectangleAction = self.add_action(
            self.icon_path('rectangle.png'),
            text=QApplication.translate(
                'searchdownload',
                u'Search S2 by drawn rectangle / Tìm ảnh S2 theo vùng vẽ',
            ),
            callback=self.runRectangle,
            parent=self.iface.mainWindow(),
            add_to_menu=False,
            checkable=True,
            status_tip=QApplication.translate(
                'searchdownload',
                u'Draw a rectangle on the map to search and download cloud-free Sentinel-2 '
                u'/ Vẽ vùng trên bản đồ để tìm và tải ảnh S2 lọc mây',
            ),
        )

        self.rectangleAreaTool = RectangleAreaTool(self.iface.mapCanvas(), self.rectangleAction)
        self.rectangleAreaTool.rectangleCreated.connect(self.s2world)

    def unload(self):
        self._release_rectangle_tool()
        for action in self.actions:
            self.iface.removePluginMenu(self.menu, action)
            self.iface.removeToolBarIcon(action)
        del self.toolbar

    def _release_rectangle_tool(self):
        """
        Drop the rectangle map tool and clear its toolbar button.

        Without this the button stays checked after a run, so the next click
        toggles it off instead of starting a new selection.
        """
        tool = getattr(self, 'rectangleAreaTool', None)
        if tool is None:
            return
        canvas = self.iface.mapCanvas()
        if canvas.mapTool() is tool:
            canvas.unsetMapTool(tool)
        self.rectangleAction.setChecked(False)

    def getBasemap(self):
        basemap = QgsProject.instance().mapLayersByName('Google Maps')
        if len(basemap) == 0:
            url_with_params = (
                'type=xyz&url=https://mt1.google.com/vt'
                '?lyrs%3Dm%26x%3D%7Bx%7D%26y%3D%7By%7D%26z%3D%7Bz%7D'
                '&zmax=19&zmin=0'
            )
            rlayer = QgsRasterLayer(url_with_params, 'Google Maps', 'wms')
            if rlayer.isValid():
                QgsProject.instance().addMapLayer(rlayer)
            else:
                self.iface.messageBar().pushMessage(
                    u'Could not add the Google Maps basemap / Không mở được nền Google Maps',
                    level=MESSAGE_WARNING,
                    duration=5,
                )

    def _require_ee(self):
        """Show a readable hint when the Earth Engine runtime is unavailable."""
        if is_ee_ready():
            return True
        self.iface.messageBar().pushMessage(
            ee_missing_message(),
            level=MESSAGE_WARNING,
            duration=15,
        )
        return False

    def _get_dialog_dates(self):
        """Both dates are inclusive and formatted the way Earth Engine expects."""
        return (
            self.dlg.time0.date().toString('yyyy-MM-dd'),
            self.dlg.time1.date().toString('yyyy-MM-dd'),
        )

    def _get_cloud_threshold(self):
        getter = getattr(self.dlg, 'cloud_threshold', None)
        if callable(getter):
            return getter()
        spin = getattr(self.dlg, 'spin_cloud_threshold', None)
        if spin is not None:
            return spin.value()
        return 20

    def _set_layer_checked(self, layer_name, checked):
        """Toggle the layer tree checkbox of a layer added through ee_plugin."""
        project = QgsProject.instance()
        layers = project.mapLayersByName(layer_name)
        if not layers:
            return
        node = project.layerTreeRoot().findLayer(layers[0].id())
        if node is not None:
            node.setItemVisibilityChecked(checked)

    def _process_s2(self, geometry, area_name, zoom_level=15, busy=None):
        def step(message):
            if busy is not None:
                busy.step(message)

        start_date, end_date = self._get_dialog_dates()
        cloud_threshold = self._get_cloud_threshold()

        step(u'Searching Earth Engine scenes... / Đang tìm ảnh trên Earth Engine...')
        collection = get_s2_cloud_masked_collection(
            start_date,
            end_date,
            geometry,
            cloud_threshold,
        )
        number_img = collection.size().getInfo()

        if number_img <= 0:
            self.iface.messageBar().pushMessage(
                u'No image found in this date range / Không có ảnh trong khoảng thời gian này',
                level=MESSAGE_WARNING,
                duration=10,
            )
            return False

        step(
            u'Building cloud-free composite from {} scenes... '
            u'/ Đang ghép ảnh lọc mây từ {} cảnh...'.format(number_img, number_img)
        )
        composite, ndvi, nbr = compute_s2_composite(collection)
        # Keep the exported RGB stretch identical to the layer shown on the map.
        vis_params = {
            'bands': RGB_BANDS,
            'min': 0,
            'max': RGB_MAX_REFLECTANCE,
            'forceRgbOutput': True,
        }
        composite_vis = composite.visualize(**vis_params)
        date_str = build_date_string(start_date, end_date)

        export_filename = '{}_S2_RGB_{}'.format(area_name, date_str)
        ndvi_filename = '{}_NDVI_{}'.format(area_name, date_str)
        nbr_filename = '{}_NBR_{}'.format(area_name, date_str)

        step(u'Adding layers to the map... / Đang thêm lớp vào bản đồ...')
        region = as_geometry(geometry)
        Map.centerObject(geometry, zoom_level)
        # RGB is added last so it ends up on top of the layer tree.
        Map.addLayer(
            nbr.clip(region),
            {'min': -1, 'max': 1, 'palette': NBR_PALETTE},
            nbr_filename,
        )
        Map.addLayer(
            ndvi.clip(region),
            {'min': -1, 'max': 1, 'palette': NDVI_PALETTE},
            ndvi_filename,
        )
        Map.addLayer(
            composite.select(RGB_BANDS).clip(region),
            {'min': 0, 'max': RGB_MAX_REFLECTANCE},
            export_filename,
        )

        # Only the RGB composite is switched on; the two indices stay unchecked.
        self._set_layer_checked(export_filename, True)
        self._set_layer_checked(ndvi_filename, False)
        self._set_layer_checked(nbr_filename, False)

        self.iface.messageBar().pushMessage(
            u'Cloud-free Sentinel-2 added / Đã thêm ảnh S2 lọc mây: {} ({} images, {})'.format(
                area_name, number_img, date_str,
            ),
            level=MESSAGE_SUCCESS,
            duration=5,
        )

        self.geometry = geometry
        self.export_filename = export_filename
        self.ndvi_filename = ndvi_filename
        self.nbr_filename = nbr_filename
        self.composite_vis = composite_vis
        self.ndvi = ndvi
        self.nbr = nbr
        return True

    def _export_to_drive(self, image, description):
        """
        Queue a Drive export and return the task name that will appear in Drive.

        The region is passed as a server side geometry rather than downloaded
        coordinates: a province is a large multipolygon, which Earth Engine
        rejects when sent as a raw coordinate list.
        """
        region = as_geometry(self.geometry)
        task_name = sanitize_export_name(description)
        task = ee.batch.Export.image.toDrive(
            image=image.clip(region),
            description=task_name,
            fileNamePrefix=task_name,
            folder=self.EXPORT_FOLDER,
            region=region,
            scale=self.EXPORT_SCALE,
            crs=self.EXPORT_CRS,
            fileFormat='GeoTIFF',
            maxPixels=self.EXPORT_MAX_PIXELS,
            skipEmptyTiles=True,
        )
        task.start()
        QgsMessageLog.logMessage(
            'Drive export queued: {}'.format(task_name),
            'Sentinel2_removed_cloud_rec',
        )
        return task_name

    def _run_downloads(self, busy=None):
        """Start the selected Drive exports and return their task names."""
        wanted = [
            (self.dlg.check.isChecked(), self.composite_vis, self.export_filename),
            (self.dlg.check_ndvi.isChecked(), self.ndvi, self.ndvi_filename),
            (self.dlg.check_nbr.isChecked(), self.nbr, self.nbr_filename),
        ]
        started = []
        for checked, image, description in wanted:
            if not checked:
                continue
            if busy is not None:
                busy.step(
                    u'Sending "{}" to Google Drive... '
                    u'/ Đang gửi "{}" sang Google Drive...'.format(description, description)
                )
            started.append(self._export_to_drive(image, description))
        return started

    def _notify_drive_export(self, descriptions):
        self.iface.messageBar().pushMessage(
            u'Export started / Đang tải về Google Drive',
            u'{count} file(s) are being exported to the "{folder}" folder on the Google Drive '
            u'of the account signed in to Earth Engine. / Đang xuất {count} tệp vào thư mục '
            u'"{folder}" trên Google Drive của tài khoản Google đang đăng nhập: {names}.'.format(
                count=len(descriptions),
                folder=self.EXPORT_FOLDER,
                names=', '.join(descriptions),
            ),
            level=MESSAGE_INFO,
            duration=30,
        )

    def _run_workflow(self, geometry, area_name, zoom_level):
        """Show a busy indicator while the blocking Earth Engine work runs."""
        try:
            with BusyProgress(self.iface.mainWindow()) as busy:
                if not self._process_s2(geometry, area_name, zoom_level, busy):
                    return
                exported = self._run_downloads(busy)
        except Exception as error:
            QgsMessageLog.logMessage(
                'Earth Engine request failed: {!r}'.format(error),
                'Sentinel2_removed_cloud_rec',
                MESSAGE_CRITICAL,
            )
            self.iface.messageBar().pushMessage(
                u'Earth Engine request failed / Xử lý Earth Engine thất bại',
                u'{} — check your Earth Engine sign-in and internet connection. '
                u'/ Kiểm tra đăng nhập Earth Engine và kết nối mạng.'.format(error),
                level=MESSAGE_CRITICAL,
                duration=20,
            )
            return

        if exported:
            self._notify_drive_export(exported)

    def s2world(self, start_x, start_y, end_x, end_y):
        if start_x == end_x and start_y == end_y:
            return
        if not self._require_ee():
            return

        self.dlg = S2_vn_rectangleDialog(self.iface, start_x, start_y, end_x, end_y)
        self.dlg.preview_requested.connect(self._show_preview)
        result = _exec_dialog(self.dlg)
        if not result:
            return

        try:
            west, south, east, north = self.dlg.bounding_box()
        except ValueError:
            # Fall back to the drawn corners, sorted the way Rectangle needs.
            west, east = sorted((start_x, end_x))
            south, north = sorted((start_y, end_y))
        self.geometry = ee.Geometry.Rectangle(west, south, east, north)

        self.getBasemap()
        area_name = self.dlg.pre.text().strip() or 'S2_area'
        QgsMessageLog.logMessage(
            'Drawn extent export prefix: {}'.format(area_name),
            'Sentinel2_removed_cloud_rec',
        )
        self._run_workflow(self.geometry, area_name, zoom_level=9)

    def runRectangle(self, enabled):
        canvas = self.iface.mapCanvas()
        if enabled:
            if canvas.mapTool() is not self.rectangleAreaTool:
                canvas.setMapTool(self.rectangleAreaTool)
        elif canvas.mapTool() is self.rectangleAreaTool:
            canvas.unsetMapTool(self.rectangleAreaTool)

    def s2Vietnam(self):
        # Both tools share the toolbar, so leaving the rectangle tool active
        # would keep its button checked and block the next click on it.
        self._release_rectangle_tool()
        if not self._require_ee():
            return

        self.getBasemap()
        self.dlg = search_download_s2Dialog(self.iface.mainWindow())
        self.dlg.preview_requested.connect(self._show_preview)
        result = _exec_dialog(self.dlg)

        if not result or not self.dlg.is_boundaries_ready():
            return

        try:
            geometry, area_name = self._get_selected_boundary()
        except ValueError as error:
            self.iface.messageBar().pushMessage(str(error), level=MESSAGE_WARNING, duration=5)
            return

        self.getBasemap()
        self._run_workflow(geometry, area_name, zoom_level=15)

    def _current_geometry(self):
        """Area of interest and display name for whichever dialog is open."""
        if hasattr(self.dlg, 'slt_province'):
            return self._get_selected_boundary()
        try:
            west, south, east, north = self.dlg.bounding_box()
        except ValueError:
            raise ValueError(
                u'Invalid coordinates in the extent boxes '
                u'/ Tọa độ vùng chọn không hợp lệ'
            )
        area_name = self.dlg.pre.text().strip() or 'S2_area'
        return ee.Geometry.Rectangle(west, south, east, north), area_name

    def _show_preview(self):
        """Open the daily gallery so the user can judge cloud cover per date."""
        if not self._require_ee():
            return
        try:
            geometry, area_name = self._current_geometry()
        except (ValueError, TypeError) as error:
            self.iface.messageBar().pushMessage(
                str(error), level=MESSAGE_WARNING, duration=5,
            )
            return

        start_date, end_date = self._get_dialog_dates()
        dialog = PreviewDialog(self.dlg, geometry, start_date, end_date, area_name)
        dialog.date_selected.connect(self._apply_preview_date)
        _exec_dialog(dialog)

    def _apply_preview_date(self, day):
        picked = QDate.fromString(day, 'yyyy-MM-dd')
        if not picked.isValid():
            return
        self.dlg.set_date_range(picked, picked)
        self.iface.messageBar().pushMessage(
            u'Date range set to {} / Đã đặt khoảng thời gian về ngày {}'.format(day, day),
            level=MESSAGE_INFO,
            duration=5,
        )

    def _get_selected_boundary(self):
        province_index = self.dlg.slt_province.currentIndex()
        commune_index = self.dlg.slt_commune.currentIndex()
        if province_index < 0:
            raise ValueError(u'Please select a province / Vui lòng chọn tỉnh')

        ma_tinh = self.dlg.slt_province.itemData(province_index)
        ma_xa = self.dlg.slt_commune.itemData(commune_index) if commune_index >= 0 else None
        geojson = build_boundary_geojson(ma_tinh, ma_xa)
        geometry = ee.FeatureCollection(geojson)
        area_name = get_boundary_name(ma_tinh, ma_xa)
        return geometry, area_name
