"""
wivexx_deliver.py — client photo delivery pipeline  (v1.14)
═══════════════════════════════════════════════════════════════════════
Two independent halves, so convert-only and upload-only both work:

  A. RAW → JPG conversion  (rawpy + Pillow)
       convert_folder(src_dir, ...) develops every .ARW/.RAW in a folder
       into full-quality JPGs, either by truly developing the sensor data
       ("develop", camera white balance) or by lifting the full-size JPEG
       the camera already embedded in the RAW ("preview", instant).

  B. Google Drive delivery  (google-api-python-client + google-auth-oauthlib)
       upload_folder(local_dir, gallery_name, ...) creates a fresh folder
       on Drive, uploads the photos (optionally renamed into a clean
       sequence), makes the folder public, and returns the share link.

Everything heavy (rawpy, the google libs) is imported lazily inside the
functions, so this module always imports — callers degrade gracefully and
show a "pip install …" message when a dependency is missing (same pattern
as opencv in the Caption Studio).

Public API (used by wivexx_gui.py / wivexx_api.py):
    raw_available() -> bool
    count_raws(src_dir) -> int
    convert_folder(src_dir, out_dir=None, mode="develop", quality=95,
                   include_jpegs=False, progress=None, cancel_evt=None) -> dict
    drive_available() -> bool
    drive_auth_status() -> "ready" | "need_credentials" | "need_consent"
    drive_setup_help() -> str
    upload_folder(local_dir, gallery_name, name_sequence=None, start=1,
                  make_public=True, progress=None, cancel_evt=None) -> dict
"""

import json
import os
import shutil
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Callable, Optional

# ── config (best-effort: module must import even if config is odd) ───────────────
try:
    import wivexx_config as _cfg
except Exception:                                    # pragma: no cover
    _cfg = None


def _c(name: str, default):
    return getattr(_cfg, name, default) if _cfg else default


_HERE = os.path.dirname(os.path.abspath(__file__))

# RAW extensions — reuse the camera-RAW set already curated for media import.
try:
    from wivexx_media_import import PHOTO_EXTS as _PHOTO_EXTS
    _RAW_EXTS = {e for e in _PHOTO_EXTS if e not in {
        ".jpg", ".jpeg", ".png", ".heic", ".heif", ".webp", ".bmp",
        ".tif", ".tiff", ".gif"}}
except Exception:
    _RAW_EXTS = {
        ".raw", ".arw", ".sr2", ".cr2", ".cr3", ".nef", ".nrw", ".orf",
        ".rw2", ".dng", ".raf", ".srw", ".pef", ".gpr", ".x3f"}

_JPEG_PASSTHRU = {".jpg", ".jpeg", ".png"}

# Google Drive: app manages only the files it creates (least privilege).
_SCOPES = ["https://www.googleapis.com/auth/drive.file"]
_DRIVE_STATE_FILE = os.path.join(_HERE, "wivexx_drive_state.json")

_logger = None


def set_logger(logger) -> None:
    global _logger
    _logger = logger


def _log(level: str, msg: str, *a) -> None:
    if _logger:
        getattr(_logger, level, _logger.info)(msg, *a)


# ══════════════════════════════════════════════════════════════════════════════
#  A.  RAW → JPG conversion
# ══════════════════════════════════════════════════════════════════════════════

def raw_available() -> bool:
    """True if rawpy + Pillow can be imported (RAW develop is possible)."""
    try:
        import rawpy            # noqa: F401
        import PIL.Image        # noqa: F401
        return True
    except Exception:
        return False


def list_raws(src_dir: str) -> list:
    """Top-level RAW files in *src_dir* (not recursive — a shoot folder)."""
    try:
        return sorted(
            os.path.join(src_dir, f) for f in os.listdir(src_dir)
            if os.path.splitext(f)[1].lower() in _RAW_EXTS
            and os.path.isfile(os.path.join(src_dir, f)))
    except Exception:
        return []


def count_raws(src_dir: str) -> int:
    return len(list_raws(src_dir))


