#!/usr/bin/env python3
"""Create an Akamai DataStream2 stream and attach required behaviors to a property.

Standalone script (no IO River code imports).

Dependencies:
  pip install requests edgegrid-python

Example:
  python scripts/setup_akamai_datastream2.py \
    --property-name my-property 
    --host akab-xxxx.luna.akamaiapis.net \
    --client-token <client_token> \
    --client-secret <client_secret> \
    --access-token <access_token> \
    --service-uid <service_uid>
"""

from __future__ import annotations

import argparse
import json
import sys
from dataclasses import dataclass
from typing import Any
from urllib.parse import urljoin, urlparse

import requests
from akamai.edgegrid import EdgeGridAuth


DEFAULT_DATASET_FIELD_IDS = [
    # Step 2 Data Sets in the manual flow.
    1000,  # CP Code
    1100,  # Request ID
    1002,  # Request Time
    1023,  # Client IP
    1008,  # HTTP status code
    1009,  # Protocol type
    1011,  # Request host
    1012,  # Request method
    1013,  # Request path
    1015,  # Response Content-Length
    1017,  # User-Agent
    1032,  # Total bytes
    2012,  # Cookie
    1006,  # Turn around time
    2010,  # Referer
    2020,  # Asnum
    2008,  # Cache status
    1066,  # Breadcrumbs
    1102,  # Country/Region
    2022,  # City
    1082,  # Edge IP
    2014,  # CMCD
]


class AkamaiApiError(RuntimeError):
    """Raised when Akamai API returns an error."""


@dataclass
class PropertyRef:
    property_id: str
    property_name: str
    contract_id: str
    group_id: str


class AkamaiClient:
    def __init__(self, host: str, client_token: str, client_secret: str, access_token: str, timeout: int = 60):
        base = host if host.startswith("https://") else f"https://{host}"
        self.base_url = base.rstrip("/") + "/"
        self.timeout = timeout

        session = requests.Session()
        session.auth = EdgeGridAuth(
            client_token=client_token,
            client_secret=client_secret,
            access_token=access_token,
        )
        session.headers.update({"Accept": "application/json"})
        self.session = session

    def _request(self, method: str, path: str, *, params: dict[str, Any] | None = None,
                 body: dict[str, Any] | None = None) -> dict[str, Any]:
        url = urljoin(self.base_url, path.lstrip("/"))
        response = self.session.request(
            method=method,
            url=url,
            params=params,
            json=body,
            timeout=self.timeout,
        )

        if response.status_code >= 400:
            detail = response.text
            try:
                parsed = response.json()
                detail = json.dumps(parsed)
            except ValueError:
                pass
            raise AkamaiApiError(
                f"{method} {path} failed with {response.status_code}: {detail}"
            )

        if not response.content:
            return {}

        return response.json()

    def list_groups(self) -> list[dict[str, Any]]:
        payload = self._request("GET", "/papi/v1/groups")
        return payload.get("groups", {}).get("items", [])

    def list_properties(self, contract_id: str, group_id: str, property_name: str | None = None) -> list[dict[str, Any]]:
        params: dict[str, Any] = {
            "contractId": contract_id,
            "groupId": group_id,
        }
        if property_name:
            params["propertyName"] = property_name

        payload = self._request("GET", "/papi/v1/properties", params=params)
        return payload.get("properties", {}).get("items", [])

    def find_property_by_name(self, property_name: str,
                              group_id_filter: str | None = None,
                              contract_id_filter: str | None = None) -> PropertyRef:
        groups = self.list_groups()

        candidates: list[PropertyRef] = []
        for group in groups:
            group_id = group.get("groupId")
            if group_id_filter and group_id != group_id_filter:
                continue

            for contract_id in group.get("contractIds", []):
                if contract_id_filter and contract_id != contract_id_filter:
                    continue

                properties = self.list_properties(contract_id, group_id, property_name)
                for item in properties:
                    if item.get("propertyName") == property_name:
                        candidates.append(
                            PropertyRef(
                                property_id=item["propertyId"],
                                property_name=item["propertyName"],
                                contract_id=contract_id,
                                group_id=group_id,
                            )
                        )

        if not candidates:
            raise AkamaiApiError(
                f"Property '{property_name}' was not found in accessible groups/contracts."
            )

        if len(candidates) > 1:
            refs = [
                {
                    "propertyId": c.property_id,
                    "contractId": c.contract_id,
                    "groupId": c.group_id,
                }
                for c in candidates
            ]
            raise AkamaiApiError(
                "Property name is not unique. Re-run with --contract-id and/or --group-id. "
                f"Matches: {json.dumps(refs)}"
            )

        return candidates[0]

    def get_property(self, ref: PropertyRef) -> dict[str, Any]:
        payload = self._request(
            "GET",
            f"/papi/v1/properties/{ref.property_id}",
            params={"contractId": ref.contract_id, "groupId": ref.group_id},
        )
        items = payload.get("properties", {}).get("items", [])
        if not items:
            raise AkamaiApiError(f"Could not retrieve property details for {ref.property_id}.")
        return items[0]

    def get_latest_version(self, ref: PropertyRef, activated_on: str | None = None) -> dict[str, Any]:
        params = {
            "contractId": ref.contract_id,
            "groupId": ref.group_id,
        }
        if activated_on:
            params["activatedOn"] = activated_on

        payload = self._request(
            "GET",
            f"/papi/v1/properties/{ref.property_id}/versions/latest",
            params=params,
        )

        items = payload.get("versions", {}).get("items", [])
        if not items:
            raise AkamaiApiError(f"No versions found for property {ref.property_id}.")

        return items[0]

    def create_new_version(self, ref: PropertyRef, create_from_version: int, create_from_etag: str) -> int:
        payload = self._request(
            "POST",
            f"/papi/v1/properties/{ref.property_id}/versions",
            params={"contractId": ref.contract_id, "groupId": ref.group_id},
            body={
                "createFromVersion": create_from_version,
                "createFromVersionEtag": create_from_etag,
            },
        )

        version_link = payload.get("versionLink", "")
        if not version_link:
            latest = self.get_latest_version(ref)
            return int(latest["propertyVersion"])

        # versionLink may include query params; parse only the path segment.
        parsed_link = urlparse(version_link)
        path_parts = parsed_link.path.rstrip("/").split("/")
        new_version = int(path_parts[-1])
        return new_version

    def get_rules(self, ref: PropertyRef, version: int) -> dict[str, Any]:
        payload = self._request(
            "GET",
            f"/papi/v1/properties/{ref.property_id}/versions/{version}/rules",
            params={"contractId": ref.contract_id, "groupId": ref.group_id},
        )
        rules = payload.get("rules")
        if not rules:
            raise AkamaiApiError(f"Rules are missing for property {ref.property_id} version {version}.")
        return rules

    def update_rules(self, ref: PropertyRef, version: int, rules: dict[str, Any]) -> None:
        self._request(
            "PUT",
            f"/papi/v1/properties/{ref.property_id}/versions/{version}/rules",
            params={
                "contractId": ref.contract_id,
                "groupId": ref.group_id,
                "validateRules": "true",
                "validateMode": "full",
                "dryRun": "false",
            },
            body={"rules": rules},
        )

    def list_datastreams(self, group_id: str) -> list[dict[str, Any]]:
        group_id_number = int(group_id.replace("grp_", ""))
        payload = self._request(
            "GET",
            "/datastream-config-api/v2/log/streams",
            params={"groupId": group_id_number},
        )
        if isinstance(payload, list):
            return payload
        return payload.get("items", payload.get("streams", []))

    def create_datastream(self, payload: dict[str, Any]) -> int:
        response = self._request(
            "POST",
            "/datastream-config-api/v2/log/streams",
            params={"activate": "true"},
            body=payload,
        )
        if "streamId" not in response:
            raise AkamaiApiError(f"Unexpected DataStream create response: {json.dumps(response)}")
        return int(response["streamId"])


