from __future__ import annotations

import io
import logging

from muleline.storage.base import StorageBackend

log = logging.getLogger(__name__)


class GoogleDriveStorage(StorageBackend):
    """Google Drive storage backend using a service account."""

    def __init__(self, credentials_json: str, folder_id: str = ""):
        self.folder_id = folder_id
        self._cache: dict[str, str] = {}

        try:
            from google.oauth2 import service_account
            from googleapiclient.discovery import build
        except ImportError:
            raise ImportError("Install Google API client: pip install muleline[gdrive]")

        creds = service_account.Credentials.from_service_account_file(
            credentials_json,
            scopes=["https://www.googleapis.com/auth/drive"],
        )
        self._service = build("drive", "v3", credentials=creds, cache_discovery=False)

    def _files(self):
        return self._service.files()

    def _find_id(self, key: str) -> str | None:
        if key in self._cache:
            return self._cache[key]
        safe_key = key.replace("\\", "\\\\").replace("'", "\\'")
        query = f"name='{safe_key}' and trashed=false"
        if self.folder_id:
            safe_folder = self.folder_id.replace("\\", "\\\\").replace("'", "\\'")
            query += f" and '{safe_folder}' in parents"
        try:
            result = self._files().list(q=query, fields="files(id)", pageSize=1).execute()
            files = result.get("files", [])
            if files:
                file_id = files[0]["id"]
                self._cache[key] = file_id
                return file_id
        except Exception as e:
            log.error("Google Drive lookup failed for %s: %s", key, e)
        return None

    def save(self, key: str, content: bytes, mime_type: str = "") -> bool:
        try:
            from googleapiclient.http import MediaInMemoryUpload

            media = MediaInMemoryUpload(content, mimetype=mime_type or "application/octet-stream", resumable=False)
            existing_id = self._find_id(key)

            if existing_id:
                self._files().update(fileId=existing_id, media_body=media).execute()
            else:
                meta: dict = {"name": key}
                if self.folder_id:
                    meta["parents"] = [self.folder_id]
                result = self._files().create(body=meta, media_body=media, fields="id").execute()
                self._cache[key] = result["id"]

            return True
        except Exception as e:
            log.error("Google Drive save failed for %s: %s", key, e)
            return False

    def read(self, key: str) -> bytes | None:
        file_id = self._find_id(key)
        if not file_id:
            return None
        try:
            request = self._files().get_media(fileId=file_id)
            buf = io.BytesIO()
            from googleapiclient.http import MediaIoBaseDownload
            downloader = MediaIoBaseDownload(buf, request)
            done = False
            while not done:
                _, done = downloader.next_chunk()
            return buf.getvalue()
        except Exception as e:
            log.error("Google Drive read failed for %s: %s", key, e)
            return None

    def delete(self, key: str) -> bool:
        file_id = self._find_id(key)
        if not file_id:
            return False
        try:
            self._files().delete(fileId=file_id).execute()
            self._cache.pop(key, None)
            return True
        except Exception as e:
            log.error("Google Drive delete failed for %s: %s", key, e)
            return False

    def exists(self, key: str) -> bool:
        return self._find_id(key) is not None

    def url(self, key: str, expires: int = 3600) -> str | None:
        file_id = self._find_id(key)
        if not file_id:
            return None
        try:
            result = self._files().get(fileId=file_id, fields="webViewLink").execute()
            return result.get("webViewLink")
        except Exception as e:
            log.error("Google Drive url failed for %s: %s", key, e)
            return None