def _copy_exif(src_raw: str, dst_jpg: str) -> None:
    """Best-effort: transplant capture date / camera tags from the RAW into the
    developed JPG.  Silently skipped if piexif isn't installed or can't read
    the RAW (developed pixels are unaffected either way)."""
    try:
        import piexif
        src_exif = piexif.load(src_raw)
        keep = {"0th": {}, "Exif": {}, "1st": {}, "GPS": src_exif.get("GPS", {}),
                "thumbnail": None}
        for tag in (piexif.ImageIFD.Make, piexif.ImageIFD.Model,
                    piexif.ImageIFD.Software, piexif.ImageIFD.Artist,
                    piexif.ImageIFD.Copyright):
            if tag in src_exif.get("0th", {}):
                keep["0th"][tag] = src_exif["0th"][tag]
        for tag in (piexif.ExifIFD.DateTimeOriginal,
                    piexif.ExifIFD.DateTimeDigitized,
                    piexif.ExifIFD.LensModel, piexif.ExifIFD.FNumber,
                    piexif.ExifIFD.ExposureTime, piexif.ExifIFD.ISOSpeedRatings,
                    piexif.ExifIFD.FocalLength):
            if tag in src_exif.get("Exif", {}):
                keep["Exif"][tag] = src_exif["Exif"][tag]
        piexif.insert(piexif.dump(keep), dst_jpg)
    except Exception:
        pass


def _convert_one(src: str, out_dir: str, mode: str, quality: int,
                 auto_bright: bool = True, exposure: float = 1.0) -> str:
    """Develop / extract one RAW → JPG.  Returns 'converted' | 'skipped' | 'failed'."""
    import io as _io
    import rawpy
    from PIL import Image, ImageEnhance

    base = os.path.splitext(os.path.basename(src))[0]
    dst = os.path.join(out_dir, base + ".jpg")
    if os.path.exists(dst):
        return "skipped"
    tmp = dst + ".part"
    try:
        with rawpy.imread(src) as raw:
            if mode == "preview":
                # Lift the full-size JPEG the camera embedded — instant, matches
                # what was seen on the camera screen.  Fall back to a develop if
                # the embedded thumb isn't a full JPEG.
                try:
                    thumb = raw.extract_thumb()
                    if thumb.format == rawpy.ThumbFormat.JPEG:
                        if abs(exposure - 1.0) < 1e-3:
                            with open(tmp, "wb") as fh:
                                fh.write(thumb.data)
                        else:
                            im = ImageEnhance.Brightness(
                                Image.open(_io.BytesIO(thumb.data)).convert("RGB")
                            ).enhance(exposure)
                            im.save(tmp, "JPEG", quality=int(quality),
                                    optimize=True)
                        os.replace(tmp, dst)
                        _copy_exif(src, dst)
                        return "converted"
                except Exception:
                    pass
            # True develop: demosaic with the camera's white balance.  Auto-
            # brightness matches the camera's own rendering — with it OFF the
            # JPGs came out ~0.8 EV darker than the embedded preview (v1.15.1 fix).
            rgb = raw.postprocess(use_camera_wb=True,
                                  no_auto_bright=not auto_bright,
                                  output_bps=8)
        img = Image.fromarray(rgb)
        if abs(exposure - 1.0) > 1e-3:
            img = ImageEnhance.Brightness(img).enhance(exposure)
        img.save(tmp, "JPEG", quality=int(quality), optimize=True)
        os.replace(tmp, dst)
        _copy_exif(src, dst)
        return "converted"
    except Exception as e:
        _log("warning", "deliver: convert failed %s: %s", src, e)
        try:
            if os.path.exists(tmp):
                os.remove(tmp)
        except Exception:
            pass
        return "failed"


