Yandex Wordstat

# Environment with secrets
config/.env

# Cloud mode credentials
config/config.json
config/service_account_key.json

# Cache files (can be regenerated)
cache/regions.json
cache/iam_token.json

# macOS
.DS_Store

# =====================================================================
# Yandex Wordstat skill — credentials
# =====================================================================
# The skill supports two backends. See config/README.md for full setup.
#
# Backend auto-selection (priority):
#   1. YANDEX_WORDSTAT_BACKEND=legacy|cloud  (explicit override)
#   2. config/config.json + service_account_key.json present  → cloud
#   3. YANDEX_WORDSTAT_TOKEN set                              → legacy
#   4. Nothing                                                → error
# Cloud wins on tie. Pin legacy explicitly with the override below.

# --- Legacy mode (DEPRECATED — Yandex no longer onboards new users) ---
# OAuth token from https://oauth.yandex.ru/authorize?response_type=token&client_id=...
# Get it via: bash scripts/get_token.sh --client-id YOUR_CLIENT_ID
# Tokens expire after 1 year.
YANDEX_WORDSTAT_TOKEN=your_token_here

# Optional: pin backend explicitly even when both credential sets are present.
# Values: legacy | cloud
# YANDEX_WORDSTAT_BACKEND=legacy

# --- Cloud mode ---
# For cloud setup, copy config/config.example.json → config/config.json
# and place your service account key file at config/service_account_key.json
# (or any path you set in auth.service_account_key_file).
# See config/README.md → "Cloud mode" for the full 6-step flow.

{
  "yandex_cloud_folder_id": "b1g...",
  "auth": {
    "service_account_key_file": "config/service_account_key.json",
    "openssl_bin": "openssl"
  }
}

#!/bin/sh
# Common functions for Yandex Wordstat skill — dual backend (legacy + cloud)
#
# Public API (sourced by other scripts):
#   load_config            — picks backend, exports WORDSTAT_BACKEND, _DETECTED_VIA, _CLOUD_*
#   wordstat_request M P   — request to Wordstat API, always returns LEGACY-shaped JSON
#   print_backend_info     — backend-aware diagnostic block (used by quota.sh)
#   die_with_help MSG      — structured error pointing user at config README
#   json_escape, format_number, json_value, json_string  — legacy helpers (unchanged)
#
# Backend dispatch:
#   - WORDSTAT_BACKEND=legacy → POST api.wordstat.yandex.net/v1/{method} (Bearer OAuth)
#   - WORDSTAT_BACKEND=cloud  → POST searchapi.api.cloud.yandex.net/v2/wordstat/{method}
#                              with IAM Bearer + folderId, response normalized back to legacy
#                              shape so existing parsers in callers don't change.
#
# Selection in load_config is STRUCTURAL ONLY — no network, no IAM preflight.
# IAM/network errors surface on the first wordstat_request call.

# Resolve directories. Use $0 because we're sourced from many shells (sh + bash).
# Tests can pre-set WORDSTAT_SCRIPT_DIR / WORDSTAT_SKILL_DIR / WORDSTAT_CONFIG_DIR
# to override the auto-resolution (POSIX sh has no portable way to get the path
# of a sourced script when $0 isn't reliable).
if [ -z "${WORDSTAT_SCRIPT_DIR:-}" ]; then
    WORDSTAT_SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
fi
if [ -z "${WORDSTAT_SKILL_DIR:-}" ]; then
    WORDSTAT_SKILL_DIR="$(cd "$WORDSTAT_SCRIPT_DIR/.." && pwd)"
fi
WORDSTAT_CONFIG_DIR="${WORDSTAT_CONFIG_DIR:-$WORDSTAT_SKILL_DIR/config}"
WORDSTAT_CACHE_DIR="${WORDSTAT_CACHE_DIR:-$WORDSTAT_SKILL_DIR/cache}"

WORDSTAT_LEGACY_API="https://api.wordstat.yandex.net/v1"
WORDSTAT_CLOUD_API="https://searchapi.api.cloud.yandex.net/v2/wordstat"
WORDSTAT_IAM_API="https://iam.api.cloud.yandex.net/iam/v1/tokens"
WORDSTAT_README_URL="https://github.com/artwist-polyakov/polyakov-claude-skills/blob/main/plugins/yandex-wordstat/skills/yandex-wordstat/config/README.md"

# Exported by load_config so callers and die_with_help can read them
WORDSTAT_BACKEND=""
WORDSTAT_BACKEND_DETECTED_VIA=""
WORDSTAT_CLOUD_FOLDER_ID=""
WORDSTAT_CLOUD_SA_KEY_PATH=""
WORDSTAT_CLOUD_OPENSSL_BIN=""

# ---------------------------------------------------------------------
# Error helper
# ---------------------------------------------------------------------

die_with_help() {
    _msg="$1"
    _extra="${2:-}"

    {
        printf '[wordstat] %s\n' "$_msg"
        if [ -n "$WORDSTAT_BACKEND" ]; then
            printf 'Backend: %s' "$WORDSTAT_BACKEND"
            [ -n "$WORDSTAT_BACKEND_DETECTED_VIA" ] && \
                printf ' (%s)' "$WORDSTAT_BACKEND_DETECTED_VIA"
            printf '\n'
        fi
        [ -n "$_extra" ] && printf '%s\n' "$_extra"
        printf '\n'
        printf 'Likely the plugin config needs updating. See:\n'
        printf '  %s\n\n' "$WORDSTAT_README_URL"
        printf 'Quick checks:\n'
        printf '  - cloud mode:  config/config.json has yandex_cloud_folder_id?\n'
        if [ -n "$WORDSTAT_CLOUD_SA_KEY_PATH" ]; then
            printf '                 SA key file: %s\n' "$WORDSTAT_CLOUD_SA_KEY_PATH"
            printf '                 (resolved from auth.service_account_key_file) — present and readable?\n'
        else
            printf '                 SA key file from auth.service_account_key_file — present and readable?\n'
        fi
        printf "                 SA has role 'search-api.webSearch.user'?\n"
        printf '  - legacy mode: YANDEX_WORDSTAT_TOKEN still valid? (tokens expire after 1 year)\n'
        printf '  - to switch:   set YANDEX_WORDSTAT_BACKEND=legacy|cloud in config/.env\n'
    } >&2
    exit 1
}

# ---------------------------------------------------------------------
# Legacy helpers (kept for compatibility with bash callers)
# ---------------------------------------------------------------------

json_escape() {
    printf '%s' "$1" | sed 's/\\/\\\\/g; s/"/\\"/g; s/	/\\t/g'
}

format_number() {
    printf "%'d" "$1" 2>/dev/null || echo "$1"
}

# Extract a numeric/literal JSON value (no string quoting)
json_value() {
    _jv_json="$1"; _jv_key="$2"
    printf '%s' "$_jv_json" | grep -o "\"$_jv_key\":[^,}]*" | head -1 | sed 's/.*://' | tr -d '"[:space:]'
}

# Extract a JSON string value
json_string() {
    _js_json="$1"; _js_key="$2"
    printf '%s' "$_js_json" | grep -o "\"$_js_key\":\"[^\"]*\"" | head -1 | sed 's/.*:"//' | tr -d '"'
}

# ---------------------------------------------------------------------
# Backend selection — load_config
# ---------------------------------------------------------------------

# Read .env if present (legacy creds + override). Sourced into current shell.
_load_env_file() {
    _env_file="$WORDSTAT_CONFIG_DIR/.env"
    if [ -f "$_env_file" ]; then
        # shellcheck disable=SC1090
        . "$_env_file"
    fi
}

# Read a value from config.json. Usage: _cfg_get "key" or "auth.openssl_bin"
# Returns empty string on missing key, missing file, or parse error.
_cfg_get() {
    _cfg_file="$WORDSTAT_CONFIG_DIR/config.json"
    [ -f "$_cfg_file" ] || { echo ""; return 0; }
    _CFG_FILE="$_cfg_file" _CFG_KEY="$1" python3 - <<'PYEOF' 2>/dev/null
import json, os, sys
try:
    with open(os.environ["_CFG_FILE"]) as f:
        cfg = json.load(f)
except Exception:
    print("")
    sys.exit(0)
v = cfg
for part in os.environ["_CFG_KEY"].split("."):
    if isinstance(v, dict) and part in v:
        v = v[part]
    else:
        v = None
        break
print("" if v is None else v)
PYEOF
}

