#!/usr/bin/env python3
"""Aria Usage Tracker - aggregiert Token-Nutzung aus Claude Code JSONL Logs"""

import json
import os
import glob
from datetime import datetime, timedelta
from pathlib import Path

PROJECTS_DIR = "/root/.claude/projects/-root"

def parse_jsonl(filepath):
    """Parse JSONL und extrahiere Usage-Daten"""
    tokens = {"input": 0, "output": 0, "cache_read": 0, "cache_create": 0}
    msg_count = 0
    first_ts = None
    last_ts = None

    try:
        with open(filepath) as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    d = json.loads(line)
                    msg = d.get("message", {})
                    usage = msg.get("usage", {})

                    if usage:
                        tokens["input"] += usage.get("input_tokens", 0)
                        tokens["output"] += usage.get("output_tokens", 0)
                        tokens["cache_read"] += usage.get("cache_read_input_tokens", 0)
                        tokens["cache_create"] += usage.get("cache_creation_input_tokens", 0)
                        msg_count += 1

                    # Timestamps
                    snapshot = d.get("snapshot", {})
                    ts_str = snapshot.get("timestamp")
                    if ts_str:
                        try:
                            ts = datetime.fromisoformat(ts_str.replace("Z", "+00:00"))
                            if first_ts is None or ts < first_ts:
                                first_ts = ts
                            if last_ts is None or ts > last_ts:
                                last_ts = ts
                        except Exception:
                            pass
                except json.JSONDecodeError:
                    continue
    except Exception:
        pass

    return tokens, msg_count, first_ts, last_ts


def get_file_date(filepath):
    """Holt das Erstelldatum einer Datei"""
    try:
        stat = os.stat(filepath)
        return datetime.fromtimestamp(stat.st_mtime)
    except Exception:
        return None


def aggregate_usage(days=None):
    """Aggregiert Usage ueber alle Sessions, optional nach Tagen gefiltert"""
    cutoff = None
    if days:
        cutoff = datetime.now() - timedelta(days=days)

    total = {"input": 0, "output": 0, "cache_read": 0, "cache_create": 0}
    total_msgs = 0
    session_count = 0

    jsonl_files = glob.glob(os.path.join(PROJECTS_DIR, "*.jsonl"))

    # Auch Subagent-Logs
    jsonl_files += glob.glob(os.path.join(PROJECTS_DIR, "*/subagents/*.jsonl"))

    for filepath in jsonl_files:
        file_date = get_file_date(filepath)
        if cutoff and file_date and file_date < cutoff:
            continue

        tokens, msgs, first_ts, last_ts = parse_jsonl(filepath)
        if msgs > 0:
            session_count += 1
            total_msgs += msgs
            for k in total:
                total[k] += tokens[k]

    total_all = total["input"] + total["output"] + total["cache_read"] + total["cache_create"]

    return {
        "sessions": session_count,
        "messages": total_msgs,
        "input_tokens": total["input"],
        "output_tokens": total["output"],
        "cache_read_tokens": total["cache_read"],
        "cache_create_tokens": total["cache_create"],
        "total_tokens": total_all,
        "total_tokens_m": round(total_all / 1_000_000, 2),
    }


def get_current_session_usage():
    """Holt Usage der aktuellsten Session"""
    jsonl_files = glob.glob(os.path.join(PROJECTS_DIR, "*.jsonl"))
    if not jsonl_files:
        return None

    # Neueste Datei
    latest = max(jsonl_files, key=os.path.getmtime)
    tokens, msgs, first_ts, last_ts = parse_jsonl(latest)

    total_all = tokens["input"] + tokens["output"] + tokens["cache_read"] + tokens["cache_create"]

    return {
        "session_file": os.path.basename(latest),
        "messages": msgs,
        "input_tokens": tokens["input"],
        "output_tokens": tokens["output"],
        "cache_read_tokens": tokens["cache_read"],
        "cache_create_tokens": tokens["cache_create"],
        "total_tokens": total_all,
        "total_tokens_k": round(total_all / 1_000, 1),
    }


if __name__ == "__main__":
    print("=== Aktuelle Session ===")
    current = get_current_session_usage()
    if current:
        print(f"  Messages:     {current['messages']}")
        print(f"  Input:        {current['input_tokens']:,} tokens")
        print(f"  Output:       {current['output_tokens']:,} tokens")
        print(f"  Cache Read:   {current['cache_read_tokens']:,} tokens")
        print(f"  Cache Create: {current['cache_create_tokens']:,} tokens")
        print(f"  Total:        {current['total_tokens_k']}k tokens")

    print("\n=== Heute ===")
    today = aggregate_usage(days=1)
    print(f"  Sessions: {today['sessions']}")
    print(f"  Messages: {today['messages']}")
    print(f"  Total:    {today['total_tokens_m']}M tokens")

    print("\n=== Diese Woche ===")
    week = aggregate_usage(days=7)
    print(f"  Sessions: {week['sessions']}")
    print(f"  Messages: {week['messages']}")
    print(f"  Total:    {week['total_tokens_m']}M tokens")

    print("\n=== Gesamt ===")
    all_time = aggregate_usage()
    print(f"  Sessions: {all_time['sessions']}")
    print(f"  Messages: {all_time['messages']}")
    print(f"  Total:    {all_time['total_tokens_m']}M tokens")