def convert_folder(src_dir: str, out_dir: Optional[str] = None,
                   mode: str = "develop", quality: int = 95,
                   include_jpegs: bool = False,
                   auto_bright: Optional[bool] = None,
                   exposure: Optional[float] = None,
                   progress: Optional[Callable] = None,
                   cancel_evt: Optional[threading.Event] = None) -> dict:
    """Develop every RAW in *src_dir* into JPGs under *out_dir*
    (default <src_dir>\\<DELIVER_OUTPUT_SUBDIR>).

    mode          'develop' (rawpy demosaic) | 'preview' (embedded camera JPEG)
    include_jpegs also copy existing .jpg/.png through, so a mixed folder
                  delivers whole.
    progress(done, total, name) is called after each file.
    Returns {ok, out_dir, converted, skipped, failed, cancelled, error}.
    """
    res = {"ok": False, "out_dir": "", "converted": 0, "skipped": 0,
           "failed": 0, "cancelled": False, "error": ""}
    if not src_dir or not os.path.isdir(src_dir):
        res["error"] = "That folder doesn't exist."
        return res
    if not raw_available():
        res["error"] = ("RAW conversion needs the rawpy library. Install it "
                        "once with:  pip install rawpy")
        return res

    raws = list_raws(src_dir)
    extras = []
    if include_jpegs:
        try:
            extras = [os.path.join(src_dir, f) for f in os.listdir(src_dir)
                      if os.path.splitext(f)[1].lower() in _JPEG_PASSTHRU
                      and os.path.isfile(os.path.join(src_dir, f))]
        except Exception:
            extras = []
    if not raws and not extras:
        res["error"] = "I found no .ARW / RAW photos in that folder."
        return res

    if not out_dir:
        out_dir = os.path.join(src_dir, _c("DELIVER_OUTPUT_SUBDIR", "JPG"))
    try:
        os.makedirs(out_dir, exist_ok=True)
    except Exception as e:
        res["error"] = f"Couldn't create the output folder: {e}"
        return res
    res["out_dir"] = out_dir

    if auto_bright is None:
        auto_bright = bool(_c("DELIVER_AUTO_BRIGHTEN", True))
    if exposure is None:
        exposure = float(_c("DELIVER_BRIGHTNESS", 1.0))

    total = len(raws) + len(extras)
    done = 0
    workers = max(1, int(_c("DELIVER_CONVERT_WORKERS", 4)))

    # Copy any pass-through JPG/PNG first (cheap), then develop the RAWs.
    for p in extras:
        if cancel_evt and cancel_evt.is_set():
            res["cancelled"] = True
            break
        dst = os.path.join(out_dir, os.path.basename(p))
        try:
            if not os.path.exists(dst):
                shutil.copy2(p, dst)
                res["converted"] += 1
            else:
                res["skipped"] += 1
        except Exception:
            res["failed"] += 1
        done += 1
        if progress:
            progress(done, total, os.path.basename(p))

    if not (cancel_evt and cancel_evt.is_set()) and raws:
        with ThreadPoolExecutor(max_workers=workers) as ex:
            futs = {ex.submit(_convert_one, p, out_dir, mode, quality,
                              auto_bright, exposure): p
                    for p in raws}
            for fut in as_completed(futs):
                done += 1
                try:
                    res[fut.result()] += 1
                except Exception:
                    res["failed"] += 1
                if progress:
                    progress(done, total, os.path.basename(futs[fut]))
                if cancel_evt and cancel_evt.is_set():
                    res["cancelled"] = True
                    try:
                        ex.shutdown(wait=False, cancel_futures=True)
                    except TypeError:               # Python < 3.9
                        pass
                    break

    res["ok"] = res["converted"] > 0 or res["skipped"] > 0
    return res


# ══════════════════════════════════════════════════════════════════════════════
#  B.  Google Drive delivery
# ══════════════════════════════════════════════════════════════════════════════

def _credentials_path() -> str:
    p = _c("DRIVE_CREDENTIALS_FILE", "credentials.json")
    return p if os.path.isabs(p) else os.path.join(_HERE, p)


def _token_path() -> str:
    p = _c("DRIVE_TOKEN_FILE", "wivexx_drive_token.json")
    return p if os.path.isabs(p) else os.path.join(_HERE, p)


def drive_available() -> bool:
    """True if the google client libraries are installed."""
    try:
        import googleapiclient.discovery     # noqa: F401
        import google_auth_oauthlib.flow     # noqa: F401
        import google.oauth2.credentials      # noqa: F401
        return True
    except Exception:
        return False


def drive_auth_status() -> str:
    """'ready' (token cached & usable) | 'need_consent' (creds present, browser
    sign-in still required) | 'need_credentials' (no credentials.json / libs)."""
    if not drive_available():
        return "need_credentials"
    if not os.path.exists(_credentials_path()):
        return "need_credentials"
    if not os.path.exists(_token_path()):
        return "need_consent"
    try:
        from google.oauth2.credentials import Credentials
        creds = Credentials.from_authorized_user_file(_token_path(), _SCOPES)
        if creds and creds.valid:
            return "ready"
        if creds and creds.expired and creds.refresh_token:
            return "ready"                          # refresh happens silently
        return "need_consent"
    except Exception:
        return "need_consent"


def drive_setup_help() -> str:
    return (
        "Google Drive needs a one-time setup (about 5 minutes) — see "
        "DRIVE_SETUP.md.  In short:\n"
        "  1. console.cloud.google.com → create a free project.\n"
        "  2. Enable the “Google Drive API”.\n"
        "  3. OAuth consent screen → External → add your email as a Test user.\n"
        "  4. Credentials → Create OAuth client ID → type “Desktop app”.\n"
        f"  5. Download it as credentials.json into:\n     {_HERE}\n"
        "Then run a delivery once — a browser opens so you can grant access, "
        "and it's automatic from then on.")


