#!/usr/bin/env python3
"""
Claude API Usage Cost Fetcher - Simple JSON dumper with pagination
"""

import argparse
import json
import os
import sys
import urllib.error
import urllib.parse
import urllib.request
from datetime import datetime, timedelta, timezone


def fetch_data(endpoint, params, admin_key):
    """Fetch data from Claude API with pagination."""
    base_url = 'https://api.anthropic.com'

    # Build query string
    query_parts = []
    for key, value in params.items():
        if isinstance(value, list):
            for item in value:
                query_parts.append(f'{urllib.parse.quote(key)}={urllib.parse.quote(str(item))}')
        else:
            query_parts.append(f'{urllib.parse.quote(key)}={urllib.parse.quote(str(value))}')

    query_string = '&'.join(query_parts)
    url = f'{base_url}{endpoint}?{query_string}'

    headers = {'anthropic-version': '2023-06-01', 'x-api-key': admin_key}

    all_data = None

    while url:
        req = urllib.request.Request(url, headers=headers)

        try:
            with urllib.request.urlopen(req) as response:
                data = json.loads(response.read().decode())

                if all_data is None:
                    all_data = data
                else:
                    if 'data' in data:
                        all_data['data'].extend(data['data'])
                    all_data['has_more'] = data.get('has_more', False)
                    all_data['next_page'] = data.get('next_page')

                if data.get('has_more') and data.get('next_page'):
                    next_page = data['next_page']
                    base_url_only = url.split('?')[0] if '?' in url else url
                    url = f'{base_url_only}?{query_string}&page={urllib.parse.quote(next_page)}'
                else:
                    url = None

        except urllib.error.HTTPError as e:
            error_body = e.read().decode() if e.fp else 'No error details'
            raise Exception(f'API Error {e.code}: {error_body}')

    return all_data


def main():
    parser = argparse.ArgumentParser(description='Fetch Claude API usage/cost data')

    parser.add_argument('--admin-key', help='Admin API key (or use CLAUDE_ADMIN_KEY env var)')
    parser.add_argument('--report-type', choices=['usage', 'cost'], default='usage')
    parser.add_argument('--api-key-id', help='Filter by API key ID')
    parser.add_argument('--workspace-id', help='Filter by workspace ID')
    parser.add_argument('--days', type=int, default=7, help='Days back (default: 7)')
    parser.add_argument('--start', help='Start timestamp (ISO 8601)')
    parser.add_argument('--end', help='End timestamp (ISO 8601)')
    parser.add_argument('--bucket-width', choices=['1m', '1h', '1d'], default='1d')
    parser.add_argument('--output', help='Output file path')

    args = parser.parse_args()

    admin_key = args.admin_key or os.environ.get('CLAUDE_ADMIN_KEY')
    if not admin_key:
        print('Error: Admin API key required via --admin-key or CLAUDE_ADMIN_KEY', file=sys.stderr)
        sys.exit(1)

    # Determine time range
    if args.start and args.end:
        starting_at = args.start
        ending_at = args.end
    else:
        ending_at_dt = datetime.now(timezone.utc)
        starting_at_dt = ending_at_dt - timedelta(days=args.days)
        starting_at = starting_at_dt.isoformat().replace('+00:00', 'Z')
        ending_at = ending_at_dt.isoformat().replace('+00:00', 'Z')

    params = {'starting_at': starting_at, 'ending_at': ending_at, 'bucket_width': args.bucket_width}

    if args.api_key_id:
        params['api_key_ids[]'] = [args.api_key_id]
    if args.workspace_id:
        params['workspace_ids[]'] = [args.workspace_id]

    if args.report_type == 'usage':
        if args.api_key_id:
            params['group_by[]'] = ['model', 'api_key_id']
        endpoint = '/v1/organizations/usage_report/messages'
    else:
        endpoint = '/v1/organizations/cost_report'

    try:
        data = fetch_data(endpoint, params, admin_key)

        output = json.dumps(data, indent=2)

        if args.output:
            with open(args.output, 'w') as f:
                f.write(output)
        else:
            print(output)

    except Exception as e:
        print(f'Error: {e}', file=sys.stderr)
        sys.exit(1)


if __name__ == '__main__':
    main()