def ensure_behavior(behaviors: list[dict[str, Any]], name: str, options: dict[str, Any]) -> None:
    for behavior in behaviors:
        if behavior.get("name") == name:
            behavior["options"] = options
            return
    behaviors.append({"name": name, "options": options})


def make_datastream_behavior(stream_id: int, sampling_percentage: int) -> dict[str, Any]:
    stream_id_str = str(stream_id)
    return {
        "streamType": "LOG",
        "enabled": True,
        "datastreamIds": stream_id_str,
        "logEnabled": True,
        "logStreamName": [stream_id_str],
        "samplingPercentage": int(sampling_percentage),
        "collectMidgressTraffic": False,
    }


def make_report_behavior(cookie_mode: str, custom_log_field: str | None) -> dict[str, Any]:
    options: dict[str, Any] = {
        "logHost": True,
        "logReferer": True,
        "logUserAgent": True,
        "logCookies": cookie_mode,
        "logEdgeIP": True,
    }

    if custom_log_field is not None:
        options["logCustomLogField"] = True
        options["customLogField"] = custom_log_field

    return options


def normalize_cookie_mode(cookie_mode: str) -> str:
    # Akamai report behavior accepts OFF, ALL, SOME.
    aliases = {
        "NONE": "OFF",
        "ONLY_WHITELISTED": "SOME",
    }
    upper_mode = cookie_mode.upper()
    return aliases.get(upper_mode, upper_mode)