def _ensure_service():
    """Build an authorised Drive API client, running the one-time browser
    consent if needed.  Raises with a friendly message if creds are missing."""
    from google.oauth2.credentials import Credentials
    from google.auth.transport.requests import Request
    from google_auth_oauthlib.flow import InstalledAppFlow
    from googleapiclient.discovery import build

    creds = None
    tok = _token_path()
    if os.path.exists(tok):
        try:
            creds = Credentials.from_authorized_user_file(tok, _SCOPES)
        except Exception:
            creds = None
    if not creds or not creds.valid:
        if creds and creds.expired and creds.refresh_token:
            creds.refresh(Request())
        else:
            cred_file = _credentials_path()
            if not os.path.exists(cred_file):
                raise FileNotFoundError(
                    "credentials.json not found — see DRIVE_SETUP.md.")
            flow = InstalledAppFlow.from_client_secrets_file(cred_file, _SCOPES)
            creds = flow.run_local_server(port=0)
        try:
            with open(tok, "w", encoding="utf-8") as fh:
                fh.write(creds.to_json())
        except Exception as e:
            _log("warning", "deliver: could not cache Drive token: %s", e)
    return build("drive", "v3", credentials=creds, cache_discovery=False)


def _drive_state() -> dict:
    try:
        with open(_DRIVE_STATE_FILE, encoding="utf-8") as fh:
            d = json.load(fh)
        return d if isinstance(d, dict) else {}
    except Exception:
        return {}


def _save_drive_state(d: dict) -> None:
    try:
        with open(_DRIVE_STATE_FILE, "w", encoding="utf-8") as fh:
            json.dump(d, fh, indent=2)
    except Exception:
        pass


def _ensure_parent(service, name: str) -> Optional[str]:
    """Return the id of the top-level gallery parent folder, creating it once
    and caching its id.  None disables the parent (galleries land in My Drive)."""
    if not name:
        return None
    state = _drive_state()
    cached = state.get("parent_id")
    if cached:
        try:                                        # verify it still exists
            service.files().get(fileId=cached, fields="id, trashed").execute()
            return cached
        except Exception:
            pass
    # Look for one we created earlier (drive.file scope sees app-created files).
    try:
        q = ("mimeType='application/vnd.google-apps.folder' and trashed=false "
             f"and name='{name}'")
        hits = service.files().list(q=q, spaces="drive",
                                    fields="files(id,name)").execute()
        files = hits.get("files", [])
        if files:
            state["parent_id"] = files[0]["id"]
            _save_drive_state(state)
            return files[0]["id"]
    except Exception:
        pass
    try:
        folder = service.files().create(
            body={"name": name,
                  "mimeType": "application/vnd.google-apps.folder"},
            fields="id").execute()
        state["parent_id"] = folder["id"]
        _save_drive_state(state)
        return folder["id"]
    except Exception as e:
        _log("warning", "deliver: parent folder failed: %s", e)
        return None


def _seq_name(seq: str, idx: int, width: int, ext: str) -> str:
    return f"{seq}-{idx:0{width}d}{ext}"


