#!/usr/bin/env python3

import json
import sys
from datetime import datetime, timedelta, timezone
from pathlib import Path
from urllib.request import Request, urlopen
from urllib.error import URLError, HTTPError

# Constants
CREDENTIALS_PATH = Path.home() / ".claude" / ".credentials.json"
API_ENDPOINT = "https://api.anthropic.com/api/oauth/usage"
REQUEST_TIMEOUT = 10

def read_credentials():
    """Read OAuth access token and expiration from Claude Code credentials file."""
    try:
        with open(CREDENTIALS_PATH, 'r') as f:
            creds = json.load(f)
            oauth_data = creds.get('claudeAiOauth', {})
            access_token = oauth_data.get('accessToken')
            expires_at = oauth_data.get('expiresAt')
            return (access_token, expires_at)
    except FileNotFoundError:
        return (None, None)
    except json.JSONDecodeError:
        return (None, None)

def is_token_expired(expires_at):
    """Check if the OAuth token has expired."""
    if expires_at is None:
        return True

    try:
        import time
        current_time_ms = int(time.time() * 1000)
        return current_time_ms >= expires_at
    except (ValueError, TypeError):
        return True

def fetch_usage_data(access_token):
    """Fetch usage data from Anthropic API."""
    request = Request(API_ENDPOINT)
    request.add_header('Authorization', f'Bearer {access_token}')
    request.add_header('Accept', 'application/json')
    request.add_header('Content-Type', 'application/json')
    request.add_header('User-Agent', 'claude-usage-waybar/1.0')
    request.add_header('anthropic-beta', 'oauth-2025-04-20')

    try:
        with urlopen(request, timeout=REQUEST_TIMEOUT) as response:
            data = json.loads(response.read().decode())
            return data
    except HTTPError as e:
        if e.code == 401:
            return None  # Auth failed
        raise
    except (URLError, Exception):
        return None

def calculate_time_elapsed_percentage(reset_time_str, window_hours):
    """Calculate what percentage of the time window has elapsed."""
    reset_time = datetime.fromisoformat(reset_time_str.replace('Z', '+00:00'))
    current_time = datetime.now(timezone.utc)

    window_start = reset_time - timedelta(hours=window_hours)
    elapsed = current_time - window_start
    total = timedelta(hours=window_hours)

    percentage = (elapsed.total_seconds() / total.total_seconds()) * 100
    return max(0, min(100, percentage))  # Clamp to 0-100

def calculate_difference(utilization, reset_time_str, window_hours):
    """Calculate the difference between usage rate and time elapsed rate."""
    usage_percentage = utilization
    time_percentage = calculate_time_elapsed_percentage(reset_time_str, window_hours)
    difference = usage_percentage - time_percentage
    return usage_percentage, time_percentage, difference

def format_waybar_output(weekly_data, session_data):
    """Format usage data as Waybar-compatible JSON."""
    # Calculate differences
    weekly_usage, weekly_time, weekly_diff = calculate_difference(
        weekly_data['utilization'],
        weekly_data['resets_at'],
        24 * 7  # 7 days in hours
    )

    session_usage, session_time, session_diff = calculate_difference(
        session_data['utilization'],
        session_data['resets_at'],
        5  # 5 hours
    )

    # Format text output
    weekly_sign = '+' if weekly_diff >= 0 else ''
    session_sign = '+' if session_diff >= 0 else ''
    text = f"W: {weekly_sign}{weekly_diff:.1f}% | S: {session_sign}{session_diff:.1f}%"

    # Format tooltip with detailed breakdown
    weekly_reset = datetime.fromisoformat(weekly_data['resets_at'].replace('Z', '+00:00')).astimezone()
    session_reset = datetime.fromisoformat(session_data['resets_at'].replace('Z', '+00:00')).astimezone()

    tooltip = (
        f"Weekly: {weekly_usage:.1f}% used ({weekly_time:.1f}% through week) = {weekly_sign}{weekly_diff:.1f}%\n"
        f"Session: {session_usage:.1f}% used ({session_time:.1f}% through 5h) = {session_sign}{session_diff:.1f}%\n"
        f"\n"
        f"Weekly resets: {weekly_reset.strftime('%Y-%m-%d %H:%M')}\n"
        f"Session resets: {session_reset.strftime('%Y-%m-%d %H:%M')}"
    )

    return {
        "text": text,
        "tooltip": tooltip,
        "class": "claude-usage",
        "percentage": int(weekly_usage)
    }

def error_output(text, tooltip):
    """Return error state as Waybar JSON."""
    return {
        "text": text,
        "tooltip": tooltip,
        "class": "claude-usage-error"
    }

def main():
    """Main entry point."""
    # Read credentials
    access_token, expires_at = read_credentials()
    if not access_token:
        output = error_output("No creds", "Claude credentials not found or invalid")
        print(json.dumps(output))
        return

    # Check if token is expired
    if is_token_expired(expires_at):
        output = error_output("Token expired", "OAuth token has expired")
        print(json.dumps(output))
        return

    # Fetch usage data
    try:
        usage_data = fetch_usage_data(access_token)
    except Exception:
        output = error_output("API error", "Failed to fetch usage data")
        print(json.dumps(output))
        return

    if usage_data is None:
        output = error_output("Auth failed", "Token expired or invalid")
        print(json.dumps(output))
        return

    # Validate response structure
    if 'seven_day' not in usage_data or 'five_hour' not in usage_data:
        output = error_output("API error", "Invalid response structure")
        print(json.dumps(output))
        return

    # Format and output
    try:
        output = format_waybar_output(usage_data['seven_day'], usage_data['five_hour'])
        print(json.dumps(output))
    except Exception as e:
        output = error_output("Error", f"Failed to process data: {str(e)}")
        print(json.dumps(output))

if __name__ == '__main__':
    main()