#!/usr/bin/env python3
"""Read CLIProxyAPI account quotas and render a MiniToo animation. No token copies."""
import argparse
import concurrent.futures
from datetime import datetime, timezone
import json
import io
import math
import os
from pathlib import Path
import struct
import time

from PIL import Image, ImageDraw, ImageFont
import requests
import yaml

ROOT = Path(__file__).resolve().parent
STATE = ROOT / 'state'
CONFIG = Path(os.environ['CLIPROXY_CONFIG']).expanduser()
FONT = '/System/Library/Fonts/Monaco.ttf'

def percent(value):
    if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value):
        return None
    return max(0, min(100, 100 - value))

def window(label, item, used_key, reset_key):
    if not isinstance(item, dict) or percent(item.get(used_key)) is None:
        return None
    return {'label': label, 'remaining': percent(item[used_key]), 'reset': item.get(reset_key)}

def parse_quota(kind, body):
    if kind == 'claude':
        candidates = [window(label, body.get(key), 'utilization', 'resets_at')
                      for key, label in [('five_hour', '5h'), ('seven_day', '7d')]]
    else:
        candidates = []
        rate = body.get('rate_limit') or {}
        for key in ['primary_window', 'secondary_window']:
            item = rate.get(key)
            if not isinstance(item, dict):
                continue
            seconds = item.get('limit_window_seconds')
            label = 'limit'
            if isinstance(seconds, (int, float)) and seconds > 0:
                label = f'{seconds / 86400:g}d' if seconds % 86400 == 0 else f'{seconds / 3600:g}h'
            candidates.append(window(label, item, 'used_percent', 'reset_at'))
    return [x for x in candidates if x is not None]

def fetch_account(item):
    path, label = item
    account = {'label': label, 'windows': [], 'status': 'unavailable'}
    try:
        auth = json.loads(path.read_text())
        if auth.get('disabled'):
            account['status'] = 'disabled'
            return account
        kind = auth['type']
        headers = {'Authorization': 'Bearer ' + auth['access_token'], 'Accept': 'application/json', 'User-Agent': 'divoom-usage/1.0'}
        if kind == 'codex':
            url = 'https://chatgpt.com/backend-api/wham/usage'
            headers['ChatGPT-Account-Id'] = auth['account_id']
        else:
            url = 'https://api.anthropic.com/api/oauth/usage'
            headers['anthropic-beta'] = 'oauth-2025-04-20'
        response = requests.get(url, headers=headers, timeout=20, allow_redirects=False)
        if response.status_code != 200:
            account['status'] = 'login needed' if response.status_code in [401, 403] else f'HTTP {response.status_code}'
            return account
        account['windows'] = parse_quota(kind, response.json())
        account['status'] = 'ok' if account['windows'] else 'no quota data'
    except (requests.RequestException, ValueError, KeyError, OSError):
        account['status'] = 'read failed'
    return account

def accounts():
    cfg = yaml.safe_load(CONFIG.read_text())
    auth_dir = Path(cfg['auth-dir']).expanduser()
    labels_path = ROOT / 'accounts.json'
    labels = json.loads(labels_path.read_text()) if labels_path.exists() else {}
    found = []
    counts = {'codex': 0, 'claude': 0}
    for path in sorted(auth_dir.glob('*.json')):
        try:
            kind = json.loads(path.read_text()).get('type')
        except (ValueError, OSError):
            continue
        if kind not in counts:
            continue
        counts[kind] += 1
        label = labels.get(path.name)
        if label is None:
            number = 1
            while f'{kind.upper()} {number}' in labels.values():
                number += 1
            label = f'{kind.upper()} {number}'
            labels[path.name] = label
        found.append((path, label))
    labels_path.write_text(json.dumps(labels, indent=2) + '\n')
    labels_path.chmod(0o600)
    return found

def reset_text(value):
    if value is None:
        return 'reset not given'
    try:
        date = datetime.fromisoformat(value.replace('Z', '+00:00')) if isinstance(value, str) else datetime.fromtimestamp(value, timezone.utc)
        seconds = max(0, int(date.timestamp() - time.time()))
        if seconds >= 86400:
            return f'reset {seconds//86400}d {seconds%86400//3600}h'
        return f'reset {seconds//3600}h {seconds%3600//60}m'
    except (ValueError, TypeError, OverflowError):
        return 'reset unknown'

