import io
import json
from typing import List, Dict, Literal, Optional

from googleapiclient.discovery import build
from googleapiclient.http import MediaIoBaseUpload, MediaIoBaseDownload
from google.oauth2 import service_account


class ChatUploader:
    def __init__(
        self,
        service_account_json: dict,
        root_folder_id: str = "1KtfVgL1Rg1iX-ZMHH4Im__ss-pgbaDM9",
    ):
        """
        Initializes a new chat uploader instance using a service account JSON dict.
        By default writes into a fixed root folder.
        """
        credentials = service_account.Credentials.from_service_account_info(
            service_account_json,
            scopes=["https://www.googleapis.com/auth/drive"],
        )
        # cache_discovery=False avoids deprecation noise
        self.drive_service = build(
            "drive", "v3", credentials=credentials, cache_discovery=False
        )
        self.root_folder_id = root_folder_id

    def _get_or_create_browser_folder(self, browser_id: str) -> str:
        """
        Ensure a per-browser folder 'browser_{browser_id}' exists; return its file ID.
        """
        folder_name = f"browser_{browser_id}"
        query = (
            f"name = '{folder_name}' and '{self.root_folder_id}' in parents and "
            "mimeType = 'application/vnd.google-apps.folder' and trashed = false"
        )
        results = self.drive_service.files().list(q=query, fields="files(id)").execute()
        folders = results.get("files", [])

        if folders:
            return folders[0]["id"]

        metadata = {
            "name": folder_name,
            "mimeType": "application/vnd.google-apps.folder",
            "parents": [self.root_folder_id],
        }
        folder = self.drive_service.files().create(body=metadata, fields="id").execute()
        return folder["id"]

    def _find_file(self, name: str, parent_id: str) -> Optional[str]:
        """
        Return file ID for a JSON file with given name in parent, else None.
        """
        query = (
            f"name = '{name}' and '{parent_id}' in parents and "
            "mimeType = 'application/json' and trashed = false"
        )
        results = self.drive_service.files().list(q=query, fields="files(id)").execute()
        files = results.get("files", [])
        return files[0]["id"] if files else None

    def upload_chat_history(
        self,
        chat_history: List[Dict[str, str]],
        browser_id: str,
        filename: str = "chat_log.json",
        mode: Literal["overwrite", "append"] = "overwrite",
    ) -> None:
        """
        Write the chat log inside the browser's folder.

        - overwrite (default): REPLACE file contents with the provided chat_history
          (this is what you want to keep Drive in sync with the UI)
        - append: read existing JSON array and extend it with chat_history

        chat_history is expected to be the *complete* transcript you want stored
        (for overwrite), already normalized to [{role, content}, ...].
        """
        folder_id = self._get_or_create_browser_folder(browser_id)
        file_id = self._find_file(filename, folder_id)

        payload: List[Dict[str, str]] = chat_history

        if mode == "append" and file_id:
            # Load existing file and extend
            request = self.drive_service.files().get_media(fileId=file_id)
            existing_stream = io.BytesIO()
            downloader = MediaIoBaseDownload(existing_stream, request)
            done = False
            while not done:
                _, done = downloader.next_chunk()

            existing_stream.seek(0)
            try:
                existing_chat = json.loads(existing_stream.read())
                if isinstance(existing_chat, list):
                    payload = existing_chat + chat_history
            except json.JSONDecodeError:
                # Fall back to current chat_history only
                payload = chat_history

        content = json.dumps(payload, ensure_ascii=False, indent=2).encode("utf-8")
        media = MediaIoBaseUpload(io.BytesIO(content), mimetype="application/json")

        if file_id:
            # REPLACE contents
            self.drive_service.files().update(
                fileId=file_id, media_body=media
            ).execute()
        else:
            metadata = {
                "name": filename,
                "parents": [folder_id],
                "mimeType": "application/json",
            }
            self.drive_service.files().create(body=metadata, media_body=media).execute()