from __future__ import annotations

import json
import logging
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path
from typing import Any

from niannian.storage.cloud_base import parse_remote_root_parts
from niannian.storage.cloud_http import HttpOpener, open_with_strip_auth_redirect
from niannian.storage.onedrive_api import onedrive_path
from niannian.storage.upload_checkpoint import ChunkedUploadHooks, UploadCheckpointStore
from niannian.storage.upload_progress import UploadProgressTracker

logger = logging.getLogger(__name__)

DEFAULT_API_BASE = "https://www.googleapis.com/drive/v3"
DEFAULT_UPLOAD_BASE = "https://www.googleapis.com/upload/drive/v3"
DEFAULT_UPLOAD_MAX_BYTES = 5 * 1024 * 1024
DEFAULT_CHUNK_SIZE_BYTES = 8 * 1024 * 1024
DEFAULT_MAX_UPLOAD_BYTES = 512 * 1024 * 1024

def _split_chunks(data: bytes, chunk_size: int) -> list[bytes]:
    if chunk_size <= 0:
        raise ValueError("chunk_size must be positive")
    return [data[i : i + chunk_size] for i in range(0, len(data), chunk_size)]


class GoogleDriveAPIError(Exception):
    pass


def google_drive_path(rel_path: str, remote_root: str) -> str:
    return onedrive_path(rel_path, remote_root)


