src/shopify_client.py

import mimetypes
import os
import time
import uuid
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import requests
from dotenv import load_dotenv

load_dotenv()


class ShopifyConfigError(RuntimeError):
    pass


_token_cache: Dict[str, Any] = {"token": None, "expires_at": 0.0}


def _get_config() -> Tuple[str, str, str, str]:
    shop_name = os.getenv("SHOPIFY_SHOP_NAME")
    api_version = os.getenv("SHOPIFY_API_VERSION", "2024-10")
    client_id = os.getenv("SHOPIFY_CLIENT_ID")
    client_secret = os.getenv("SHOPIFY_CLIENT_SECRET")

    if not shop_name or not client_id or not client_secret:
        raise ShopifyConfigError(
            "Missing SHOPIFY_SHOP_NAME, SHOPIFY_CLIENT_ID, or SHOPIFY_CLIENT_SECRET in the environment (.env)."
        )

    return shop_name, api_version, client_id, client_secret


def _get_access_token() -> str:
    shop_name, _, client_id, client_secret = _get_config()

    if _token_cache["token"] and time.time() < _token_cache["expires_at"] - 60:
        return _token_cache["token"]

    response = requests.post(
        f"https://{shop_name}.myshopify.com/admin/oauth/access_token",
        headers={"Content-Type": "application/x-www-form-urlencoded"},
        data={
            "grant_type": "client_credentials",
            "client_id": client_id,
            "client_secret": client_secret,
        },
        timeout=30,
    )
    if response.status_code != 200:
        raise RuntimeError(
            f"Failed to obtain Shopify access token ({response.status_code}): {response.text}"
        )

    data = response.json()
    _token_cache["token"] = data["access_token"]
    _token_cache["expires_at"] = time.time() + data.get("expires_in", 0)
    return _token_cache["token"]


def _graphql_url() -> str:
    shop_name, api_version, _, _ = _get_config()
    return f"https://{shop_name}.myshopify.com/admin/api/{api_version}/graphql.json"


def _headers() -> Dict[str, str]:
    return {
        "Content-Type": "application/json",
        "X-Shopify-Access-Token": _get_access_token(),
    }


