#!/usr/bin/env python3
import csv
import json
import os
import re
import sys
import tempfile
import urllib.request
import zipfile
from pathlib import Path

from tableauhyperapi import Connection, HyperProcess, Telemetry


SERVER = os.environ.get("TABLEAU_SERVER", "https://mars.ushareit.org").rstrip("/")
API_VERSION = os.environ.get("TABLEAU_API_VERSION", "3.15")
SITE_CONTENT_URL = os.environ.get("TABLEAU_SITE_CONTENT_URL", "")
WORKBOOK_ID = os.environ.get("TABLEAU_WORKBOOK_ID", "07c673fb-bac8-4c78-83e6-c60f28e04e13")


def request_json(url, *, body=None, headers=None):
    request = urllib.request.Request(url, data=body, headers=headers or {})
    with urllib.request.urlopen(request, timeout=120) as response:
        return json.load(response)


def main():
    if len(sys.argv) != 3:
        raise SystemExit("usage: refresh_tableau.py YYYY-MM-DD OUTPUT.csv")
    target_date, output_path = sys.argv[1:]
    if not re.fullmatch(r"\d{4}-\d{2}-\d{2}", target_date):
        raise SystemExit("invalid target date")

    pat_name = os.environ.get("TABLEAU_PAT_NAME", "")
    pat_secret = os.environ.get("TABLEAU_PAT_SECRET", "")
    if not pat_name or not pat_secret:
        raise SystemExit("TABLEAU_PAT_NAME and TABLEAU_PAT_SECRET are required")

    signin_body = json.dumps({"credentials": {
        "personalAccessTokenName": pat_name,
        "personalAccessTokenSecret": pat_secret,
        "site": {"contentUrl": SITE_CONTENT_URL},
    }}).encode()
    signin = request_json(
        f"{SERVER}/api/{API_VERSION}/auth/signin",
        body=signin_body,
        headers={"Content-Type": "application/json", "Accept": "application/json"},
    )["credentials"]
    auth_token = signin["token"]
    site_id = signin["site"]["id"]

    try:
        with tempfile.TemporaryDirectory(prefix="manta-tableau-") as temp_dir:
            archive_path = Path(temp_dir) / "workbook.twbx"
            download = urllib.request.Request(
                f"{SERVER}/api/{API_VERSION}/sites/{site_id}/workbooks/{WORKBOOK_ID}/content",
                headers={"X-Tableau-Auth": auth_token},
            )
            with urllib.request.urlopen(download, timeout=180) as response:
                archive_path.write_bytes(response.read())
            with zipfile.ZipFile(archive_path) as archive:
                hyper_members = [name for name in archive.namelist() if name.endswith(".hyper")]
                if len(hyper_members) != 1:
                    raise SystemExit(f"expected one Hyper extract, found {len(hyper_members)}")
                archive.extract(hyper_members[0], temp_dir)
            hyper_path = Path(temp_dir) / hyper_members[0]

            query = f'''
                SELECT "pkg", "source", "dt", "campaign_name",
                       SUM("impression"), SUM("click"), SUM("conversions"), SUM("spend")
                FROM "Extract"."Extract"
                WHERE "dt" = '{target_date}'
                  AND LOWER("source") IN ('bigo', 'kwai', 'mintegral')
                GROUP BY "pkg", "source", "dt", "campaign_name"
                ORDER BY "pkg", "source", "campaign_name"
            '''
            with HyperProcess(Telemetry.DO_NOT_SEND_USAGE_DATA_TO_TABLEAU) as process:
                with Connection(process.endpoint, str(hyper_path)) as connection:
                    rows = connection.execute_list_query(query)
            if not rows:
                raise SystemExit(f"no Tableau rows for {target_date}")
            with open(output_path, "w", encoding="utf-8", newline="") as handle:
                writer = csv.writer(handle)
                writer.writerow(("pkg", "source", "dt", "campaign_name",
                                 "impression", "click", "conversions", "spend"))
                writer.writerows(rows)
            print(json.dumps({"date": target_date, "rawRows": len(rows)}))
    finally:
        try:
            signout = urllib.request.Request(
                f"{SERVER}/api/{API_VERSION}/auth/signout",
                data=b"",
                headers={"X-Tableau-Auth": auth_token},
            )
            urllib.request.urlopen(signout, timeout=30).close()
        except Exception:
            pass


if __name__ == "__main__":
    main()