# Resolve a path: absolute as-is, relative resolved against skill dir
_resolve_path() {
    _rp="$1"
    case "$_rp" in
        /*) printf '%s\n' "$_rp" ;;
        *)  printf '%s/%s\n' "$WORDSTAT_SKILL_DIR" "$_rp" ;;
    esac
}

# Detect cloud structural config. Sets WORDSTAT_CLOUD_* variables on success.
# Returns 0 if cloud is structurally configured, 1 if not, 2 if config.json is
# present but malformed (caller should die loudly).
_detect_cloud_config() {
    _cfg_file="$WORDSTAT_CONFIG_DIR/config.json"
    [ -f "$_cfg_file" ] || return 1

    _folder=$(_cfg_get yandex_cloud_folder_id)
    _sa_rel=$(_cfg_get auth.service_account_key_file)
    _ossl=$(_cfg_get auth.openssl_bin)

    if [ -z "$_folder" ]; then
        WORDSTAT_BACKEND_DETECTED_VIA="cloud (config.json present but yandex_cloud_folder_id missing)"
        return 2
    fi
    if [ -z "$_sa_rel" ]; then
        WORDSTAT_BACKEND_DETECTED_VIA="cloud (config.json present but auth.service_account_key_file missing)"
        return 2
    fi

    _sa_resolved=$(_resolve_path "$_sa_rel")
    if [ ! -r "$_sa_resolved" ]; then
        WORDSTAT_CLOUD_SA_KEY_PATH="$_sa_resolved"
        WORDSTAT_BACKEND_DETECTED_VIA="cloud (SA key file not found at resolved path)"
        return 2
    fi

    WORDSTAT_CLOUD_FOLDER_ID="$_folder"
    WORDSTAT_CLOUD_SA_KEY_PATH="$_sa_resolved"
    WORDSTAT_CLOUD_OPENSSL_BIN="${_ossl:-openssl}"
    return 0
}

load_config() {
    _load_env_file

    # 1. Explicit override
    if [ -n "${YANDEX_WORDSTAT_BACKEND:-}" ]; then
        case "$YANDEX_WORDSTAT_BACKEND" in
            cloud)
                _rc=0
                _detect_cloud_config || _rc=$?
                if [ "$_rc" = "2" ]; then
                    WORDSTAT_BACKEND="cloud"
                    die_with_help "YANDEX_WORDSTAT_BACKEND=cloud but config is incomplete: $WORDSTAT_BACKEND_DETECTED_VIA"
                fi
                if [ "$_rc" = "1" ]; then
                    WORDSTAT_BACKEND="cloud"
                    die_with_help "YANDEX_WORDSTAT_BACKEND=cloud but config/config.json is missing"
                fi
                WORDSTAT_BACKEND="cloud"
                WORDSTAT_BACKEND_DETECTED_VIA="explicit override"
                return 0
                ;;
            legacy)
                if [ -z "${YANDEX_WORDSTAT_TOKEN:-}" ]; then
                    WORDSTAT_BACKEND="legacy"
                    WORDSTAT_BACKEND_DETECTED_VIA="explicit override"
                    die_with_help "YANDEX_WORDSTAT_BACKEND=legacy but YANDEX_WORDSTAT_TOKEN is not set"
                fi
                WORDSTAT_BACKEND="legacy"
                WORDSTAT_BACKEND_DETECTED_VIA="explicit override"
                return 0
                ;;
            *)
                die_with_help "Invalid YANDEX_WORDSTAT_BACKEND='$YANDEX_WORDSTAT_BACKEND' (expected 'legacy' or 'cloud')"
                ;;
        esac
    fi

    # 2. Cloud structurally configured → cloud (cloud wins on tie)
    _rc=0
    _detect_cloud_config || _rc=$?
    if [ "$_rc" = "0" ]; then
        WORDSTAT_BACKEND="cloud"
        WORDSTAT_BACKEND_DETECTED_VIA="auto: config.json present"
        return 0
    fi
    if [ "$_rc" = "2" ]; then
        # Malformed cloud config → fail loudly, do NOT silently fall back
        WORDSTAT_BACKEND="cloud"
        die_with_help "config/config.json present but invalid: $WORDSTAT_BACKEND_DETECTED_VIA"
    fi

    # 3. Legacy creds present → legacy
    if [ -n "${YANDEX_WORDSTAT_TOKEN:-}" ]; then
        WORDSTAT_BACKEND="legacy"
        WORDSTAT_BACKEND_DETECTED_VIA="auto: YANDEX_WORDSTAT_TOKEN set"
        return 0
    fi

    # 4. Nothing
    die_with_help "No Wordstat credentials found"
}

# ---------------------------------------------------------------------
# Backend info — used by quota.sh
# ---------------------------------------------------------------------

print_backend_info() {
    case "$WORDSTAT_BACKEND" in
        legacy)
            echo "Backend: legacy ($WORDSTAT_BACKEND_DETECTED_VIA)"
            echo ""
            echo "=== Endpoints ==="
            echo "  POST $WORDSTAT_LEGACY_API/topRequests"
            echo "  POST $WORDSTAT_LEGACY_API/dynamics"
            echo "  POST $WORDSTAT_LEGACY_API/regions"
            echo ""
            echo "=== API Limits ==="
            echo "  - Rate limit: 10 requests/second"
            echo "  - Daily quota: 1000 requests"
            echo ""
            echo "Note: This API is deprecated for new users. Existing tokens still work."
            ;;
        cloud)
            echo "Backend: cloud ($WORDSTAT_BACKEND_DETECTED_VIA)"
            echo "  folder_id: $WORDSTAT_CLOUD_FOLDER_ID"
            echo "  SA key:    $WORDSTAT_CLOUD_SA_KEY_PATH"
            echo ""
            echo "=== Endpoints ==="
            echo "  POST $WORDSTAT_CLOUD_API/topRequests"
            echo "  POST $WORDSTAT_CLOUD_API/dynamics"
            echo "  POST $WORDSTAT_CLOUD_API/regions"
            echo ""
            echo "=== API Limits ==="
            echo "  Wordstat in Search API is currently in Preview."
            echo "  See https://yandex.cloud/ru/docs/search-api/pricing for current limits and billing."
            ;;
        *)
            echo "Backend: (not configured)"
            ;;
    esac
}

# ---------------------------------------------------------------------
# IAM token — JWT PS256 with SA key (inline-copied from yandex-search-api)
# ---------------------------------------------------------------------

_make_secure_tmpdir() {
    _old_umask=$(umask)
    umask 077
    _td=$(mktemp -d "${TMPDIR:-/tmp}/wordstat_XXXXXX")
    umask "$_old_umask"
    echo "$_td"
}

_check_openssl() {
    _ossl="$1"
    if ! command -v "$_ossl" >/dev/null 2>&1; then
        die_with_help "openssl not found at '$_ossl'" \
            "Install OpenSSL 1.1.1+ or set auth.openssl_bin in config/config.json"
    fi
    _ossl_ver=$("$_ossl" version 2>/dev/null || true)
    case "$_ossl_ver" in
        LibreSSL*)
            die_with_help "LibreSSL detected ($_ossl_ver) — OpenSSL 1.1.1+ required for PS256" \
                "macOS users: brew install openssl@3 and set auth.openssl_bin to the homebrew path"
            ;;
        "OpenSSL 0."*|"OpenSSL 1.0."*)
            die_with_help "OpenSSL too old ($_ossl_ver), need 1.1.1+"
            ;;
    esac
}

_get_cached_iam_token() {
    _cf="$WORDSTAT_CACHE_DIR/iam_token.json"
    [ -f "$_cf" ] || return 0
    _CACHE_FILE="$_cf" python3 - <<'PYEOF' 2>/dev/null
import json, os, time
cf = os.environ["_CACHE_FILE"]
try:
    with open(cf) as f:
        d = json.load(f)
    exp = d.get("expires_at", 0)
    if exp - time.time() > 300:
        print(d["iam_token"])
except Exception:
    pass
PYEOF
}

_save_iam_token() {
    _tok="$1"
    _exp="$2"
    mkdir -p "$WORDSTAT_CACHE_DIR"
    _cf="$WORDSTAT_CACHE_DIR/iam_token.json"
    _old_umask=$(umask)
    umask 077
    _tmp="$WORDSTAT_CACHE_DIR/.iam_token_tmp_$$.json"
    _SAVE_TOKEN="$_tok" _SAVE_EXP="$_exp" _TMP_FILE="$_tmp" python3 - <<'PYEOF'
import json, os
d = {"iam_token": os.environ["_SAVE_TOKEN"], "expires_at": int(os.environ["_SAVE_EXP"])}
with open(os.environ["_TMP_FILE"], "w") as f:
    json.dump(d, f)
PYEOF
    mv "$_tmp" "$_cf"
    umask "$_old_umask"
}

# Issue a fresh IAM token from the SA key. Echoes token on stdout.
_iam_token_issue() {
    _check_openssl "$WORDSTAT_CLOUD_OPENSSL_BIN"

    if [ ! -r "$WORDSTAT_CLOUD_SA_KEY_PATH" ]; then
        die_with_help "Service account key file not readable: $WORDSTAT_CLOUD_SA_KEY_PATH"
    fi

    _tmp=$(_make_secure_tmpdir)
    # shellcheck disable=SC2064
    trap "rm -rf '$_tmp'" EXIT INT TERM

    # Build JWT header + payload, write key.pem and signing_input.txt
    _SA_KEY="$WORDSTAT_CLOUD_SA_KEY_PATH" _TMP="$_tmp" python3 - <<'PYEOF' || die_with_help "Failed to build JWT from SA key"
import json, base64, time, os, sys
sa_key_file = os.environ["_SA_KEY"]
tmp_dir = os.environ["_TMP"]
try:
    with open(sa_key_file) as f:
        sa = json.load(f)
    sa_id = sa["service_account_id"]
    key_id = sa["id"]
    private_key = sa["private_key"]
except Exception as e:
    print(f"SA key parse error: {e}", file=sys.stderr)
    sys.exit(1)

with open(os.path.join(tmp_dir, "key.pem"), "w") as f:
    f.write(private_key)

header = json.dumps({"typ": "JWT", "alg": "PS256", "kid": key_id}, separators=(",", ":"))
header_b64 = base64.urlsafe_b64encode(header.encode()).rstrip(b"=").decode()

now = int(time.time())
payload = json.dumps({
    "iss": sa_id,
    "aud": "https://iam.api.cloud.yandex.net/iam/v1/tokens",
    "iat": now,
    "exp": now + 3600,
}, separators=(",", ":"))
payload_b64 = base64.urlsafe_b64encode(payload.encode()).rstrip(b"=").decode()

signing_input = f"{header_b64}.{payload_b64}"
with open(os.path.join(tmp_dir, "signing_input.txt"), "w") as f:
    f.write(signing_input)
with open(os.path.join(tmp_dir, "header_payload.txt"), "w") as f:
    f.write(signing_input)
PYEOF

    "$WORDSTAT_CLOUD_OPENSSL_BIN" dgst -sha256 \
        -sigopt rsa_padding_mode:pss \
        -sigopt rsa_pss_saltlen:-1 \
        -sign "$_tmp/key.pem" \
        -out "$_tmp/signature.bin" \
        "$_tmp/signing_input.txt" 2>/dev/null \
        || die_with_help "openssl PS256 signing failed"

    _sig=$(python3 -c "
import base64, sys
with open('$_tmp/signature.bin', 'rb') as f:
    print(base64.urlsafe_b64encode(f.read()).rstrip(b'=').decode())
")
    _hp=$(cat "$_tmp/header_payload.txt")
    _jwt="${_hp}.${_sig}"

    _resp=$(curl -s -X POST "$WORDSTAT_IAM_API" \
        -H "Content-Type: application/json" \
        -d "{\"jwt\":\"$_jwt\"}")

    if [ -z "$_resp" ]; then
        die_with_help "Empty response from IAM API"
    fi

    # Parse token + expiry
    _result=$(printf '%s' "$_resp" | python3 -c "
import json, sys
from datetime import datetime
try:
    d = json.load(sys.stdin)
except Exception as e:
    print('PARSE_ERROR:' + str(e))
    sys.exit(0)
tok = d.get('iamToken', '')
exp_s = d.get('expiresAt', '')
if not tok:
    print('NO_TOKEN:' + json.dumps(d)[:300])
    sys.exit(0)
if exp_s:
    try:
        ts = datetime.fromisoformat(exp_s.replace('Z', '+00:00')).timestamp()
        exp = int(ts)
    except Exception:
        import time
        exp = int(time.time()) + 43200
else:
    import time
    exp = int(time.time()) + 43200
print(f'{tok}|{exp}')
")
    case "$_result" in
        PARSE_ERROR:*) die_with_help "IAM response parse error: ${_result#PARSE_ERROR:}" "Raw: $_resp" ;;
        NO_TOKEN:*)    die_with_help "IAM response missing iamToken" "${_result#NO_TOKEN:}" ;;
    esac

    _tok=$(printf '%s' "$_result" | cut -d'|' -f1)
    _exp=$(printf '%s' "$_result" | cut -d'|' -f2)
    _save_iam_token "$_tok" "$_exp"

    rm -rf "$_tmp"
    trap - EXIT INT TERM
    printf '%s' "$_tok"
}

_iam_token_get() {
    _cached=$(_get_cached_iam_token)
    if [ -n "$_cached" ]; then
        printf '%s' "$_cached"
        return 0
    fi
    _iam_token_issue
}

# ---------------------------------------------------------------------
# Request translation + response normalization (cloud ↔ legacy)
# ---------------------------------------------------------------------

# Translate legacy-shape params JSON → cloud request body JSON.
# Args: $1 = method (topRequests|dynamics|regions), $2 = legacy params JSON
# Output: cloud-shape JSON on stdout.
# Exits 1 with die_with_help on dynamics preflight failure.
_xlate_request() {
    _method="$1"
    _params="$2"
    _METHOD="$_method" _PARAMS="$_params" _FOLDER="$WORDSTAT_CLOUD_FOLDER_ID" \
    python3 - <<'PYEOF'
import json, os, re, sys

method = os.environ["_METHOD"]
params = json.loads(os.environ["_PARAMS"])
folder = os.environ["_FOLDER"]

DEVICE_MAP = {
    "all": "DEVICE_ALL",
    "desktop": "DEVICE_DESKTOP",
    "phone": "DEVICE_PHONE",
    "tablet": "DEVICE_TABLET",
}
PERIOD_MAP = {
    "monthly": "PERIOD_MONTHLY",
    "weekly": "PERIOD_WEEKLY",
    "daily": "PERIOD_DAILY",
}
REGION_TYPE_MAP = {
    "all": "REGION_ALL",
    "cities": "REGION_CITIES",
    "regions": "REGION_REGIONS",
}

def map_devices(d):
    if d is None:
        return ["DEVICE_ALL"]
    if isinstance(d, list):
        return [DEVICE_MAP.get(x, x) if isinstance(x, str) and not x.startswith("DEVICE_") else x for x in d]
    return [DEVICE_MAP.get(d, "DEVICE_ALL")]

def map_regions(r):
    if r is None:
        return None
    return [str(x) for x in r]

def to_rfc3339(d):
    # Accept either YYYY-MM-DD or already-RFC3339
    if not d:
        return d
    if "T" in d:
        return d
    return d + "T00:00:00Z"

if method == "topRequests":
    body = {"phrase": params["phrase"]}
    if "numPhrases" in params:
        body["numPhrases"] = str(params["numPhrases"])
    if "regions" in params:
        body["regions"] = map_regions(params["regions"])
    if "devices" in params:
        body["devices"] = map_devices(params["devices"])
    body["folderId"] = folder

elif method == "dynamics":
    # ---- Preflight: cloud only allows '+' operator at weekly/monthly ----
    period = params.get("period", "monthly")
    phrase = params.get("phrase", "")
    if period != "daily":
        # Token-boundary detection of operators that cloud rejects at non-daily.
        # Hyphen inside word (санкт-петербург, премиум-класс) MUST pass.
        # Token-leading -, !, or any of " ( | ) trigger the failure. + is allowed.
        bad_ops = []
        if re.search(r'(^|\s)-\S', phrase):
            bad_ops.append("- (minus-word)")
        if re.search(r'(^|\s)!\S', phrase):
            bad_ops.append("! (exact form)")
        if '"' in phrase:
            bad_ops.append('" (exact phrase)')
        if "(" in phrase or ")" in phrase or "|" in phrase:
            bad_ops.append("( | ) (grouping)")
        if bad_ops:
            print("PREFLIGHT_FAIL:" + ", ".join(bad_ops), file=sys.stderr)
            sys.exit(2)

    body = {
        "phrase": phrase,
        "period": PERIOD_MAP.get(period, period),
        "fromDate": to_rfc3339(params["fromDate"]),
    }
    if "toDate" in params and params["toDate"]:
        body["toDate"] = to_rfc3339(params["toDate"])
    if "regions" in params:
        body["regions"] = map_regions(params["regions"])
    if "devices" in params:
        body["devices"] = map_devices(params["devices"])
    body["folderId"] = folder

elif method == "regions":
    body = {"phrase": params["phrase"]}
    if "regionType" in params:
        body["region"] = REGION_TYPE_MAP.get(params["regionType"], params["regionType"])
    if "devices" in params:
        body["devices"] = map_devices(params["devices"])
    body["folderId"] = folder

else:
    print(f"UNKNOWN_METHOD:{method}", file=sys.stderr)
    sys.exit(2)

print(json.dumps(body, ensure_ascii=False))
PYEOF
}

# Normalize cloud response JSON → legacy shape JSON.
# Args: $1 = method, $2 = path to cloud response file (optional; if missing, spool stdin)
# Output: legacy-shape JSON on stdout
#
# Implementation note: cloud responses for topRequests --limit 2000 can be
# multi-MB. Passing through env var is unsafe (ARG_MAX / E2BIG). We use a file
# path. If the caller already has the response in a file (e.g. _cloud_request),
# pass it as $2 to skip the spool step.
_normalize_response() {
    _method="$1"
    _nr_owns_tmp=0
    if [ -n "${2:-}" ]; then
        _nr_tmp="$2"
    else
        _nr_tmp="${TMPDIR:-/tmp}/wordstat_norm_$$.json"
        cat > "$_nr_tmp"
        _nr_owns_tmp=1
    fi
    _METHOD="$_method" _RESP_FILE="$_nr_tmp" python3 - <<'PYEOF'
import json, os, sys
method = os.environ["_METHOD"]
try:
    with open(os.environ["_RESP_FILE"], "r", encoding="utf-8") as f:
        d = json.load(f)
except Exception as e:
    print(json.dumps({"error": f"Cloud response parse error: {e}"}))
    sys.exit(0)

# Translate cloud error JSON to legacy {"error": ...}
if "code" in d and "message" in d and "results" not in d and "topRequests" not in d:
    print(json.dumps({"error": d.get("message", "cloud error"), "code": d.get("code")}))
    sys.exit(0)

def to_int(v):
    if v is None:
        return 0
    try:
        return int(v)
    except (TypeError, ValueError):
        return v

if method == "topRequests":
    out = {}
    if "totalCount" in d:
        out["totalCount"] = to_int(d["totalCount"])
    out["topRequests"] = [
        {"phrase": r.get("phrase", ""), "count": to_int(r.get("count", 0))}
        for r in d.get("results", [])
    ]
    out["associations"] = [
        {"phrase": r.get("phrase", ""), "count": to_int(r.get("count", 0))}
        for r in d.get("associations", [])
    ]

elif method == "dynamics":
    out = {
        "data": [
            {
                "date": r.get("date", ""),
                "count": to_int(r.get("count", 0)),
                "share": r.get("share", 0),
            }
            for r in d.get("results", [])
        ]
    }

elif method == "regions":
    out = {
        "regions": [
            {
                "regionId": to_int(r.get("region", 0)),
                "count": to_int(r.get("count", 0)),
                "share": r.get("share", 0),
                "affinity": to_int(r.get("affinityIndex", r.get("affinity", 0))),
            }
            for r in d.get("results", [])
        ]
    }
else:
    out = d

# Compact separators — no spaces. Matches the legacy API JSON shape that
# existing grep/sed parsers in top_requests.sh, dynamics.sh, regions_stats.sh expect.
# E.g. "topRequests":[{"phrase":"...","count":123}] not "topRequests": [{"phrase": "...", "count": 123}]
print(json.dumps(out, ensure_ascii=False, separators=(",", ":")))
PYEOF
    [ "$_nr_owns_tmp" = "1" ] && rm -f "$_nr_tmp"
    return 0
}

# ---------------------------------------------------------------------
# wordstat_request — public dispatcher
# ---------------------------------------------------------------------

# Legacy backend: direct curl to api.wordstat.yandex.net/v1
_legacy_request() {
    _method="$1"
    _params="$2"
    curl -s -X POST "$WORDSTAT_LEGACY_API/$_method" \
        -H "Authorization: Bearer $YANDEX_WORDSTAT_TOKEN" \
        -H "Content-Type: application/json; charset=utf-8" \
        -H "Accept-Language: ru" \
        -d "$_params"
}

# Cloud backend: translate, sign, POST, normalize
_cloud_request() {
    _method="$1"
    _params="$2"

    # 1. Translate request
    _xlate_out=$(_xlate_request "$_method" "$_params" 2>&1)
    _xlate_rc=$?
    if [ "$_xlate_rc" != "0" ]; then
        case "$_xlate_out" in
            *PREFLIGHT_FAIL:*)
                _ops=${_xlate_out#*PREFLIGHT_FAIL:}
                die_with_help \
                    "Cloud Wordstat dynamics: at weekly/monthly granularity, only '+' operator is allowed. Found: $_ops" \
                    "Either switch --period to daily, or remove these operators from --phrase. See: https://aistudio.yandex.ru/docs/ru/search-api/operations/wordstat-getdynamics.html"
                ;;
            *UNKNOWN_METHOD:*)
                die_with_help "Unknown wordstat method: ${_xlate_out#*UNKNOWN_METHOD:}"
                ;;
            *)
                die_with_help "Request translation failed" "$_xlate_out"
                ;;
        esac
    fi
    _cloud_body="$_xlate_out"

    # 2. Get IAM token (uses cache, falls back to issue)
    _tok=$(_iam_token_get)
    if [ -z "$_tok" ]; then
        die_with_help "Failed to obtain IAM token"
    fi

    # 3. POST with retry on 5xx and refresh on 401
    _attempt=0
    _max_attempts=3
    _backoff=2
    while [ "$_attempt" -lt "$_max_attempts" ]; do
        _attempt=$((_attempt + 1))
        _tmp=$(_make_secure_tmpdir)
        _resp_file="$_tmp/resp"
        _status=$(curl -s -o "$_resp_file" -w '%{http_code}' \
            -X POST "$WORDSTAT_CLOUD_API/$_method" \
            -H "Authorization: Bearer $_tok" \
            -H "Content-Type: application/json" \
            -d "$_cloud_body")

        case "$_status" in
            2[0-9][0-9])
                _normalize_response "$_method" "$_resp_file"
                rm -rf "$_tmp"
                return 0
                ;;
            401)
                # Refresh once and retry
                if [ "$_attempt" = "1" ]; then
                    rm -f "$WORDSTAT_CACHE_DIR/iam_token.json"
                    _tok=$(_iam_token_issue)
                    rm -rf "$_tmp"
                    continue
                fi
                _err=$(cat "$_resp_file" 2>/dev/null)
                rm -rf "$_tmp"
                die_with_help "Cloud Wordstat 401 Unauthorized after token refresh" "$_err"
                ;;
            403)
                _err=$(cat "$_resp_file" 2>/dev/null)
                rm -rf "$_tmp"
                die_with_help "Cloud Wordstat 403 Forbidden" \
                    "Check that your service account has the role 'search-api.webSearch.user' on folder $WORDSTAT_CLOUD_FOLDER_ID. Raw: $_err"
                ;;
            5[0-9][0-9]|000)
                if [ "$_attempt" -lt "$_max_attempts" ]; then
                    rm -rf "$_tmp"
                    sleep "$_backoff"
                    _backoff=$((_backoff * 2))
                    continue
                fi
                _err=$(cat "$_resp_file" 2>/dev/null)
                rm -rf "$_tmp"
                die_with_help "Cloud Wordstat $_status after $_max_attempts retries" "$_err"
                ;;
            *)
                _err=$(cat "$_resp_file" 2>/dev/null)
                rm -rf "$_tmp"
                die_with_help "Cloud Wordstat HTTP $_status" "$_err"
                ;;
        esac
    done
}

# Public entry point
wordstat_request() {
    _method="$1"
    _params="$2"

    if [ -z "$WORDSTAT_BACKEND" ]; then
        die_with_help "wordstat_request called before load_config"
    fi

    case "$WORDSTAT_BACKEND" in
        legacy) _legacy_request "$_method" "$_params" ;;
        cloud)  _cloud_request "$_method" "$_params" ;;
        *)      die_with_help "Unknown backend: $WORDSTAT_BACKEND" ;;
    esac
}

#!/bin/bash
# Get search volume dynamics from Yandex Wordstat

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# common.sh is POSIX sh; bash sources it without issue.
. "$SCRIPT_DIR/common.sh"

# Defaults
PHRASE=""
PERIOD="monthly"
FROM_DATE=""
TO_DATE=""
REGIONS=""
DEVICES="all"

# Parse args
while [[ $# -gt 0 ]]; do
    case $1 in
        --phrase|-p) PHRASE="$2"; shift 2 ;;
        --period) PERIOD="$2"; shift 2 ;;
        --from-date|-f) FROM_DATE="$2"; shift 2 ;;
        --to-date|-t) TO_DATE="$2"; shift 2 ;;
        --regions|-r) REGIONS="$2"; shift 2 ;;
        --devices|-d) DEVICES="$2"; shift 2 ;;
        *) echo "Unknown option: $1"; exit 1 ;;
    esac
done

if [[ -z "$PHRASE" ]]; then
    echo "Usage: dynamics.sh --phrase \"search query\" [options]"
    echo ""
    echo "Options:"
    echo "  --phrase, -p    Search phrase (required)"
    echo "  --period        Grouping: daily, weekly, monthly (default: monthly)"
    echo "  --from-date, -f Start date YYYY-MM-DD (required)"
    echo "  --to-date, -t   End date YYYY-MM-DD (default: today)"
    echo "  --regions, -r   Region IDs, comma-separated (optional)"
    echo "  --devices, -d   Device filter: all, desktop, phone, tablet (default: all)"
    echo ""
    echo "Examples:"
    echo "  bash scripts/dynamics.sh --phrase \"юрист дтп\" --from-date 2025-01-01"
    echo "  bash scripts/dynamics.sh --phrase \"юрист\" --period weekly --from-date 2025-06-01"
    exit 1
fi

# Set default from_date if not provided
if [[ -z "$FROM_DATE" ]]; then
    FROM_DATE=$(date -v-1y +%Y-%m-%d 2>/dev/null || date -d "1 year ago" +%Y-%m-%d 2>/dev/null || echo "2025-01-01")
fi

load_config

# Escape phrase for JSON
PHRASE_ESCAPED=$(json_escape "$PHRASE")

# Build JSON params
PARAMS="{\"phrase\":\"$PHRASE_ESCAPED\",\"period\":\"$PERIOD\",\"fromDate\":\"$FROM_DATE\""

if [[ -n "$TO_DATE" ]]; then
    PARAMS="$PARAMS,\"toDate\":\"$TO_DATE\""
fi

if [[ -n "$REGIONS" ]]; then
    PARAMS="$PARAMS,\"regions\":[$REGIONS]"
fi

if [[ "$DEVICES" != "all" ]]; then
    PARAMS="$PARAMS,\"devices\":\"$DEVICES\""
fi

PARAMS="$PARAMS}"

echo "=== Yandex Wordstat: Dynamics ==="
echo "Phrase: $PHRASE"
echo "Period: $PERIOD"
echo "From: $FROM_DATE"
[[ -n "$TO_DATE" ]] && echo "To: $TO_DATE"
[[ -n "$REGIONS" ]] && echo "Regions: $REGIONS"
echo "Devices: $DEVICES"
echo ""
echo "Fetching data..."

result=$(wordstat_request "dynamics" "$PARAMS")

# Check for error
if echo "$result" | grep -q '"error"'; then
    echo "Error:"
    echo "$result"
    exit 1
fi

echo ""
echo "=== Results ==="
echo ""

echo "| Date | Count |"
echo "|------|-------|"

# Extract dynamics data
echo "$result" | grep -o '{"date":"[^"]*","count":[0-9]*,"share":[^}]*}' | while IFS= read -r entry; do
    dt=$(echo "$entry" | grep -o '"date":"[^"]*"' | sed 's/"date":"//' | tr -d '"')
    cnt=$(echo "$entry" | grep -o '"count":[0-9]*' | sed 's/"count"://')

    echo "| $dt | $(format_number "$cnt") |"
done

echo ""
echo "=== Raw JSON ==="
echo "$result" | head -c 2000
echo ""
echo "[truncated if > 2000 chars]"

#!/bin/bash
# Get Yandex OAuth token for Wordstat API
#
# DEPRECATED — Yandex no longer onboards new users to the legacy Wordstat OAuth API.
# This script only works for users who already have a Yandex OAuth client_id from
# before the deprecation. For new setups, use the cloud backend instead — see
# config/README.md → "Cloud mode (recommended)".

set -e

echo "[NOTICE] Legacy mode only. Yandex stopped onboarding new Wordstat OAuth users."
echo "[NOTICE] For cloud setup see config/README.md → 'Cloud mode'."
echo ""

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
CONFIG_DIR="$SCRIPT_DIR/../config"
ENV_FILE="$CONFIG_DIR/.env"

CLIENT_ID=""
CLIENT_SECRET=""

# Parse args
while [[ $# -gt 0 ]]; do
    case $1 in
        --client-id|-i) CLIENT_ID="$2"; shift 2 ;;
        --client-secret|-s) CLIENT_SECRET="$2"; shift 2 ;;
        *) echo "Unknown option: $1"; exit 1 ;;
    esac
done

if [[ -z "$CLIENT_ID" ]]; then
    echo "Usage: get_token.sh --client-id YOUR_CLIENT_ID [--client-secret YOUR_SECRET]"
    echo ""
    echo "Options:"
    echo "  --client-id, -i      OAuth client ID (required)"
    echo "  --client-secret, -s  OAuth client secret (optional, for code exchange)"
    echo ""
    echo "Get credentials at: https://oauth.yandex.ru/client/new"
    echo ""
    echo "Two modes:"
    echo "  1. Without secret: Opens browser URL, you copy token manually"
    echo "  2. With secret: Exchange authorization code for token"
    exit 1
fi

echo "=== Yandex OAuth Token Setup ==="
echo ""

if [[ -z "$CLIENT_SECRET" ]]; then
    # Simple mode: token in URL fragment
    echo "Step 1: Open this URL in your browser:"
    echo ""
    echo "  https://oauth.yandex.ru/authorize?response_type=token&client_id=$CLIENT_ID"
    echo ""
    echo "Step 2: Authorize the application"
    echo ""
    echo "Step 3: Copy the token from the redirect URL:"
    echo "  https://oauth.yandex.ru/#access_token=YOUR_TOKEN_HERE&..."
    echo ""
    echo -n "Paste your token here: "
    read -r TOKEN

    if [[ -z "$TOKEN" ]]; then
        echo "Error: No token provided"
        exit 1
    fi
else
    # Code exchange mode
    echo "Step 1: Open this URL in your browser:"
    echo ""
    echo "  https://oauth.yandex.ru/authorize?response_type=code&client_id=$CLIENT_ID"
    echo ""
    echo "Step 2: Authorize the application"
    echo ""
    echo "Step 3: Copy the code from the redirect URL or page"
    echo ""
    echo -n "Paste the authorization code here: "
    read -r AUTH_CODE

    if [[ -z "$AUTH_CODE" ]]; then
        echo "Error: No code provided"
        exit 1
    fi

    echo ""
    echo "Exchanging code for token..."

    RESPONSE=$(curl -s -X POST "https://oauth.yandex.ru/token" \
        -H "Content-Type: application/x-www-form-urlencoded" \
        -d "grant_type=authorization_code" \
        -d "code=$AUTH_CODE" \
        -d "client_id=$CLIENT_ID" \
        -d "client_secret=$CLIENT_SECRET")

    # Extract token
    TOKEN=$(echo "$RESPONSE" | grep -o '"access_token":"[^"]*"' | sed 's/"access_token":"//' | tr -d '"')

    if [[ -z "$TOKEN" ]]; then
        echo "Error: Failed to get token"
        echo "$RESPONSE"
        exit 1
    fi

    # Show expiration
    EXPIRES=$(echo "$RESPONSE" | grep -o '"expires_in":[0-9]*' | sed 's/"expires_in"://')
    if [[ -n "$EXPIRES" ]]; then
        DAYS=$((EXPIRES / 86400))
        echo "Token expires in: $DAYS days"
    fi
fi

echo ""
echo "Token received!"
echo ""

# Save to .env
if [[ -f "$ENV_FILE" ]]; then
    # Update existing file
    if grep -q "^YANDEX_WORDSTAT_TOKEN=" "$ENV_FILE"; then
        # Replace existing token
        sed -i.bak "s/^YANDEX_WORDSTAT_TOKEN=.*/YANDEX_WORDSTAT_TOKEN=$TOKEN/" "$ENV_FILE"
        rm -f "$ENV_FILE.bak"
        echo "Updated token in: $ENV_FILE"
    else
        # Append
        echo "YANDEX_WORDSTAT_TOKEN=$TOKEN" >> "$ENV_FILE"
        echo "Added token to: $ENV_FILE"
    fi
else
    # Create new file
    echo "YANDEX_WORDSTAT_TOKEN=$TOKEN" > "$ENV_FILE"
    echo "Created: $ENV_FILE"
fi

echo ""
echo "Verifying token..."
echo ""

# Test the token
bash "$SCRIPT_DIR/quota.sh"

#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.10"
# dependencies = ["openpyxl"]
# ///
"""Missed demand analysis for Yandex Direct campaigns.

Subcommands:
  parse-xlsx   Parse Yandex Direct XLSX export into groups/phrases/minus
  build-query  Build OR-query from slot structure
  merge-slots  Merge batch segmentation results into unified slots
  query-total  Extract totalCount from a pre-fetched Wordstat response file

NOTE: query-total is transport-agnostic. It does NOT make HTTP calls.
The caller (query_total.sh) routes the request through common.sh:wordstat_request,
which handles backend selection (legacy vs cloud) and writes the legacy-shape
JSON response to a file. This file is what query-total reads.
"""
import argparse
import json
import sys
import re
from collections import OrderedDict


# --- Stop words that need + prefix in Wordstat queries ---
STOP_WORDS = frozenset([
    "на", "в", "к", "за", "с", "по", "из", "от", "до", "для",
    "без", "при", "под", "над", "между", "через", "об", "перед",
])

# Characters forbidden in slot variants (OR-syntax operators and leading modifiers)
# Hyphen inside words is OK (б/у, санкт-петербург)
FORBIDDEN_CHARS_RE = re.compile(r'[|()"]')
LEADING_OPERATOR_RE = re.compile(r'^[+\-!]')
# Punctuation to strip from slot variants (dots, commas, semicolons, etc.)
# Keeps hyphens (б/у, санкт-петербург), slashes, and ! (Wordstat exact form operator)
PUNCTUATION_RE = re.compile(r'[.,;:?…·•]')


def cmd_parse_xlsx(args):
    """Parse Yandex Direct XLSX export."""
    import openpyxl

    wb = openpyxl.load_workbook(args.file, data_only=True)
    if "Тексты" not in wb.sheetnames:
        sys.exit(
            f"Error: лист 'Тексты' не найден. Доступные: {wb.sheetnames}"
        )
    ws = wb["Тексты"]

    # Find campaign minus-phrases (scan rows 1-20, cols A-F)
    campaign_minus = ""
    found_minus = False
    for row in ws.iter_rows(min_row=1, max_row=20, min_col=1, max_col=6, values_only=False):
        if found_minus:
            break
        for cell in row:
            if cell.value and "Минус-фразы на кампанию" in str(cell.value):
                next_cell = ws.cell(row=cell.row, column=cell.column + 1)
                if next_cell.value:
                    campaign_minus = str(next_cell.value).strip()
                found_minus = True
                break

    # Find first data row: col A == "-" and col G non-empty
    data_start = None
    for row in ws.iter_rows(min_row=1, max_row=ws.max_row, min_col=1, max_col=1, values_only=False):
        cell = row[0]
        if cell.value and str(cell.value).strip() == "-":
            # Check col G
            g_cell = ws.cell(row=cell.row, column=7)
            if g_cell.value and str(g_cell.value).strip():
                data_start = cell.row
                break

    if data_start is None:
        sys.exit("Error: не найдены строки данных (col A='-' с непустой col G)")

    # Collect groups
    groups = OrderedDict()  # group_id -> {name, group_minus_set, phrases_list}
    for row_idx in range(data_start, ws.max_row + 1):
        col_a = ws.cell(row=row_idx, column=1).value
        if col_a is None or str(col_a).strip() != "-":
            continue

        col_g = ws.cell(row=row_idx, column=7).value
        if col_g is None or not str(col_g).strip():
            continue

        col_c = ws.cell(row=row_idx, column=3).value
        if col_c is None or not str(col_c).strip():
            continue  # skip rows with empty group_id

        group_id = str(col_c).strip()
        group_name = str(ws.cell(row=row_idx, column=4).value or "").strip()
        phrase = str(col_g).strip()

        # Col AG (33) - group minus
        col_ag = ws.cell(row=row_idx, column=33).value
        group_minus_val = str(col_ag).strip() if col_ag else ""

        if group_id not in groups:
            groups[group_id] = {
                "name": group_name,
                "group_minus_set": set(),
                "phrases_ordered": OrderedDict(),
            }

        g = groups[group_id]
        # Deduplicate phrases preserving order
        g["phrases_ordered"][phrase] = None
        if group_minus_val:
            g["group_minus_set"].add(group_minus_val)

    wb.close()

    # Build output
    result_groups = []
    for gid, g in groups.items():
        result_groups.append({
            "id": gid,
            "name": g["name"],
            "group_minus": " ".join(sorted(g["group_minus_set"])),
            "phrases": list(g["phrases_ordered"].keys()),
        })

    # Filter by --group if specified
    if args.group:
        result_groups = [g for g in result_groups if g["id"] == args.group]

    result = {
        "campaign_minus": campaign_minus,
        "campaign_minus_length": len(campaign_minus),
        "groups": result_groups,
    }
    print(json.dumps(result, ensure_ascii=False, indent=2))


def sanitize_variant(variant, slot_name):
    """Sanitize a slot variant. Returns (cleaned, warnings)."""
    warnings = []
    cleaned = variant

    # Strip punctuation (муз. центр -> муз центр)
    if PUNCTUATION_RE.search(cleaned):
        cleaned = PUNCTUATION_RE.sub("", cleaned)
        cleaned = re.sub(r'\s+', ' ', cleaned).strip()

    # Remove forbidden OR-syntax characters (replace with space to avoid merging tokens)
    if FORBIDDEN_CHARS_RE.search(cleaned):
        warnings.append(
            f"slot '{slot_name}': вариант '{variant}' содержит OR-синтаксис, деградирован в многословный"
        )
        cleaned = FORBIDDEN_CHARS_RE.sub(" ", cleaned)
        # Collapse multiple spaces
        cleaned = re.sub(r'\s+', ' ', cleaned).strip()

    # Remove leading operators (+ - !) from each token
    tokens = cleaned.split()
    new_tokens = []
    for t in tokens:
        if LEADING_OPERATOR_RE.match(t):
            stripped = LEADING_OPERATOR_RE.sub("", t)
            if stripped:
                warnings.append(
                    f"slot '{slot_name}': убран ведущий оператор из '{t}'"
                )
                new_tokens.append(stripped)
        else:
            new_tokens.append(t)
    cleaned = " ".join(new_tokens)

    return cleaned.strip(), warnings


def add_stop_word_plus(variant):
    """Add + prefix to stop words in a variant string."""
    tokens = variant.split()
    result = []
    for t in tokens:
        # Don't touch tokens already starting with operator
        if t[0] in ("+", "-", "!"):
            result.append(t)
        elif t.lower() in STOP_WORDS:
            result.append(f"+{t}")
        else:
            result.append(t)
    return " ".join(result)


def normalize_token(token):
    """Normalize a token for comparison: lowercase, strip punctuation."""
    t = token.lower().strip()
    t = PUNCTUATION_RE.sub("", t)
    t = LEADING_OPERATOR_RE.sub("", t)
    t = FORBIDDEN_CHARS_RE.sub("", t)
    return t.strip()


class SlotMerger:
    """Merges slot segmentation results from multiple LLM batches.

    Does NOT perform LLM segmentation itself — only merges
    already-segmented results and provides coverage analysis.
    """

    SLOT_NAMES = ["objects", "actions", "modifiers", "additional"]
    ADDITIONAL_PATTERN_RE = re.compile(
        r'\b(' + '|'.join(STOP_WORDS) + r')\s+(\S+)',
        re.IGNORECASE,
    )

    def __init__(self, phrases):
        self._phrases = list(phrases)
        self._slots = {
            name: OrderedDict() for name in self.SLOT_NAMES
        }

    def merge_batch(self, batch_index, phrase_indexes, slots):
        """Merge slots from one batch into accumulated state.

        Args:
            batch_index: sequential index of this batch
            phrase_indexes: indices into self._phrases covered by batch
            slots: dict with keys from SLOT_NAMES, values are variant lists

        Raises:
            ValueError: if phrase_indexes are out of range
        """
        n = len(self._phrases)
        for pi in phrase_indexes:
            if not isinstance(pi, int) or pi < 0 or pi >= n:
                raise ValueError(
                    f"phrase_index {pi} вне диапазона [0..{n - 1}]"
                )
        for slot_name in self.SLOT_NAMES:
            variants = slots.get(slot_name, [])
            if not isinstance(variants, list):
                continue
            for v in variants:
                v_str = str(v).strip()
                if not v_str:
                    continue
                cleaned, _ = sanitize_variant(v_str, slot_name)
                if not cleaned:
                    continue
                canon = re.sub(r'\s+', ' ', cleaned.lower()).strip()
                slot_dict = self._slots[slot_name]
                if canon not in slot_dict:
                    slot_dict[canon] = {
                        "display": cleaned,
                        "batch_indexes": set(),
                    }
                slot_dict[canon]["batch_indexes"].add(batch_index)

    def _remove_subsets(self):
        """Remove subset variants within each slot (heuristic).

        If normalized tokens of A are a strict subset of tokens of B,
        B is redundant in Wordstat OR (A already covers wider queries).
        Comparison uses normalized tokens (lower, no punctuation).
        """
        for slot_name in self.SLOT_NAMES:
            slot_dict = self._slots[slot_name]
            keys = list(slot_dict.keys())
            token_sets = {}
            for k in keys:
                tokens = set(
                    normalize_token(t)
                    for t in k.split()
                    if normalize_token(t)
                )
                token_sets[k] = tokens
            to_remove = set()
            for i, k_short in enumerate(keys):
                for j, k_long in enumerate(keys):
                    if i == j or k_long in to_remove:
                        continue
                    ts_short = token_sets[k_short]
                    ts_long = token_sets[k_long]
                    if ts_short and ts_short < ts_long:
                        to_remove.add(k_long)
            for k in to_remove:
                del slot_dict[k]

    def _compute_coverage(self):
        """Compute coverage of original phrases by merged slot variants."""
        all_variant_tokens = set()
        for slot_name in self.SLOT_NAMES:
            for canon in self._slots[slot_name]:
                for t in canon.split():
                    nt = normalize_token(t)
                    if nt:
                        all_variant_tokens.add(nt)

        uncovered_phrases = []
        all_phrase_tokens = set()
        for idx, phrase in enumerate(self._phrases):
            phrase_tokens = set()
            for t in phrase.split():
                nt = normalize_token(t)
                if nt and nt not in STOP_WORDS:
                    phrase_tokens.add(nt)
                    all_phrase_tokens.add(nt)
            if phrase_tokens and not (phrase_tokens & all_variant_tokens):
                uncovered_phrases.append({"index": idx, "phrase": phrase})

        uncovered_tokens = sorted(
            all_phrase_tokens - all_variant_tokens - STOP_WORDS
        )

        additional_variants_lower = set(
            self._slots.get("additional", OrderedDict()).keys()
        )
        additional_patterns = []
        seen_patterns = set()
        for phrase in self._phrases:
            normalized_phrase = PUNCTUATION_RE.sub("", phrase.lower())
            normalized_phrase = re.sub(r'\s+', ' ', normalized_phrase).strip()
            for m in self.ADDITIONAL_PATTERN_RE.finditer(normalized_phrase):
                pattern = re.sub(r'\s+', ' ', m.group(0)).strip()
                if pattern not in seen_patterns:
                    seen_patterns.add(pattern)
                    if pattern not in additional_variants_lower:
                        additional_patterns.append(pattern)

        return {
            "uncovered_phrases": uncovered_phrases,
            "uncovered_tokens": uncovered_tokens,
            "additional_patterns": additional_patterns,
        }

    def finalize(self, max_query_length=4096):
        """Finalize merged slots: deduplicate, remove subsets, build query."""
        self._remove_subsets()

        slots_pre_trim = {}
        for slot_name in self.SLOT_NAMES:
            slots_pre_trim[slot_name] = [
                entry["display"]
                for entry in self._slots[slot_name].values()
            ]

        coverage = self._compute_coverage()

        debug = {}
        for slot_name in self.SLOT_NAMES:
            debug[slot_name] = {
                entry["display"]: sorted(entry["batch_indexes"])
                for entry in self._slots[slot_name].values()
            }

        slot_order = self.SLOT_NAMES
        warnings = []
        trimmed = []

        sanitized_slots = {}
        for slot_name in slot_order:
            clean_variants = []
            for entry in self._slots[slot_name].values():
                clean_variants.append(add_stop_word_plus(entry["display"]))
            sanitized_slots[slot_name] = clean_variants

        def build_or_string(slot_variants):
            if not slot_variants:
                return ""
            if len(slot_variants) == 1:
                return slot_variants[0]
            return "(" + "|".join(slot_variants) + ")"

        def assemble_query(s_slots):
            parts = []
            for name in slot_order:
                or_str = build_or_string(s_slots[name])
                if or_str:
                    parts.append(or_str)
            return " ".join(parts)

        def est_phrases(s_slots):
            product = 1
            for name in slot_order:
                n = len(s_slots[name])
                if n > 0:
                    product *= n
            return product

        trim_order = ["additional", "modifiers", "actions"]
        while est_phrases(sanitized_slots) > 200:
            trimmed_any = False
            for trim_slot in trim_order:
                if len(sanitized_slots[trim_slot]) > 1:
                    removed = sanitized_slots[trim_slot].pop()
                    trimmed.append({
                        "slot": trim_slot, "removed": removed,
                        "reason": "estimated_phrases > 200",
                    })
                    trimmed_any = True
                    break
            if not trimmed_any:
                break

        query = assemble_query(sanitized_slots)
        while len(query) > max_query_length:
            trimmed_any = False
            for trim_slot in trim_order:
                if len(sanitized_slots[trim_slot]) > 1:
                    removed = sanitized_slots[trim_slot].pop()
                    trimmed.append({
                        "slot": trim_slot, "removed": removed,
                        "reason": f"query_length > {max_query_length}",
                    })
                    trimmed_any = True
                    query = assemble_query(sanitized_slots)
                    break
            if not trimmed_any:
                warnings.append(
                    f"query_length {len(query)} > {max_query_length}"
                    " после максимальной обрезки"
                )
                break

        slots_post_trim = {}
        for slot_name in slot_order:
            slots_post_trim[slot_name] = list(sanitized_slots[slot_name])

        return {
            "slots_pre_trim": slots_pre_trim,
            "slots_post_trim": slots_post_trim,
            "query": query,
            "estimated_phrases": est_phrases(sanitized_slots),
            "query_length": len(query),
            "trimmed": trimmed,
            "warnings": warnings,
            "coverage": coverage,
            "debug": debug,
        }


def cmd_merge_slots(args):
    """Merge batch segmentation results from stdin JSON."""
    raw = sys.stdin.read()
    try:
        data = json.loads(raw)
    except json.JSONDecodeError as e:
        sys.exit(f"Error: невалидный JSON на stdin: {e}")

    phrases = data.get("phrases", [])
    batches = data.get("batches", [])

    if not isinstance(phrases, list) or not phrases:
        sys.exit("Error: phrases должен быть непустым списком строк")
    if not isinstance(batches, list) or not batches:
        sys.exit("Error: batches должен быть непустым списком объектов")
    for p in phrases:
        if not isinstance(p, str):
            sys.exit(f"Error: элемент phrases не строка: {type(p).__name__}")

    merger = SlotMerger(phrases)
    for i, batch in enumerate(batches):
        if not isinstance(batch, dict):
            sys.exit(f"Error: batch[{i}] не объект")
        phrase_indexes = batch.get("phrase_indexes", [])
        if not isinstance(phrase_indexes, list):
            sys.exit(f"Error: batch[{i}].phrase_indexes не список")
        slots = batch.get("slots", {})
        if not isinstance(slots, dict):
            sys.exit(f"Error: batch[{i}].slots не объект")
        try:
            merger.merge_batch(i, phrase_indexes, slots)
        except ValueError as e:
            sys.exit(f"Error: batch[{i}]: {e}")

    result = merger.finalize(max_query_length=args.max_query_length)
    print(json.dumps(result, ensure_ascii=False, indent=2))


def cmd_build_query(args):
    """Build OR-query from slots JSON."""
    try:
        slots = json.loads(args.slots_json)
    except json.JSONDecodeError as e:
        sys.exit(f"Error: невалидный JSON слотов: {e}")

    max_len = args.max_query_length
    warnings = []
    trimmed = []

    # Slot order: actions, objects, modifiers, additional
    slot_order = ["actions", "objects", "modifiers", "additional"]

    # Sanitize all variants
    sanitized_slots = {}
    for slot_name in slot_order:
        raw_variants = slots.get(slot_name, [])
        if not isinstance(raw_variants, list):
            raw_variants = []
        clean_variants = []
        for v in raw_variants:
            v_str = str(v).strip()
            if not v_str:
                continue
            cleaned, san_warnings = sanitize_variant(v_str, slot_name)
            warnings.extend(san_warnings)
            if cleaned:
                # Apply stop-word plus
                cleaned = add_stop_word_plus(cleaned)
                clean_variants.append(cleaned)
        sanitized_slots[slot_name] = clean_variants

    def build_or_string(slot_variants):
        """Build OR part: (a|b|c) or a."""
        if len(slot_variants) == 0:
            return ""
        if len(slot_variants) == 1:
            return slot_variants[0]
        return "(" + "|".join(slot_variants) + ")"

    def assemble_query(s_slots):
        """Assemble full query from slots dict."""
        parts = []
        for name in slot_order:
            or_str = build_or_string(s_slots[name])
            if or_str:
                parts.append(or_str)
        return " ".join(parts)

    def estimated_phrases(s_slots):
        """Calculate estimated phrase count (product of slot sizes)."""
        product = 1
        for name in slot_order:
            n = len(s_slots[name])
            if n > 0:
                product *= n
        return product

    # Trim priority: additional -> modifiers -> actions (objects untouched)
    trim_order = ["additional", "modifiers", "actions"]

    # Check estimated phrases limit
    while estimated_phrases(sanitized_slots) > 200:
        trimmed_any = False
        for trim_slot in trim_order:
            if len(sanitized_slots[trim_slot]) > 1:
                removed = sanitized_slots[trim_slot].pop()
                trimmed.append({"slot": trim_slot, "removed": removed, "reason": "estimated_phrases > 200"})
                trimmed_any = True
                break
        if not trimmed_any:
            break

    # Check query length limit
    query = assemble_query(sanitized_slots)
    while len(query) > max_len:
        trimmed_any = False
        for trim_slot in trim_order:
            if len(sanitized_slots[trim_slot]) > 1:
                removed = sanitized_slots[trim_slot].pop()
                trimmed.append({"slot": trim_slot, "removed": removed, "reason": f"query_length > {max_len}"})
                trimmed_any = True
                query = assemble_query(sanitized_slots)
                break
        if not trimmed_any:
            warnings.append(
                f"query_length {len(query)} всё ещё > {max_len} после максимальной обрезки"
            )
            break

    result = {
        "query": query,
        "estimated_phrases": estimated_phrases(sanitized_slots),
        "query_length": len(query),
        "trimmed": trimmed,
        "warnings": warnings,
    }
    print(json.dumps(result, ensure_ascii=False, indent=2))


def cmd_query_total(args):
    """Extract totalCount from a pre-fetched Wordstat response file.

    Transport-agnostic: caller (query_total.sh) is responsible for the HTTP call
    via common.sh:wordstat_request, which handles legacy vs cloud backend
    selection. This function only parses the resulting legacy-shape JSON.
    """
    try:
        with open(args.json_file, "r", encoding="utf-8") as f:
            resp_body = f.read()
    except OSError as e:
        result = {"error": f"Cannot read response file: {e}", "query": args.phrase}
        print(json.dumps(result, ensure_ascii=False))
        sys.exit(1)

    try:
        obj = json.loads(resp_body)
    except json.JSONDecodeError:
        result = {"error": "Invalid JSON response", "query": args.phrase, "raw": resp_body[:500]}
        print(json.dumps(result, ensure_ascii=False))
        sys.exit(1)

    # Check for errors (multiple possible formats)
    if "error" in obj:
        err_val = obj["error"]
        result = {"error": str(err_val), "query": args.phrase, "raw": resp_body[:500]}
        print(json.dumps(result, ensure_ascii=False))
        sys.exit(1)

    if "error_code" in obj:
        err_msg = f"{obj.get('error_str', '')}: {obj.get('error_detail', '')}"
        result = {
            "error": err_msg.strip(": "),
            "error_code": obj["error_code"],
            "query": args.phrase,
            "raw": resp_body[:500],
        }
        print(json.dumps(result, ensure_ascii=False))
        sys.exit(1)

    # Extract totalCount with fallbacks (top-level or nested under "result")
    total_count = obj.get("totalCount")
    if total_count is None:
        total_count = (obj.get("result") or {}).get("totalCount")
    if total_count is None:
        result = {
            "error": "totalCount not found in response",
            "query": args.phrase,
            "raw": resp_body[:500],
        }
        print(json.dumps(result, ensure_ascii=False))
        sys.exit(1)

    result = {"total_count": int(total_count), "query": args.phrase}
    print(json.dumps(result, ensure_ascii=False))


def main():
    parser = argparse.ArgumentParser(description="Missed demand analysis tools")
    sub = parser.add_subparsers(dest="command", required=True)

    # parse-xlsx
    p_parse = sub.add_parser("parse-xlsx", help="Parse Yandex Direct XLSX export")
    p_parse.add_argument("file", help="Path to XLSX file")
    p_parse.add_argument("--group", default=None, help="Filter by group ID")

    # build-query
    p_build = sub.add_parser("build-query", help="Build OR-query from slots")
    p_build.add_argument("slots_json", help="JSON string with slots")
    p_build.add_argument("--max-query-length", type=int, default=4096, help="Max query length (default: 4096)")

    # merge-slots
    p_merge = sub.add_parser(
        "merge-slots",
        help="Merge batch segmentation results (stdin JSON)",
    )
    p_merge.add_argument(
        "--max-query-length", type=int, default=4096,
        help="Max query length (default: 4096)",
    )

    # query-total
    p_query = sub.add_parser(
        "query-total",
        help="Extract totalCount from a pre-fetched Wordstat response file (transport-agnostic)",
    )
    p_query.add_argument(
        "--json-file",
        required=True,
        help="Path to JSON file containing legacy-shape Wordstat response (written by common.sh:wordstat_request)",
    )
    p_query.add_argument("--phrase", required=True, help="Original search phrase (for output formatting)")

    args = parser.parse_args()
    if args.command == "parse-xlsx":
        cmd_parse_xlsx(args)
    elif args.command == "build-query":
        cmd_build_query(args)
    elif args.command == "merge-slots":
        cmd_merge_slots(args)
    elif args.command == "query-total":
        cmd_query_total(args)


if __name__ == "__main__":
    main()