def upload_folder(local_dir: str, gallery_name: str,
                  name_sequence: Optional[str] = None, start: int = 1,
                  make_public: bool = True,
                  include_raw: bool = False,
                  extra_raw_dir: Optional[str] = None,
                  progress: Optional[Callable] = None,
                  cancel_evt: Optional[threading.Event] = None,
                  pause_evt: Optional[threading.Event] = None) -> dict:
    """Create a Drive folder named *gallery_name*, upload every image in
    *local_dir* (renamed to "<name_sequence>-001.jpg…" when given), make it
    public, and return the share link.

    include_raw   also upload the camera RAW originals (.ARW/.CR3/.NEF/…) — i.e.
                  deliver the unprocessed files. Off by default so a normal
                  gallery stays JPG-only.
    extra_raw_dir when delivering "RAW + JPG", the developed JPGs live in
                  *local_dir* but the RAW originals are elsewhere (the shoot
                  folder); pull the RAWs from here too, into the same gallery.
    Returns {ok, link, folder_id, uploaded, failed, cancelled, error}.
    """
    res = {"ok": False, "link": "", "folder_id": "", "uploaded": 0,
           "failed": 0, "cancelled": False, "error": ""}
    if not local_dir or not os.path.isdir(local_dir):
        res["error"] = "The folder to upload doesn't exist."
        return res
    if not drive_available():
        res["error"] = ("Google Drive needs its libraries. Install once with:  "
                        "pip install google-api-python-client "
                        "google-auth-oauthlib google-auth-httplib2")
        return res

    exts = {".jpg", ".jpeg", ".png", ".webp", ".tif", ".tiff", ".heic"}
    if include_raw:
        exts = exts | _RAW_EXTS
    files = sorted(
        os.path.join(local_dir, f) for f in os.listdir(local_dir)
        if os.path.splitext(f)[1].lower() in exts
        and os.path.isfile(os.path.join(local_dir, f)))
    # "RAW + JPG": the developed JPGs are in local_dir, but the RAW originals
    # live in the shoot folder — append them so one gallery holds both.
    if include_raw and extra_raw_dir and os.path.isdir(extra_raw_dir) \
            and os.path.abspath(extra_raw_dir) != os.path.abspath(local_dir):
        try:
            raws = sorted(
                os.path.join(extra_raw_dir, f) for f in os.listdir(extra_raw_dir)
                if os.path.splitext(f)[1].lower() in _RAW_EXTS
                and os.path.isfile(os.path.join(extra_raw_dir, f)))
            files = files + raws
        except Exception:
            pass
    if not files:
        res["error"] = ("No RAW or photo files to upload in that folder."
                        if include_raw else "No photos to upload in that folder.")
        return res

    try:
        from googleapiclient.http import MediaFileUpload
        service = _ensure_service()
    except FileNotFoundError as e:
        res["error"] = str(e)
        return res
    except Exception as e:
        res["error"] = f"Drive sign-in failed: {e}"
        return res

    parent_id = _ensure_parent(service, _c("DRIVE_PARENT_FOLDER",
                                           "Wivexx Client Galleries"))
    try:
        body = {"name": gallery_name or "Client gallery",
                "mimeType": "application/vnd.google-apps.folder"}
        if parent_id:
            body["parents"] = [parent_id]
        folder = service.files().create(
            body=body, fields="id, webViewLink").execute()
        folder_id = folder["id"]
        res["folder_id"] = folder_id
    except Exception as e:
        res["error"] = f"Couldn't create the Drive folder: {e}"
        return res

    total = len(files)
    width = max(3, len(str(start + total - 1)))
    seq = (name_sequence or "").strip()
    try:
        sizes = [os.path.getsize(f) for f in files]
    except Exception:
        sizes = [1] * total
    total_bytes = sum(sizes) or 1
    done_bytes = 0
    t0 = time.time()
    paused_total = 0.0
    for i, path in enumerate(files):
        # honour pause — and don't let paused time skew the ETA
        if pause_evt is not None and pause_evt.is_set():
            _p0 = time.time()
            while pause_evt.is_set():
                if cancel_evt and cancel_evt.is_set():
                    break
                time.sleep(0.2)
            paused_total += time.time() - _p0
        if cancel_evt and cancel_evt.is_set():
            res["cancelled"] = True
            break
        ext = os.path.splitext(path)[1].lower()
        if ext == ".jpeg":
            ext = ".jpg"
        name = (_seq_name(seq, start + i, width, ext) if seq
                else os.path.basename(path))
        try:
            media = MediaFileUpload(path, resumable=True)
            request = service.files().create(
                body={"name": name, "parents": [folder_id]},
                media_body=media, fields="id")
            response = None
            while response is None:
                _status, response = request.next_chunk()
            res["uploaded"] += 1
        except Exception as e:
            _log("warning", "deliver: upload failed %s: %s", path, e)
            res["failed"] += 1
        done_bytes += sizes[i] if i < len(sizes) else 0
        eta = None
        elapsed = time.time() - t0 - paused_total
        if done_bytes > 0 and (i + 1) < total and elapsed > 0.5:
            rate = done_bytes / elapsed
            if rate > 0:
                eta = (total_bytes - done_bytes) / rate
        if progress:
            progress(i + 1, total, name, eta)

    if make_public and res["uploaded"] > 0:
        try:
            service.permissions().create(
                fileId=folder_id,
                body={"type": "anyone", "role": "reader"}).execute()
        except Exception as e:
            _log("warning", "deliver: make-public failed: %s", e)

    try:
        meta = service.files().get(fileId=folder_id,
                                   fields="webViewLink").execute()
        res["link"] = meta.get("webViewLink",
                               f"https://drive.google.com/drive/folders/{folder_id}")
    except Exception:
        res["link"] = f"https://drive.google.com/drive/folders/{folder_id}"

    res["ok"] = res["uploaded"] > 0
    return res