def build_datastream_payload(ref: PropertyRef, property_id_number: int, stream_name: str, endpoint: str,
                             notification_emails: list[str], dataset_field_ids: list[int]) -> dict[str, Any]:
    return {
        "streamName": stream_name,
        "groupId": int(ref.group_id.replace("grp_", "")),
        "contractId": ref.contract_id.replace("ctr_", ""),
        "properties": [{"propertyId": property_id_number}],
        "collectMidgress": False,
        "datasetFields": [{"datasetFieldId": field_id} for field_id in dataset_field_ids],
        "destination": {
            "destinationType": "HTTPS",
            "authenticationType": "NONE",
            "displayName": stream_name,
            "endpoint": endpoint,
            "contentType": "application/json",
            "compressLogs": True,
        },
        "deliveryConfiguration": {
            "format": "JSON",
            "fieldDelimiter": "SPACE",
            "frequency": {
                "intervalInSeconds": 60,
            },
        },
        "notificationEmails": notification_emails,
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Create DataStream2 and attach DataStream + report behaviors to an Akamai property."
        )
    )

    parser.add_argument("--property-name", required=True, help="Akamai property name")

    parser.add_argument("--host", required=True, help="Akamai API host, e.g. akab-xxxx.luna.akamaiapis.net")
    parser.add_argument("--client-token", required=True, help="Akamai client token")
    parser.add_argument("--client-secret", required=True, help="Akamai client secret")
    parser.add_argument("--access-token", required=True, help="Akamai access token")

    parser.add_argument("--group-id", help="Optional group filter, e.g. grp_12345")
    parser.add_argument("--contract-id", help="Optional contract filter, e.g. ctr_ABC")

    parser.add_argument("--service-uid", required=True, help="IO River service UID used in destination endpoint")
    parser.add_argument(
        "--endpoint-template",
        default="https://traffic-reporting.ioriver.io/public/akamai/{service_uid}/stats/akamai-{service_uid}_stats",
        help=(
            "Destination endpoint template. Uses {service_uid} placeholder."
        ),
    )
    parser.add_argument(
        "--stream-name",
        help="DataStream name. Default: <property-name>_stats",
    )
    parser.add_argument(
        "--datastream-sampling-percentage",
        type=int,
        default=10,
        help="Sampling percentage for DataStream behavior in property rules (default: 10)",
    )
    parser.add_argument(
        "--cookie-mode",
        default="ALL",
        choices=["OFF", "ALL", "SOME", "NONE", "ONLY_WHITELISTED"],
        help="Cookie mode for report behavior. Supported by Akamai: OFF, ALL, SOME. Aliases: NONE->OFF, ONLY_WHITELISTED->SOME.",
    )
    parser.add_argument(
        "--custom-log-field",
        default='{"source":"datastream2"}',
        help="Custom field value for report behavior (default: JSON marker)",
    )
    parser.add_argument(
        "--notification-email",
        action="append",
        default=[],
        help="Notification email for DataStream (repeatable).",
    )

    return parser.parse_args()


def main() -> int:
    args = parse_args()

    client = AkamaiClient(
        host=args.host,
        client_token=args.client_token,
        client_secret=args.client_secret,
        access_token=args.access_token,
    )

    ref = client.find_property_by_name(
        property_name=args.property_name,
        group_id_filter=args.group_id,
        contract_id_filter=args.contract_id,
    )

    property_details = client.get_property(ref)
    property_id_number = int(ref.property_id.replace("prp_", ""))

    stream_name = args.stream_name or f"{args.property_name}_stats"
    endpoint = args.endpoint_template.format(service_uid=args.service_uid)

    existing_stream = next(
        (s for s in client.list_datastreams(ref.group_id) if s.get("streamName") == stream_name),
        None,
    )

    if existing_stream:
        stream_id = int(existing_stream["streamId"])
    else:
        ds_payload = build_datastream_payload(
            ref=ref,
            property_id_number=property_id_number,
            stream_name=stream_name,
            endpoint=endpoint,
            notification_emails=args.notification_email,
            dataset_field_ids=DEFAULT_DATASET_FIELD_IDS,
        )
        stream_id = client.create_datastream(ds_payload)

    latest = client.get_latest_version(ref)
    latest_version = int(latest["propertyVersion"])
    latest_etag = latest.get("etag")
    if not latest_etag:
        raise AkamaiApiError("Latest property version did not include etag.")

    new_version = client.create_new_version(
        ref=ref,
        create_from_version=latest_version,
        create_from_etag=latest_etag,
    )

    rules = client.get_rules(ref, new_version)
    behaviors = rules.setdefault("behaviors", [])

    ensure_behavior(
        behaviors,
        "datastream",
        make_datastream_behavior(
            stream_id=stream_id,
            sampling_percentage=args.datastream_sampling_percentage,
        ),
    )
    ensure_behavior(
        behaviors,
        "report",
        make_report_behavior(
            cookie_mode=normalize_cookie_mode(args.cookie_mode),
            custom_log_field=args.custom_log_field,
        ),
    )

    client.update_rules(ref, new_version, rules)

    result = {
        "property": {
            "propertyId": ref.property_id,
            "propertyName": ref.property_name,
            "contractId": ref.contract_id,
            "groupId": ref.group_id,
            "latestVersionBeforeChange": latest_version,
            "newVersion": new_version,
            "propertyDetailsName": property_details.get("propertyName"),
        },
        "datastream": {
            "streamId": stream_id,
            "streamName": stream_name,
            "endpoint": endpoint,
            "created": existing_stream is None,
        },
        "behaviorsApplied": ["datastream", "report"],
        "activation": "Not activated by this script. Activate the new property version in Akamai if needed.",
    }
    print(json.dumps(result, indent=2))

    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except AkamaiApiError as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        raise SystemExit(1) from exc