"""Google Workspace connector — Google Sheets and Google Drive.

Sheets and Drive share one OAuth2 credential, so they live in a single connector.
Auth uses the installed-app refresh-token flow: the connector exchanges a stored
refresh token for a short-lived access token on each call (via the shared HTTP
seam, so it is fully testable). Reads are non-sensitive; writes (appending a row,
which mutates a document) are sensitive and route through the policy engine.

Uploading binary files to Drive is intentionally left out for now (multipart);
this first cut covers the high-value read + Sheets-append cases.
"""

from __future__ import annotations

import httpx

from legate.agent.tooling import ToolContext, ToolResult, ToolSpec
from legate.connectors.base import BaseConnector, ConnectionResult
from legate.connectors.runtime import new_async_client

_TOKEN_URL = "https://oauth2.googleapis.com/token"
_SHEETS_BASE = "https://sheets.googleapis.com/v4/spreadsheets"
_DRIVE_BASE = "https://www.googleapis.com/drive/v3"
_MAX_ITEMS = 100


class GoogleAuthError(RuntimeError):
    """Raised when a Google access token cannot be obtained."""


class GoogleConnector(BaseConnector):
    type = "google"
    title = "Google Workspace (Sheets & Drive)"
    config_schema = {
        "type": "object",
        "properties": {
            "default_spreadsheet_id": {"type": "string"},
        },
    }
    secret_schema = {
        "type": "object",
        "properties": {
            "client_id": {"type": "string"},
            "client_secret": {"type": "string"},
            "refresh_token": {"type": "string"},
        },
        "required": ["client_id", "client_secret", "refresh_token"],
    }

    def list_tools(self) -> list[ToolSpec]:
        return [
            ToolSpec(
                name="gsheets.read_range",
                description="Read a cell range from a Google Sheet.",
                parameters={
                    "type": "object",
                    "properties": {
                        "spreadsheet_id": {"type": "string"},
                        "range": {"type": "string", "description": "e.g. 'Sheet1!A1:C10'."},
                    },
                    "required": ["range"],
                },
                is_sensitive=False,
                scopes=["google:sheets:read"],
            ),
            ToolSpec(
                name="gsheets.append_row",
                description="Append a row of values to a Google Sheet range.",
                parameters={
                    "type": "object",
                    "properties": {
                        "spreadsheet_id": {"type": "string"},
                        "range": {"type": "string"},
                        "values": {"type": "array", "description": "Cell values for the new row."},
                    },
                    "required": ["range", "values"],
                },
                is_sensitive=True,
                scopes=["google:sheets:write"],
            ),
            ToolSpec(
                name="gdrive.list_files",
                description="List Drive files, optionally filtered by a Drive query.",
                parameters={
                    "type": "object",
                    "properties": {
                        "query": {"type": "string", "description": "Drive 'q' filter."},
                        "limit": {"type": "integer"},
                    },
                },
                is_sensitive=False,
                scopes=["google:drive:read"],
            ),
            ToolSpec(
                name="gdrive.get_file",
                description="Fetch a Drive file's metadata by id.",
                parameters={
                    "type": "object",
                    "properties": {"file_id": {"type": "string"}},
                    "required": ["file_id"],
                },
                is_sensitive=False,
                scopes=["google:drive:read"],
            ),
        ]

    async def _access_token(self, secrets: dict) -> str:
        data = {
            "client_id": secrets.get("client_id", ""),
            "client_secret": secrets.get("client_secret", ""),
            "refresh_token": secrets.get("refresh_token", ""),
            "grant_type": "refresh_token",
        }
        async with new_async_client() as client:
            resp = await client.post(_TOKEN_URL, data=data)
        if resp.status_code != 200:
            raise GoogleAuthError(f"token exchange failed: HTTP {resp.status_code}")
        token = resp.json().get("access_token")
        if not token:
            raise GoogleAuthError("token response had no access_token")
        return token

    async def _get(self, url: str, token: str, params: dict | None = None) -> httpx.Response:
        async with new_async_client() as client:
            return await client.get(
                url, headers={"Authorization": f"Bearer {token}"}, params=params
            )

    async def _post(self, url: str, token: str, json: dict) -> httpx.Response:
        async with new_async_client() as client:
            return await client.post(url, headers={"Authorization": f"Bearer {token}"}, json=json)

    async def test_connection(self, config: dict, secrets: dict) -> ConnectionResult:
        try:
            token = await self._access_token(secrets)
            resp = await self._get(f"{_DRIVE_BASE}/about", token, params={"fields": "user"})
        except (httpx.HTTPError, GoogleAuthError) as exc:
            return ConnectionResult(ok=False, message=f"connection failed: {exc}")
        if resp.status_code == 200:
            return ConnectionResult(ok=True, message="Google credentials authenticated")
        return ConnectionResult(ok=False, message=f"HTTP {resp.status_code}")

    async def execute(
        self, tool: str, args: dict, config: dict, secrets: dict, ctx: ToolContext
    ) -> ToolResult:
        try:
            token = await self._access_token(secrets)
        except (httpx.HTTPError, GoogleAuthError) as exc:
            return ToolResult(ok=False, error=str(exc))

        try:
            if tool == "gsheets.read_range":
                sheet_id = args.get("spreadsheet_id") or config.get("default_spreadsheet_id")
                if not sheet_id:
                    return ToolResult(ok=False, error="no spreadsheet_id specified")
                rng = args.get("range", "")
                resp = await self._get(f"{_SHEETS_BASE}/{sheet_id}/values/{rng}", token)
                if resp.status_code != 200:
                    return ToolResult(ok=False, error=f"HTTP {resp.status_code}")
                body = resp.json()
                return ToolResult(
                    ok=True, data={"range": body.get("range"), "values": body.get("values", [])}
                )

            if tool == "gsheets.append_row":
                sheet_id = args.get("spreadsheet_id") or config.get("default_spreadsheet_id")
                if not sheet_id:
                    return ToolResult(ok=False, error="no spreadsheet_id specified")
                rng = args.get("range", "")
                url = (
                    f"{_SHEETS_BASE}/{sheet_id}/values/{rng}:append"
                    "?valueInputOption=USER_ENTERED"
                )
                resp = await self._post(url, token, {"values": [args.get("values", [])]})
                if resp.status_code not in (200, 201):
                    return ToolResult(ok=False, error=f"HTTP {resp.status_code}")
                updates = resp.json().get("updates", {})
                return ToolResult(
                    ok=True,
                    data={
                        "updated_range": updates.get("updatedRange"),
                        "updated_rows": updates.get("updatedRows"),
                    },
                )

            if tool == "gdrive.list_files":
                params = {
                    "pageSize": min(int(args.get("limit", 20)), _MAX_ITEMS),
                    "fields": "files(id,name,mimeType,modifiedTime)",
                }
                if args.get("query"):
                    params["q"] = args["query"]
                resp = await self._get(f"{_DRIVE_BASE}/files", token, params=params)
                if resp.status_code != 200:
                    return ToolResult(ok=False, error=f"HTTP {resp.status_code}")
                return ToolResult(ok=True, data={"files": resp.json().get("files", [])})

            if tool == "gdrive.get_file":
                resp = await self._get(
                    f"{_DRIVE_BASE}/files/{args['file_id']}",
                    token,
                    params={"fields": "id,name,mimeType,modifiedTime,size"},
                )
                if resp.status_code != 200:
                    return ToolResult(ok=False, error=f"HTTP {resp.status_code}")
                return ToolResult(ok=True, data=resp.json())

            return ToolResult(ok=False, error=f"unknown google tool: {tool}")
        except httpx.HTTPError as exc:
            return ToolResult(ok=False, error=f"google request failed: {exc}")