def render(snapshot):
    frames = []
    small = ImageFont.truetype(FONT, 9)
    font = ImageFont.truetype(FONT, 11)
    title = ImageFont.truetype(FONT, 13)
    for account in snapshot['accounts']:
        img = Image.new('RGB', (160, 128), '#0b111b')
        draw = ImageDraw.Draw(img)
        draw.text((6, 5), account['label'][:14], font=title, fill='#f1f5f9')
        draw.text((6, 23), 'QUOTA LEFT', font=small, fill='#94a3b8')
        if account['status'] != 'ok':
            draw.text((6, 52), account['status'], font=font, fill='#ffb454')
        for i, w in enumerate(account['windows'][:2]):
            y = 40 + i * 35
            value = w['remaining']
            color = '#41dca0' if value > 30 else '#ffb454' if value > 10 else '#ff6070'
            draw.text((6, y), w['label'], font=font, fill='#e2e8f0')
            draw.text((154, y), f'{value:.0f}%', anchor='ra', font=font, fill=color)
            draw.rectangle((6, y+14, 153, y+18), fill='#273142')
            if value > 0:
                draw.rectangle((6, y+14, 6+round(147*value/100), y+18), fill=color)
            draw.text((6, y+21), reset_text(w['reset']), font=small, fill='#94a3b8')
        draw.line((6, 113, 154, 113), fill='#273142')
        stamp = datetime.fromtimestamp(snapshot['fetched_at']).strftime('%d %b %H:%M')
        draw.text((6, 117), stamp, font=small, fill='#94a3b8')
        frames.append(img)
    if not frames:
        img = Image.new('RGB', (160,128), '#0b111b')
        ImageDraw.Draw(img).text((5,50), 'NO ACCOUNTS', font=font, fill='#ffb454')
        frames.append(img)
    return frames

def packet(body):
    data = struct.pack('<H', len(body) + 3) + b'\x8b' + body
    return b'\x01' + data + struct.pack('<H', sum(data) & 0xffff) + b'\x02'

def encode(frames):
    if not 1 <= len(frames) <= 255:
        raise ValueError('Device supports 1..255 frames')
    # Full-width MiniToo LCD format: 8 rows x 10 columns of 16px cells.
    payload = struct.pack('>BBHBB', 0x23, len(frames), 5000, 8, 10)
    for frame in frames:
        if frame.size != (160, 128):
            raise ValueError('Full-width MiniToo frames must be 160x128')
        image_bytes = io.BytesIO()
        frame.convert('RGB').save(image_bytes, format='JPEG', quality=95, subsampling=0)
        jpeg = image_bytes.getvalue()
        payload += b'\x01' + struct.pack('>I', len(jpeg)) + jpeg
    size = struct.pack('<I', len(payload))
    packets = [packet(b'\x00' + size)]
    packets += [packet(b'\x01' + size + struct.pack('<H', n) + payload[pos:pos+256])
                for n, pos in enumerate(range(0, len(payload), 256))]
    return b''.join(struct.pack('<H', len(p)) + p for p in packets)

def update():
    STATE.mkdir(exist_ok=True, mode=0o700)
    with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool:
        result = list(pool.map(fetch_account, accounts()))
    snapshot = {'fetched_at': time.time(), 'accounts': result}
    (STATE/'usage.json').write_text(json.dumps(snapshot, indent=2) + '\n')
    frames = render(snapshot)
    frames[0].save(STATE/'preview.png')
    frames[0].save(STATE/'preview.gif', save_all=True, append_images=frames[1:], duration=5000, loop=0)
    temp = STATE/'display.tmp'
    temp.write_bytes(encode(frames))
    temp.replace(STATE/'display.bin')
    print(json.dumps({'event':'quota_updated','accounts':len(result),'ok':sum(a['status']=='ok' for a in result)}), flush=True)

def main():
    os.umask(0o077)
    parser = argparse.ArgumentParser()
    parser.add_argument('--watch', action='store_true')
    args = parser.parse_args()
    while True:
        update()
        if not args.watch:
            break
        time.sleep(300)

if __name__ == '__main__':
    main()