"""Draw a chart with the ChartKit MCP server and put it on a Google Slides slide.

The chart is drawn for the box under the slide title and lands as an image
whose alt text points back to the ChartKit chart, so anyone with the ChartKit
extension can open it in the deck and edit it.

Setup:
    pip install requests google-api-python-client google-auth
    export CHARTKIT_API_KEY=ck_team_...
    export GOOGLE_APPLICATION_CREDENTIALS=/path/to/service-account.json
    export CHARTKIT_MCP_URL=https://...   # optional, defaults to the production server

Share the presentation with the service account's email address as an editor.

Usage:
    python chartkit_slides.py PRESENTATION_ID            # adds a new slide
    python chartkit_slides.py PRESENTATION_ID SLIDE_ID   # uses an existing slide
"""

import os
import sys
import uuid

import google.auth
import requests
from googleapiclient.discovery import build

MCP_URL = os.environ.get("CHARTKIT_MCP_URL", "https://www.getchartkit.app/mcp")
SLIDES_SCOPES = ["https://www.googleapis.com/auth/presentations"]
EMU_PER_PT = 12700

SLIDE_TITLE = "Online sales grew 17% a year since 2021"
CHART = {
    "chart_type": "clusteredBar",
    "table": {
        "categories": ["2021", "2022", "2023", "2024", "2025"],
        "series": [{"name": "Online sales"}],
        "values": [[41.2, 48.9, 57.3, 66.0, 78.4]],
    },
    "unit": "EUR m",
    "annotations": [
        {"kind": "cagr", "from": {"category_index": 0}, "to": {"category_index": 4}},
    ],
}


class ChartKitError(RuntimeError):
    pass


def call_tool(name, arguments, api_key, timeout=60):
    """Calls one ChartKit MCP tool and returns the tool result."""
    response = requests.post(
        MCP_URL,
        headers={
            "Authorization": f"Bearer {api_key}",
            "Accept": "application/json, text/event-stream",
        },
        json={
            "jsonrpc": "2.0",
            "id": str(uuid.uuid4()),
            "method": "tools/call",
            "params": {"name": name, "arguments": arguments},
        },
        timeout=timeout,
    )
    if response.status_code == 401:
        raise ChartKitError("ChartKit rejected the API key (HTTP 401). Check CHARTKIT_API_KEY.")
    response.raise_for_status()

    message = response.json()
    if "error" in message:
        error = message["error"]
        raise ChartKitError(f"{name} failed with JSON-RPC error {error['code']}: {error['message']}")

    result = message["result"]
    if result.get("isError"):
        text = " ".join(block.get("text", "") for block in result.get("content", []))
        raise ChartKitError(f"{name} returned an error: {text}")
    return result


def render_png(image_url):
    # The first request renders the PNG, so Google's fetch during createImage gets a finished image.
    response = requests.get(image_url, timeout=60)
    response.raise_for_status()


def page_size_pt(slides, presentation_id):
    page_size = slides.presentations().get(presentationId=presentation_id, fields="pageSize").execute()["pageSize"]

    def to_pt(dimension):
        return dimension["magnitude"] / EMU_PER_PT if dimension.get("unit") == "EMU" else dimension["magnitude"]

    return to_pt(page_size["width"]), to_pt(page_size["height"])


def add_title_slide(slides, presentation_id, title):
    """Appends a Title only slide with the given title and returns its object id."""
    slide_id = f"slide_{uuid.uuid4().hex[:12]}"
    title_id = f"title_{uuid.uuid4().hex[:12]}"
    slides.presentations().batchUpdate(
        presentationId=presentation_id,
        body={
            "requests": [
                {
                    "createSlide": {
                        "objectId": slide_id,
                        "slideLayoutReference": {"predefinedLayout": "TITLE_ONLY"},
                        "placeholderIdMappings": [
                            {"layoutPlaceholder": {"type": "TITLE", "index": 0}, "objectId": title_id},
                        ],
                    }
                },
                {"insertText": {"objectId": title_id, "text": title}},
            ]
        },
    ).execute()
    return slide_id


def insert_chart(slides, presentation_id, slide_id, chart, left, top, width, height):
    """Places the chart PNG in the given box (points) and sets the alt text ChartKit reads."""
    image_id = f"chartkit_{uuid.uuid4().hex[:16]}"
    slides.presentations().batchUpdate(
        presentationId=presentation_id,
        body={
            "requests": [
                {
                    "createImage": {
                        "objectId": image_id,
                        "url": chart["imageUrl"],
                        "elementProperties": {
                            "pageObjectId": slide_id,
                            "size": {
                                "width": {"magnitude": width, "unit": "PT"},
                                "height": {"magnitude": height, "unit": "PT"},
                            },
                            "transform": {
                                "scaleX": 1,
                                "scaleY": 1,
                                "translateX": left,
                                "translateY": top,
                                "unit": "PT",
                            },
                        },
                    }
                },
                {
                    "updatePageElementAltText": {
                        "objectId": image_id,
                        "title": "ChartKit",
                        "description": chart["slidesAltText"],
                    }
                },
            ]
        },
    ).execute()
    return image_id


def chart_box(page_width, page_height, top=110, margin=36):
    """The box under the title, in points."""
    return margin, top, page_width - 2 * margin, page_height - top - margin


def centred(box, size):
    """The image's position and size: the chart's size, centred in the box it was drawn for."""
    left, top, width, height = box
    return left + (width - size["width"]) / 2, top + (height - size["height"]) / 2, size["width"], size["height"]


def main():
    if len(sys.argv) not in (2, 3):
        print(__doc__)
        sys.exit(2)
    presentation_id = sys.argv[1]
    slide_id = sys.argv[2] if len(sys.argv) == 3 else None

    # Reads GOOGLE_APPLICATION_CREDENTIALS, or the service account attached to the
    # Cloud Run service, GKE pod or VM this runs on.
    credentials, _ = google.auth.default(scopes=SLIDES_SCOPES)
    slides = build("slides", "v1", credentials=credentials, cache_discovery=False)
    box = chart_box(*page_size_pt(slides, presentation_id))

    # ChartKit lays the chart out for the box and returns sizePt, the size to insert it at. The
    # extension reopens the chart at the size it was drawn for, so the edit view matches the slide.
    arguments = {**CHART, "size": {"width_pt": box[2], "height_pt": box[3]}}
    chart = call_tool("create_chart", arguments, api_key=os.environ["CHARTKIT_API_KEY"])["structuredContent"]
    render_png(chart["imageUrl"])

    if slide_id is None:
        slide_id = add_title_slide(slides, presentation_id, SLIDE_TITLE)
    image_id = insert_chart(slides, presentation_id, slide_id, chart, *centred(box, chart["sizePt"]))

    print(f"Chart {chart['chartId']} is on slide {slide_id} as image {image_id}.")
    print(f"https://docs.google.com/presentation/d/{presentation_id}/edit#slide=id.{slide_id}")


if __name__ == "__main__":
    main()