class GoogleDriveClient:
    """Google Drive v3 small-file upload/download under remote_root folder chain."""

    def __init__(
        self,
        access_token: str,
        *,
        remote_root: str | None = None,
        api_base: str = DEFAULT_API_BASE,
        upload_max_bytes: int = DEFAULT_UPLOAD_MAX_BYTES,
        chunk_size_bytes: int = DEFAULT_CHUNK_SIZE_BYTES,
        max_upload_bytes: int = DEFAULT_MAX_UPLOAD_BYTES,
        timeout: float = 60.0,
        http_opener: HttpOpener | None = None,
        state_path: Path | None = None,
        credential_manager: Any | None = None,
    ) -> None:
        self._access_token = access_token
        self._root_parts = parse_remote_root_parts(remote_root)
        self._remote_root = "/".join(self._root_parts)
        self._api_base = api_base.rstrip("/")
        self._upload_base = DEFAULT_UPLOAD_BASE
        self._upload_max_bytes = int(upload_max_bytes)
        self._chunk_size_bytes = int(chunk_size_bytes)
        self._max_upload_bytes = int(max_upload_bytes)
        self._timeout = timeout
        self._open = http_opener or open_with_strip_auth_redirect
        self._state_path = state_path
        self._credential_manager = credential_manager
        self._folder_cache: dict[str, str] = {}
        self._load_state()

    def physical_path(self, rel_path: str) -> str:
        return f"/{google_drive_path(rel_path, self._remote_root)}"

    def upload_object(
        self,
        rel_path: str,
        data: bytes,
        *,
        progress: UploadProgressTracker | None = None,
        checkpoint: UploadCheckpointStore | None = None,
        job_id: str | None = None,
    ) -> None:
        if len(data) > self._max_upload_bytes:
            logger.warning(
                "Google Drive upload rejected: %s exceeds max_upload_bytes (%d > %d)",
                rel_path,
                len(data),
                self._max_upload_bytes,
            )
            raise GoogleDriveAPIError(
                f"file too large: {rel_path} ({len(data)} bytes > max {self._max_upload_bytes})"
            )
        parent_id = self._ensure_path_parent(google_drive_path(rel_path, self._remote_root))
        name = rel_path.strip("/").split("/")[-1]
        existing = self._find_child(parent_id, name)
        if len(data) <= self._upload_max_bytes:
            self._upload_simple(parent_id, name, data, existing)
            return
        hooks = ChunkedUploadHooks(
            provider="google_drive",
            rel_path=rel_path,
            remote_root=self._remote_root,
            data=data,
            chunk_size=self._chunk_size_bytes,
            progress=progress,
            checkpoint=checkpoint,
            job_id=job_id,
        )
        if hooks.enabled:
            resume = hooks.begin()
            try:
                self._upload_resumable(parent_id, name, data, existing, hooks=hooks, resume=resume)
                hooks.done()
            except Exception as exc:
                hooks.failed(str(exc))
                raise
        else:
            self._upload_resumable(parent_id, name, data, existing)

    def _upload_simple(
        self,
        parent_id: str,
        name: str,
        data: bytes,
        existing: str | None,
    ) -> None:
        if existing:
            url = f"{self._api_base}/files/{existing}?uploadType=media"
            self._request(url, method="PATCH", data=data, content_type="application/octet-stream")
            return
        meta = json.dumps({"name": name, "parents": [parent_id]}).encode("utf-8")
        boundary = "niannian_gdrive_boundary"
        body = (
            f"--{boundary}\r\n"
            "Content-Type: application/json; charset=UTF-8\r\n\r\n"
        ).encode("utf-8") + meta + (
            f"\r\n--{boundary}\r\n"
            "Content-Type: application/octet-stream\r\n\r\n"
        ).encode("utf-8") + data + f"\r\n--{boundary}--\r\n".encode("utf-8")
        url = f"{self._upload_base}/files?uploadType=multipart"
        self._request(
            url,
            method="POST",
            data=body,
            content_type=f"multipart/related; boundary={boundary}",
        )

    def _upload_resumable(
        self,
        parent_id: str,
        name: str,
        data: bytes,
        existing: str | None,
        *,
        hooks: ChunkedUploadHooks | None = None,
        resume: Any | None = None,
    ) -> None:
        upload_url: str | None = None
        start_chunk = 0
        if resume is not None and resume.provider_state.get("resumable_url"):
            upload_url = str(resume.provider_state["resumable_url"])
            start_chunk = int(resume.start_chunk_index)

        if not upload_url:
            if existing:
                init_url = f"{self._upload_base}/files/{existing}?uploadType=resumable"
                init_body = None
            else:
                init_url = f"{self._upload_base}/files?uploadType=resumable"
                init_body = {"name": name, "parents": [parent_id]}
            upload_url = self._init_resumable_session(init_url, init_body)

        total = len(data)
        offset = 0
        for chunk_index, chunk in enumerate(_split_chunks(data, self._chunk_size_bytes)):
            if chunk_index < start_chunk:
                offset += len(chunk)
                continue
            start = offset
            end = offset + len(chunk) - 1
            content_range = f"bytes {start}-{end}/{total}"
            self._request(
                upload_url,
                method="PUT",
                data=chunk,
                content_type="application/octet-stream",
                extra_headers={"Content-Range": content_range},
            )
            offset = end + 1
            if hooks is not None:
                hooks.after_chunk(
                    chunk_index,
                    provider_state={
                        "resumable_url": upload_url,
                        "next_offset": offset,
                        "parent_id": parent_id,
                        "name": name,
                        "existing": existing,
                    },
                )

    def _init_resumable_session(
        self,
        url: str,
        body: dict[str, Any] | None,
    ) -> str:
        data = None
        content_type = None
        if body is not None:
            data = json.dumps(body).encode("utf-8")
            content_type = "application/json; charset=UTF-8"
        req = urllib.request.Request(
            url,
            data=data,
            method="POST",
            headers=self._headers(content_type=content_type),
        )
        try:
            with self._open(req, timeout=self._timeout) as resp:
                location = resp.headers.get("Location") or resp.headers.get("location")
                if not location:
                    raise GoogleDriveAPIError("resumable upload missing Location header")
                return location
        except urllib.error.HTTPError as exc:
            detail = exc.read().decode("utf-8", "replace") if exc.fp else str(exc)
            raise GoogleDriveAPIError(f"HTTP {exc.code}: {detail[:200]}") from exc

    def download_object(self, rel_path: str) -> bytes | None:
        file_id = self._resolve_file_id(google_drive_path(rel_path, self._remote_root))
        if not file_id:
            return None
        url = f"{self._api_base}/files/{file_id}?alt=media"
        try:
            return self._request(url, method="GET")
        except GoogleDriveAPIError as exc:
            if "HTTP 404" in str(exc):
                return None
            raise

    def _load_state(self) -> None:
        if not self._state_path or not self._state_path.is_file():
            return
        try:
            raw = json.loads(self._state_path.read_text(encoding="utf-8"))
            if raw.get("remote_root") != self._remote_root:
                return
            self._folder_cache = dict(raw.get("folders") or {})
        except (OSError, json.JSONDecodeError):
            pass

    def _save_state(self) -> None:
        if not self._state_path:
            return
        self._state_path.parent.mkdir(parents=True, exist_ok=True)
        payload = {"remote_root": self._remote_root, "folders": self._folder_cache}
        self._state_path.write_text(
            json.dumps(payload, ensure_ascii=False, indent=2) + "\n",
            encoding="utf-8",
        )

    def _headers(
        self,
        *,
        content_type: str | None = None,
        extra_headers: dict[str, str] | None = None,
        use_auth: bool = True,
    ) -> dict[str, str]:
        headers: dict[str, str] = {"User-Agent": "niannian/0.2.9"}
        if use_auth:
            headers["Authorization"] = f"Bearer {self._access_token}"
        if content_type:
            headers["Content-Type"] = content_type
        if extra_headers:
            headers.update(extra_headers)
        return headers

    def _request(
        self,
        url: str,
        *,
        method: str = "GET",
        data: bytes | None = None,
        content_type: str | None = None,
        extra_headers: dict[str, str] | None = None,
        use_auth: bool = True,
        auth_retry: bool = True,
    ) -> bytes:
        req = urllib.request.Request(
            url,
            data=data,
            method=method,
            headers=self._headers(
                content_type=content_type,
                extra_headers=extra_headers,
                use_auth=use_auth,
            ),
        )
        try:
            with self._open(req, timeout=self._timeout) as resp:
                return resp.read()
        except urllib.error.HTTPError as exc:
            if method == "GET" and exc.code in (301, 302, 303, 307, 308):
                location = exc.headers.get("Location")
                if location:
                    return self._request(
                        location,
                        method="GET",
                        use_auth=False,
                        auth_retry=False,
                    )
            if auth_retry and use_auth and exc.code == 401 and self._credential_manager is not None:
                refresh = self._credential_manager.refresh("google_drive")
                if refresh.ok:
                    new_token = self._credential_manager.get_access_token("google_drive")
                    if new_token:
                        self._access_token = new_token
                        return self._request(
                            url,
                            method=method,
                            data=data,
                            content_type=content_type,
                            extra_headers=extra_headers,
                            auth_retry=False,
                        )
            detail = exc.read().decode("utf-8", "replace") if exc.fp else str(exc)
            raise GoogleDriveAPIError(f"HTTP {exc.code}: {detail[:200]}") from exc

    def _request_json(self, url: str, *, method: str = "GET", body: dict[str, Any] | None = None) -> dict[str, Any]:
        data = None
        content_type = None
        if body is not None:
            data = json.dumps(body).encode("utf-8")
            content_type = "application/json"
        raw = self._request(url, method=method, data=data, content_type=content_type)
        if not raw:
            return {}
        payload = json.loads(raw.decode("utf-8", "replace"))
        if isinstance(payload, dict) and payload.get("error"):
            err = payload["error"]
            msg = err.get("message") if isinstance(err, dict) else str(err)
            raise GoogleDriveAPIError(f"Drive error: {msg}")
        return payload if isinstance(payload, dict) else {}

    def _find_child(self, parent_id: str, name: str) -> str | None:
        q = urllib.parse.quote(
            f"'{parent_id}' in parents and name = '{name.replace(chr(39), chr(92)+chr(39))}' and trashed = false"
        )
        url = (
            f"{self._api_base}/files?q={q}&fields=files(id,name)&pageSize=1"
            "&supportsAllDrives=true&includeItemsFromAllDrives=true"
        )
        payload = self._request_json(url)
        files = payload.get("files") or []
        if files:
            return str(files[0]["id"])
        return None

    def _ensure_folder(self, parent_id: str, name: str) -> str:
        key = f"{parent_id}/{name}"
        cached = self._folder_cache.get(key)
        if cached:
            return cached
        existing = self._find_child(parent_id, name)
        if existing:
            self._folder_cache[key] = existing
            self._save_state()
            return existing
        payload = self._request_json(
            f"{self._api_base}/files?supportsAllDrives=true",
            method="POST",
            body={"name": name, "mimeType": "application/vnd.google-apps.folder", "parents": [parent_id]},
        )
        folder_id = str(payload["id"])
        self._folder_cache[key] = folder_id
        self._save_state()
        return folder_id

    def _ensure_path_parent(self, file_path: str) -> str:
        parts = [p for p in file_path.split("/") if p]
        if not parts:
            raise GoogleDriveAPIError("empty path")
        parent_id = "root"
        built: list[str] = []
        for part in parts[:-1]:
            built.append(part)
            cache_key = "/".join(built)
            if cache_key in self._folder_cache:
                parent_id = self._folder_cache[cache_key]
                continue
            parent_id = self._ensure_folder(parent_id, part)
            self._folder_cache[cache_key] = parent_id
            self._save_state()
        return parent_id

    def _resolve_file_id(self, file_path: str) -> str | None:
        parts = [p for p in file_path.split("/") if p]
        if not parts:
            return None
        parent_id = self._ensure_path_parent(file_path)
        return self._find_child(parent_id, parts[-1])
