# -*- coding: utf-8 -*-
"""
Google Earth Engine + Planet download engine.

Handles GEE authentication, image collection queries,
cloud masking, compositing, index computation, and
GeoTIFF / PNG export.  Also supports Planet Data API.

Author: Loc Vu Trung / Slow Forest
License: GPL-3.0
"""

import json
import os
import zipfile
from datetime import datetime

from .satellites import SATELLITES


def _force_ipv4():
    """Monkey-patch socket to prefer IPv4 over IPv6."""
    import socket
    if getattr(socket, "_lvt_ipv4_patched", False):
        return
    _orig = socket.getaddrinfo

    def _ipv4_only(host, port, family=0, type_=0, proto=0, flags=0):
        return _orig(host, port, socket.AF_INET, type_, proto, flags)

    socket.getaddrinfo = _ipv4_only
    socket._lvt_ipv4_patched = True


class GEEDownloader:
    """Core engine for satellite image download via Google Earth Engine."""

    def __init__(self):
        _force_ipv4()
        import ee
        self.ee = ee

        # Try to initialize with multiple fallback strategies
        # so any authenticated user can use it regardless of project
        for attempt in self._init_strategies():
            try:
                attempt()
                return  # success
            except Exception:
                continue

        # All strategies failed — trigger auth flow
        ee.Authenticate()
        # After auth, try again
        for attempt in self._init_strategies():
            try:
                attempt()
                return
            except Exception:
                continue
        # Last resort
        ee.Initialize()

    def _init_strategies(self):
        """Yield initialization callables in priority order."""
        ee = self.ee
        return [
            lambda: ee.Initialize(),  # user default
            lambda: ee.Initialize(project="earthengine-legacy"),  # public
            lambda: ee.Initialize(project="earthengine-public"),  # public alt
        ]

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def search_images(self, sat_key, aoi_geojson, start_date, end_date,
                      max_cloud=20):
        """Query individual images and return per-image metadata."""
        ee = self.ee
        cfg = SATELLITES[sat_key]
        aoi = self._make_aoi(aoi_geojson)

        col = ee.ImageCollection(cfg["collection"]) \
            .filterBounds(aoi) \
            .filterDate(start_date, end_date) \
            .sort("system:time_start")

        if cfg.get("cloud_field") and max_cloud < 100:
            col = col.filter(
                ee.Filter.lte(cfg["cloud_field"], max_cloud))

        if cfg.get("extra_filter"):
            for k, v in cfg["extra_filter"].items():
                col = col.filter(ee.Filter.eq(k, v))

        count = col.size().getInfo()
        if count == 0:
            return []
        if count > 200:
            count = 200

        info_list = col.limit(count).getInfo()
        results = []
        for feat in (info_list or {}).get("features", []):
            props = feat.get("properties", {}) or {}
            img_id = feat.get("id", "")
            ts = props.get("system:time_start", 0)
            date_str = datetime.utcfromtimestamp(
                ts / 1000).strftime("%Y-%m-%d") if ts else "?"

            cloud = None
            if cfg.get("cloud_field"):
                cloud = props.get(cfg["cloud_field"])
                if cloud is not None:
                    cloud = round(float(cloud), 1)

            results.append({
                "id": img_id,
                "date": date_str,
                "cloud": cloud,
            })

        return results

    def get_thumb_url(self, image_id, sat_key, aoi_geojson, bands_key):
        """Generate thumbnail URL for a single image (low-res preview)."""
        ee = self.ee
        cfg = SATELLITES[sat_key]
        bands_cfg = cfg["bands"][bands_key]
        aoi = self._make_aoi(aoi_geojson)

        image = ee.Image(image_id)

        if bands_cfg.get("index"):
            image = self._apply_cloud_mask_single(image, sat_key)
            image = self._compute_index(image, sat_key, bands_key)
            vis = {
                "min": bands_cfg.get("min", -0.2),
                "max": bands_cfg.get("max", 0.9),
                "palette": bands_cfg.get("palette", []),
                "dimensions": "64x64",
                "region": aoi,
                "format": "png",
            }
        else:
            bands = bands_cfg["ee_bands"][:3]
            image = image.select(bands)
            vis = {
                "bands": bands,
                "min": bands_cfg.get("min", 0),
                "max": bands_cfg.get("max", 3000),
                "dimensions": "64x64",
                "region": aoi,
                "format": "png",
            }

        return image.clip(aoi).getThumbURL(vis)

    def download_single(self, image_id, sat_key, aoi_geojson,
                        bands_key, scale, output_dir, fmt="GeoTIFF",
                        crop_to_polygon=False):
        """Download a single image by ID (full resolution)."""
        ee = self.ee
        cfg = SATELLITES[sat_key]
        bands_cfg = cfg["bands"][bands_key]
        aoi = self._make_aoi(aoi_geojson)

        image = ee.Image(image_id)
        image = self._apply_cloud_mask_single(image, sat_key)

        if bands_cfg.get("index"):
            image = self._compute_index(image, sat_key, bands_key)
        else:
            image = image.select(bands_cfg["ee_bands"])

        if crop_to_polygon:
            image = image.clip(aoi)
        else:
            image = image.clip(aoi.bounds())

        # Build filename from image ID
        parts = image_id.split("/")
        date_str = parts[-1][:8] if parts else "img"
        safe_bands = bands_key.replace(" ", "_").replace(
            "(", "").replace(")", "").replace("-", "")
        fname = f"{sat_key}_{date_str}_{safe_bands}"

        if fmt == "PNG":
            return self._download_png(image, aoi, bands_cfg,
                                      scale, output_dir, fname)
        else:
            return self._download_geotiff(image, aoi, bands_cfg,
                                          scale, output_dir, fname,
                                          crop_to_polygon)

    # ------------------------------------------------------------------
    # AOI Handling (robust — handles None, bad geometry, etc.)
    # ------------------------------------------------------------------

    def _make_aoi(self, aoi_geojson):
        """Convert GeoJSON dict to ee.Geometry.

        Handles: Polygon, MultiPolygon, FeatureCollection,
        GeometryCollection, simplification.
        """
        ee = self.ee

        if aoi_geojson is None:
            raise ValueError("AOI GeoJSON is None")

        geom_type = aoi_geojson.get("type", "")

        if geom_type == "FeatureCollection":
            fc = ee.FeatureCollection(aoi_geojson)
            return fc.geometry().dissolve(1)

        if geom_type == "Feature":
            geom_data = aoi_geojson.get("geometry")
            if geom_data is None:
                raise ValueError("Feature has no geometry")
            return ee.Geometry(geom_data)

        if geom_type in ("Polygon", "MultiPolygon", "Point",
                         "LineString", "GeometryCollection"):
            coords = aoi_geojson.get("coordinates", [])
            if not coords:
                raise ValueError("Geometry has no coordinates")
            geom = ee.Geometry(aoi_geojson)
            # Simplify if > 5000 vertices
            coords_str = json.dumps(coords)
            vcount = coords_str.count(",") // 2
            if vcount > 5000:
                geom = geom.simplify(maxError=10)
            return geom

        raise ValueError(f"Unsupported geometry type: {geom_type}")

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _apply_cloud_mask_single(self, image, sat_key):
        cfg = SATELLITES[sat_key]

        if cfg.get("cloud_mask_band") == "SCL":
            bad = cfg["cloud_mask_values"]
            scl = image.select("SCL")
            mask = scl.neq(bad[0])
            for v in bad[1:]:
                mask = mask.And(scl.neq(v))
            return image.updateMask(mask)

        elif cfg.get("cloud_mask_bits"):
            bits = cfg["cloud_mask_bits"]
            qa = image.select("QA_PIXEL")
            mask = qa.bitwiseAnd(1 << bits[0]).eq(0)
            for b in bits[1:]:
                mask = mask.And(qa.bitwiseAnd(1 << b).eq(0))
            return image.updateMask(mask)

        return image

    def _compute_index(self, image, sat_key, bands_key):
        cfg = SATELLITES[sat_key]
        bands_cfg = cfg["bands"][bands_key]
        fb = bands_cfg["formula_bands"]

        if bands_key in ("NDVI", "NDWI"):
            return image.normalizedDifference(
                [fb[0], fb[1]]).rename("index")
        elif bands_key == "EVI":
            nir = image.select(fb[0])
            red = image.select(fb[1])
            blue = image.select(fb[2])
            evi = nir.subtract(red).multiply(2.5).divide(
                nir.add(red.multiply(6)).subtract(
                    blue.multiply(7.5)).add(1))
            return evi.rename("index")
        return image

    def _download_geotiff(self, image, aoi, bands_cfg, scale,
                          output_dir, fname, crop_to_polygon=False):
        """Download GeoTIFF via getDownloadURL.

        Uses explicit band config + filePerBand=False to get a single
        multi-band TIFF (avoids ZIP corruption issues).
        """
        import requests

        region = aoi if crop_to_polygon else aoi.bounds()

        # Build explicit bands list for reliable multi-band TIFF
        band_names = image.bandNames().getInfo()
        bands_list = [{"id": b, "scale": scale} for b in band_names]

        url = image.getDownloadURL({
            "bands": bands_list,
            "region": region,
            "filePerBand": False,
            "format": "GEO_TIFF",
            "maxPixels": 1e9,
        })

        # Stream download for large files
        resp = requests.get(url, timeout=600, stream=True)
        resp.raise_for_status()

        content_type = resp.headers.get("Content-Type", "")
        tif_path = os.path.join(output_dir, f"{fname}.tif")

        # Check if response is a ZIP (multi-file)
        first_chunk = b""
        chunks = []
        for chunk in resp.iter_content(chunk_size=65536):
            chunks.append(chunk)
            if len(first_chunk) == 0:
                first_chunk = chunk[:4]

        raw = b"".join(chunks)

        if first_chunk[:2] == b"PK" or "zip" in content_type:
            # ZIP file — extract the TIFF
            zip_path = tif_path.replace(".tif", ".zip")
            with open(zip_path, "wb") as f:
                f.write(raw)

            with zipfile.ZipFile(zip_path) as zf:
                tif_names = [n for n in zf.namelist()
                             if n.lower().endswith(".tif")]
                if tif_names:
                    # Extract all TIFs and use the largest one
                    best = max(tif_names,
                               key=lambda n: zf.getinfo(n).file_size)
                    zf.extract(best, output_dir)
                    extracted = os.path.join(output_dir, best)
                    if extracted != tif_path:
                        if os.path.exists(tif_path):
                            os.remove(tif_path)
                        os.rename(extracted, tif_path)
            try:
                os.remove(zip_path)
            except OSError:
                pass
        else:
            # Direct TIFF
            with open(tif_path, "wb") as f:
                f.write(raw)

        # Validate file size (corrupt files are usually < 1 KB)
        fsize = os.path.getsize(tif_path)
        if fsize < 1024:
            raise IOError(
                f"Downloaded file too small ({fsize} bytes) — "
                f"likely corrupted. Try reducing area or resolution.")

        return tif_path

    def _download_png(self, image, aoi, bands_cfg, scale,
                      output_dir, fname):
        import requests

        if bands_cfg.get("index") or bands_cfg.get("palette"):
            vis = {
                "min": bands_cfg.get("min", -0.2),
                "max": bands_cfg.get("max", 0.9),
                "palette": bands_cfg.get("palette", []),
                "dimensions": "2048x2048",
                "region": aoi,
                "format": "png",
            }
        else:
            vis = {
                "bands": bands_cfg["ee_bands"][:3],
                "min": bands_cfg.get("min", 0),
                "max": bands_cfg.get("max", 3000),
                "dimensions": "2048x2048",
                "region": aoi,
                "format": "png",
            }

        url = image.getThumbURL(vis)
        resp = requests.get(url, timeout=120)
        resp.raise_for_status()

        png_path = os.path.join(output_dir, f"{fname}.png")
        with open(png_path, "wb") as f:
            f.write(resp.content)
        return png_path


