#!/usr/bin/env python
"""One-off migration: bump the stored `maxTokens` on existing widgets and
characters in Drupal.

The interactive-chat default was 4096 (and silently dropped on the way to the
backend before the LLModel alias fix). This script rewrites the model block in
every existing widget/character node body so they pick up a larger value.

Casing matches each entity's storage convention (both are accepted on read):
  - characters store snake_case `max_tokens` (written via model_dump_json)
  - widgets store camelCase `maxTokens` (written by create_widget.py / frontend)

Entities with no `model` block are skipped (they fall back to the code default).

Run it against whichever Drupal the api/.env points at, using that server's
repo venv, e.g.:

    cd /tsai/repo/api && TSAI_ENV=production ./venv/bin/python scripts/bump_max_tokens.py --dry-run
    cd /data/tsai/api && TSAI_ENV=scike_ai  ./venv/bin/python scripts/bump_max_tokens.py
"""
import argparse
import asyncio
import json
import sys
from pathlib import Path

import requests
from dotenv import load_dotenv

# Load api/.env before importing DrupalService — drupal.py reads DRUPAL_HOST at
# import time.
API_DIR = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(API_DIR))
load_dotenv(API_DIR / ".env")

from app.services.drupal import DrupalService, DRUPAL_HOST, DRUPAL_API  # noqa: E402

DEFAULT_MAX_TOKENS = 16384


def _entity_name(attrs, body):
    return attrs.get("title") or body.get("card", {}).get("title") or "(unnamed)"


def _update_model(body, token_key):
    """Set `token_key` on body['model'] to the new value. Returns (old, new) or
    None if there is no model block to update."""
    model = body.get("model")
    if not isinstance(model, dict):
        return None
    old = model.get("max_tokens", model.get("maxTokens"))
    # Normalize to the canonical key for this entity type, dropping the other
    # casing so there is no ambiguity.
    other = "maxTokens" if token_key == "max_tokens" else "max_tokens"
    model.pop(other, None)
    model[token_key] = DEFAULT_MAX_TOKENS
    return old, DEFAULT_MAX_TOKENS


def _patch_node(session, csrf_token, node_type, node_id, body, body_format):
    headers = {
        "Content-Type": "application/vnd.api+json",
        "X-CSRF-Token": csrf_token,
        "Authorization": f"Bearer {csrf_token}",
    }
    node_data = {
        "data": {
            "type": node_type,
            "id": node_id,
            # JSON:API partial update — only the body attribute is touched.
            "attributes": {
                "body": {"value": json.dumps(body), "format": body_format}
            },
        }
    }
    url = f"{DRUPAL_HOST}{DRUPAL_API}/node/{node_type.split('--')[1]}/{node_id}"
    return session.patch(url, headers=headers, json=node_data)


def _process(session, csrf_token, label, node_type, nodes, token_key, dry_run):
    updated = skipped = failed = 0
    for node in nodes:
        node_id = node.get("id", "")
        attrs = node.get("attributes", {})
        body_field = attrs.get("body") or {}
        body_format = body_field.get("format", "plain_text")
        try:
            body = json.loads(body_field.get("value", "{}"))
        except (json.JSONDecodeError, TypeError):
            print(f"  ! {label} {node_id}: unparseable body, skipping")
            skipped += 1
            continue

        name = _entity_name(attrs, body)
        result = _update_model(body, token_key)
        if result is None:
            print(f"  - {label} '{name}': no model block, skipping")
            skipped += 1
            continue

        old, new = result
        if old == new:
            print(f"  = {label} '{name}': already {new}, skipping")
            skipped += 1
            continue

        if dry_run:
            print(f"  ~ {label} '{name}': {old} -> {new} (dry-run)")
            updated += 1
            continue

        resp = _patch_node(session, csrf_token, node_type, node_id, body, body_format)
        if resp.status_code == 200:
            print(f"  + {label} '{name}': {old} -> {new}")
            updated += 1
        else:
            print(f"  ! {label} '{name}': PATCH failed {resp.status_code} {resp.text[:200]}")
            failed += 1

    return updated, skipped, failed


async def main():
    parser = argparse.ArgumentParser(description="Bump maxTokens on existing widgets/characters.")
    parser.add_argument("--dry-run", action="store_true", help="Report changes without patching.")
    parser.add_argument("--skip-widgets", action="store_true", help="Don't touch widgets.")
    parser.add_argument("--skip-characters", action="store_true", help="Don't touch characters.")
    args = parser.parse_args()

    drupal = DrupalService()
    session = requests.Session()
    totals = [0, 0, 0]  # updated, skipped, failed
    try:
        login_data = await drupal.login(session)
        logout_token = login_data.get("logout_token")
        csrf_token = await drupal.get_csrf_token(session)

        if not args.skip_characters:
            print("Characters:")
            chars = drupal.fetch_all_characters()
            u, s, f = _process(session, csrf_token, "character", "node--character",
                               chars, "max_tokens", args.dry_run)
            totals = [totals[0] + u, totals[1] + s, totals[2] + f]

        if not args.skip_widgets:
            print("Widgets:")
            widgets = drupal.fetch_all_widgets()
            u, s, f = _process(session, csrf_token, "widget", "node--widget",
                               widgets, "maxTokens", args.dry_run)
            totals = [totals[0] + u, totals[1] + s, totals[2] + f]

        await drupal.logout(session, logout_token)
    finally:
        session.close()

    verb = "would update" if args.dry_run else "updated"
    print(f"\nDone — {verb}: {totals[0]}, skipped: {totals[1]}, failed: {totals[2]}")
    return 1 if totals[2] else 0


if __name__ == "__main__":
    sys.exit(asyncio.run(main()))