def _graphql(query: str, variables: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
    response = requests.post(
        _graphql_url(),
        headers=_headers(),
        json={"query": query, "variables": variables or {}},
        timeout=60,
    )

    if response.status_code != 200:
        raise RuntimeError(f"Shopify GraphQL request failed ({response.status_code}): {response.text}")

    data = response.json()
    if data.get("errors"):
        raise RuntimeError(f"Shopify GraphQL error: {data['errors']}")

    return data["data"]


def _check_user_errors(user_errors: List[Dict[str, str]], action: str) -> None:
    if user_errors:
        messages = "; ".join(f"{e.get('field')}: {e.get('message')}" for e in user_errors)
        raise RuntimeError(f"Shopify rejected {action}: {messages}")


def _numeric_id(gid: str) -> str:
    return gid.rsplit("/", 1)[-1]


PRODUCT_CREATE_MUTATION = """
mutation productCreate($input: ProductInput!) {
  productCreate(input: $input) {
    product {
      id
      title
      variants(first: 250) {
        edges { node { id inventoryItem { id } selectedOptions { name value } } }
      }
    }
    userErrors { field message }
  }
}
"""

VARIANTS_BULK_UPDATE_MUTATION = """
mutation productVariantsBulkUpdate($productId: ID!, $variants: [ProductVariantsBulkInput!]!) {
  productVariantsBulkUpdate(productId: $productId, variants: $variants) {
    productVariants { id price }
    userErrors { field message }
  }
}
"""

LOCATIONS_QUERY = """
query { locations(first: 1) { edges { node { id } } } }
"""

INVENTORY_SET_QUANTITY_MUTATION = """
mutation inventorySetOnHandQuantities($input: InventorySetOnHandQuantitiesInput!, $idempotencyKey: String!) {
  inventorySetOnHandQuantities(input: $input) @idempotent(key: $idempotencyKey) {
    userErrors { field message }
  }
}
"""

STAGED_UPLOADS_CREATE_MUTATION = """
mutation stagedUploadsCreate($input: [StagedUploadInput!]!) {
  stagedUploadsCreate(input: $input) {
    stagedTargets {
      url
      resourceUrl
      parameters { name value }
    }
    userErrors { field message }
  }
}
"""

PRODUCT_CREATE_MEDIA_MUTATION = """
mutation productCreateMedia($productId: ID!, $media: [CreateMediaInput!]!) {
  productCreateMedia(productId: $productId, media: $media) {
    media { alt mediaContentType status }
    mediaUserErrors { field message }
  }
}
"""

COLLECTIONS_QUERY = """
query($query: String!) {
  collections(first: 5, query: $query) {
    edges { node { id title } }
  }
}
"""

COLLECTION_ADD_PRODUCTS_MUTATION = """
mutation collectionAddProducts($id: ID!, $productIds: [ID!]!) {
  collectionAddProducts(id: $id, productIds: $productIds) {
    collection { id }
    userErrors { field message }
  }
}
"""


def _upload_images(product_gid: str, image_paths: List[str]) -> None:
    stage_inputs = []
    valid_paths = []

    for path in image_paths:
        p = Path(path)
        if not p.exists():
            continue
        mime_type = mimetypes.guess_type(p.name)[0] or "image/jpeg"
        stage_inputs.append({
            "filename": p.name,
            "mimeType": mime_type,
            "httpMethod": "POST",
            "resource": "PRODUCT_IMAGE",
            "fileSize": str(p.stat().st_size),
        })
        valid_paths.append(p)

    if not stage_inputs:
        return

    data = _graphql(STAGED_UPLOADS_CREATE_MUTATION, {"input": stage_inputs})
    result = data["stagedUploadsCreate"]
    _check_user_errors(result["userErrors"], "staged image upload")

    media_input = []
    for target, path in zip(result["stagedTargets"], valid_paths):
        form = {param["name"]: param["value"] for param in target["parameters"]}
        with open(path, "rb") as f:
            upload_response = requests.post(
                target["url"],
                data=form,
                files={"file": (path.name, f)},
                timeout=120,
            )
        if upload_response.status_code not in (200, 201, 204):
            raise RuntimeError(
                f"Failed to upload image {path.name} to Shopify staging ({upload_response.status_code})"
            )
        media_input.append({
            "originalSource": target["resourceUrl"],
            "alt": path.stem,
            "mediaContentType": "IMAGE",
        })

    data = _graphql(PRODUCT_CREATE_MEDIA_MUTATION, {"productId": product_gid, "media": media_input})
    _check_user_errors(data["productCreateMedia"]["mediaUserErrors"], "image attachment")


def _weight_unit_enum(unit: str) -> str:
    return {"kg": "KILOGRAMS", "g": "GRAMS", "lb": "POUNDS", "oz": "OUNCES"}.get(
        (unit or "kg").strip().lower(), "KILOGRAMS"
    )


def create_product(product_payload: Dict[str, Any], image_paths: List[str]) -> Dict[str, Any]:
    variants = product_payload.get("variants") or [{}]
    options = product_payload.get("options") or []

    product_input = {
        "title": product_payload["title"],
        "descriptionHtml": product_payload.get("body_html", ""),
        "vendor": product_payload.get("vendor"),
        "productType": product_payload.get("product_type"),
        "tags": [t.strip() for t in str(product_payload.get("tags", "")).split(",") if t.strip()],
        "status": str(product_payload.get("status", "draft")).upper(),
    }

    if options:
        product_input["productOptions"] = [
            {"name": opt["name"], "values": [{"name": v} for v in opt["values"]]}
            for opt in options
        ]

    data = _graphql(PRODUCT_CREATE_MUTATION, {"input": product_input})
    result = data["productCreate"]
    _check_user_errors(result["userErrors"], "product creation")
    product = result["product"]
    product_gid = product["id"]

    variant_nodes = [edge["node"] for edge in (product.get("variants", {}).get("edges") or [])]

    def _match_variant_node(variant_spec: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        if not options:
            return variant_nodes[0] if variant_nodes else None
        wanted = {opt["name"]: variant_spec.get(f"option{i + 1}") for i, opt in enumerate(options)}
        for node in variant_nodes:
            selected = {so["name"]: so["value"] for so in node.get("selectedOptions", [])}
            if selected == wanted:
                return node
        return None

    bulk_variant_inputs: List[Dict[str, Any]] = []
    quantity_updates: List[Tuple[str, int]] = []

    for variant_spec in variants:
        node = _match_variant_node(variant_spec)
        if not node:
            continue

        variant_input: Dict[str, Any] = {"id": node["id"]}
        if variant_spec.get("price") is not None:
            variant_input["price"] = str(variant_spec["price"])
        if variant_spec.get("barcode"):
            variant_input["barcode"] = variant_spec["barcode"]

        inventory_item: Dict[str, Any] = {}
        if variant_spec.get("sku"):
            inventory_item["sku"] = variant_spec["sku"]
            inventory_item["tracked"] = True
        if variant_spec.get("weight") is not None:
            inventory_item["measurement"] = {
                "weight": {
                    "value": float(variant_spec["weight"]),
                    "unit": _weight_unit_enum(variant_spec.get("weight_unit", "kg")),
                }
            }
        if variant_spec.get("requires_shipping") is not None:
            inventory_item["requiresShipping"] = variant_spec["requires_shipping"]
        if inventory_item:
            variant_input["inventoryItem"] = inventory_item

        if len(variant_input) > 1:
            bulk_variant_inputs.append(variant_input)

        quantity = variant_spec.get("inventory_quantity")
        if quantity is not None:
            quantity_updates.append((node["inventoryItem"]["id"], int(quantity)))

    if bulk_variant_inputs:
        data = _graphql(
            VARIANTS_BULK_UPDATE_MUTATION,
            {"productId": product_gid, "variants": bulk_variant_inputs},
        )
        _check_user_errors(data["productVariantsBulkUpdate"]["userErrors"], "variant update")

    if quantity_updates:
        locations = _graphql(LOCATIONS_QUERY)["locations"]["edges"]
        if locations:
            location_id = locations[0]["node"]["id"]
            inv_data = _graphql(
                INVENTORY_SET_QUANTITY_MUTATION,
                {
                    "input": {
                        "reason": "correction",
                        "setQuantities": [
                            {
                                "inventoryItemId": inventory_item_id,
                                "locationId": location_id,
                                "quantity": qty,
                                "changeFromQuantity": 0,
                            }
                            for inventory_item_id, qty in quantity_updates
                        ],
                    },
                    "idempotencyKey": str(uuid.uuid4()),
                },
            )
            _check_user_errors(inv_data["inventorySetOnHandQuantities"]["userErrors"], "inventory set")

    if image_paths:
        _upload_images(product_gid, image_paths)

    return {"id": _numeric_id(product_gid), "gid": product_gid, "title": product["title"]}


def _find_collection_gid(title: str) -> Optional[str]:
    data = _graphql(COLLECTIONS_QUERY, {"query": f"title:'{title}'"})
    for edge in data["collections"]["edges"]:
        if edge["node"]["title"].strip().lower() == title.strip().lower():
            return edge["node"]["id"]
    return None


def link_product_to_collections(product_id: Any, collection_titles: List[str]) -> Dict[str, List[str]]:
    """Attach a product to existing collections by title. Returns which titles were linked vs skipped."""
    product_gid = f"gid://shopify/Product/{product_id}"
    linked: List[str] = []
    skipped: List[str] = []

    for title in collection_titles:
        collection_gid = _find_collection_gid(title)
        if not collection_gid:
            skipped.append(title)
            continue

        data = _graphql(
            COLLECTION_ADD_PRODUCTS_MUTATION,
            {"id": collection_gid, "productIds": [product_gid]},
        )
        if data["collectionAddProducts"]["userErrors"]:
            skipped.append(title)
        else:
            linked.append(title)

    return {"linked": linked, "skipped": skipped}


def admin_product_url(product_id: Any) -> str:
    shop_name, _, _, _ = _get_config()
    return f"https://admin.shopify.com/store/{shop_name}/products/{product_id}"