# ======================================================================
# Planet Data API
# ======================================================================

class PlanetDownloader:
    """Planet Labs Data API v2 client.

    Requires a Planet API key (free tier available).
    Docs: https://developers.planet.com/docs/data/
    """

    BASE_URL = "https://api.planet.com/data/v1"

    def __init__(self, api_key):
        import requests
        self.session = requests.Session()
        self.session.auth = (api_key, "")
        self._api_key = api_key

    def validate_key(self):
        """Check if the API key is valid."""
        resp = self.session.get(f"{self.BASE_URL}/")
        return resp.status_code == 200

    def search_images(self, aoi_geojson, start_date, end_date,
                      max_cloud=20, item_type="PSScene"):
        """Search Planet catalog.

        item_type: PSScene (3m), SkySatCollect (0.5m)
        Returns list of dicts with id, date, cloud, thumb_url.
        """
        search_filter = {
            "type": "AndFilter",
            "config": [
                {
                    "type": "GeometryFilter",
                    "field_name": "geometry",
                    "config": aoi_geojson,
                },
                {
                    "type": "DateRangeFilter",
                    "field_name": "acquired",
                    "config": {
                        "gte": f"{start_date}T00:00:00Z",
                        "lte": f"{end_date}T23:59:59Z",
                    },
                },
                {
                    "type": "RangeFilter",
                    "field_name": "cloud_cover",
                    "config": {"lte": max_cloud / 100.0},
                },
            ],
        }

        payload = {
            "item_types": [item_type],
            "filter": search_filter,
        }

        resp = self.session.post(
            f"{self.BASE_URL}/quick-search",
            json=payload, timeout=30)
        resp.raise_for_status()
        data = resp.json()

        results = []
        for feat in data.get("features", [])[:100]:
            props = feat.get("properties", {})
            acquired = props.get("acquired", "")
            date_str = acquired[:10] if acquired else "?"
            cloud = props.get("cloud_cover")
            if cloud is not None:
                cloud = round(cloud * 100, 1)

            thumb = feat.get("_links", {}).get("thumbnail", "")

            results.append({
                "id": feat.get("id", ""),
                "date": date_str,
                "cloud": cloud,
                "thumb_url": thumb,
                "item_type": item_type,
            })

        return results

    def get_thumb_url(self, item_id, item_type="PSScene"):
        """Get thumbnail URL for a Planet item."""
        url = (f"{self.BASE_URL}/item-types/{item_type}"
               f"/items/{item_id}/thumb")
        return url

    def download_asset(self, item_id, item_type, asset_type,
                       output_dir):
        """Download a Planet asset (e.g., ortho_analytic_4b).

        Steps: activate -> wait -> download.
        """
        import time

        # Get asset info
        url = (f"{self.BASE_URL}/item-types/{item_type}"
               f"/items/{item_id}/assets")
        resp = self.session.get(url, timeout=30)
        resp.raise_for_status()
        assets = resp.json()

        if asset_type not in assets:
            available = list(assets.keys())
            raise ValueError(
                f"Asset '{asset_type}' not available. "
                f"Available: {available}")

        asset = assets[asset_type]
        status = asset.get("status", "")

        # Activate if needed
        if status == "inactive":
            activate_url = asset.get("_links", {}).get("activate", "")
            if activate_url:
                self.session.post(activate_url, timeout=30)

            # Wait for activation (max 5 min)
            for _ in range(60):
                time.sleep(5)
                resp = self.session.get(url, timeout=30)
                asset = resp.json().get(asset_type, {})
                if asset.get("status") == "active":
                    break
            else:
                raise TimeoutError(
                    "Asset activation timed out (5 min)")

        # Download
        download_url = asset.get("location", "")
        if not download_url:
            raise ValueError("No download URL for asset")

        resp = self.session.get(download_url, stream=True, timeout=600)
        resp.raise_for_status()

        fname = f"planet_{item_type}_{item_id}_{asset_type}.tif"
        fpath = os.path.join(output_dir, fname)
        with open(fpath, "wb") as f:
            for chunk in resp.iter_content(chunk_size=8192):
                f.write(chunk)

        return fpath
