#!/usr/bin/env python3
"""
Claude API Cost Calculator by API Key
Fetches usage data, groups by API key, enriches with metadata, and calculates costs.
"""

import argparse
import json
import os
import sys
from collections import defaultdict
from datetime import datetime, timedelta, timezone
from pathlib import Path

# Add parent directory to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent))

from src.shared.api_client import fetch_api_data

# Global debug flag (set via --debug command line argument)
DEBUG = False

# Directory for storing raw API responses
RAW_OUTPUT_DIR = Path('output/usage/raw')

# Anthropic's official pricing (per million tokens)
PRICING = {
    # Claude Opus 4.8
    'claude-opus-4-8': {
        'input': 5.00,
        'cache_write_5m': 6.25,
        'cache_write_1h': 10.00,
        'cache_read': 0.50,
        'output': 25.00,
    },
    # Claude Opus 4.7
    'claude-opus-4-7': {
        'input': 5.00,
        'cache_write_5m': 6.25,
        'cache_write_1h': 10.00,
        'cache_read': 0.50,
        'output': 25.00,
    },
    # Claude Opus 4.6
    'claude-opus-4-6': {
        'input': 5.00,
        'cache_write_5m': 6.25,
        'cache_write_1h': 10.00,
        'cache_read': 0.50,
        'output': 25.00,
    },
    # Claude Opus 4.5
    'claude-opus-4-5-20251101': {
        'input': 5.00,
        'cache_write_5m': 6.25,
        'cache_write_1h': 10.00,
        'cache_read': 0.50,
        'output': 25.00,
    },
    # Claude Opus 4.1
    'claude-opus-4-1-20250805': {
        'input': 15.00,
        'cache_write_5m': 18.75,
        'cache_write_1h': 30.00,
        'cache_read': 1.50,
        'output': 75.00,
    },
    'claude-opus-4-1': {
        'input': 15.00,
        'cache_write_5m': 18.75,
        'cache_write_1h': 30.00,
        'cache_read': 1.50,
        'output': 75.00,
    },
    'claude-opus-4': {
        'input': 15.00,
        'cache_write_5m': 18.75,
        'cache_write_1h': 30.00,
        'cache_read': 1.50,
        'output': 75.00,
        'deprecated': True,
    },
    # Claude Sonnet 4.x
    'claude-sonnet-4-6': {
        'input': 3.00,
        'cache_write_5m': 3.75,
        'cache_write_1h': 6.00,
        'cache_read': 0.30,
        'output': 15.00,
    },
    'claude-sonnet-4-5-20250929': {
        'input': 3.00,
        'cache_write_5m': 3.75,
        'cache_write_1h': 6.00,
        'cache_read': 0.30,
        'output': 15.00,
    },
    'claude-sonnet-4-20250514': {
        'input': 3.00,
        'cache_write_5m': 3.75,
        'cache_write_1h': 6.00,
        'cache_read': 0.30,
        'output': 15.00,
        'deprecated': True,
    },
    'claude-sonnet-3-7': {
        'input': 3.00,
        'cache_write_5m': 3.75,
        'cache_write_1h': 6.00,
        'cache_read': 0.30,
        'output': 15.00,
        'deprecated': True,
    },
    'claude-3-5-sonnet-20241022': {
        'input': 3.00,
        'cache_write_5m': 3.75,
        'cache_write_1h': 6.00,
        'cache_read': 0.30,
        'output': 15.00,
        'deprecated': True,
    },
    # Claude Haiku 4.x and 3.x
    'claude-haiku-4-5-20251001': {
        'input': 1.00,
        'cache_write_5m': 1.25,
        'cache_write_1h': 2.00,
        'cache_read': 0.10,
        'output': 5.00,
    },
    'claude-haiku-4-5': {
        'input': 1.00,
        'cache_write_5m': 1.25,
        'cache_write_1h': 2.00,
        'cache_read': 0.10,
        'output': 5.00,
    },
    'claude-3-5-haiku-20241022': {
        'input': 0.80,
        'cache_write_5m': 1.00,
        'cache_write_1h': 1.60,
        'cache_read': 0.08,
        'output': 4.00,
        'deprecated': True,
    },
    'claude-haiku-3': {
        'input': 0.25,
        'cache_write_5m': 0.30,
        'cache_write_1h': 0.50,
        'cache_read': 0.03,
        'output': 1.25,
        'deprecated': True,
    },
    # Claude Opus 3 (deprecated)
    'claude-opus-3': {
        'input': 15.00,
        'cache_write_5m': 18.75,
        'cache_write_1h': 30.00,
        'cache_read': 1.50,
        'output': 75.00,
        'deprecated': True,
    },
}


def fetch_usage_for_date(admin_key: str, date_str: str, bucket_width: str = '1d') -> dict:
    """
    Fetch usage data for a single date.

    Args:
        admin_key: Admin API key
        date_str: Date in YYYY-MM-DD format
        bucket_width: Bucket width (1m, 1h, 1d)

    Returns:
        Usage data for that date
    """
    # Convert date to full day range (ending_at should be start of next day)
    dt = datetime.fromisoformat(date_str)
    next_day = dt + timedelta(days=1)

    starting_at = f'{date_str}T00:00:00Z'
    ending_at = next_day.strftime('%Y-%m-%dT00:00:00Z')

    endpoint = '/v1/organizations/usage_report/messages'
    params = {
        'starting_at': starting_at,
        'ending_at': ending_at,
        'bucket_width': bucket_width,
        'group_by[]': ['api_key_id', 'model'],
    }

    if DEBUG:
        print(f'DEBUG: Fetching usage data for {date_str}', file=sys.stderr)

    # Fetch data using shared client (handles pagination automatically)
    response = fetch_api_data(endpoint, params, admin_key, debug=DEBUG)

    # Save raw response to file
    RAW_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    filename = f'messages_{date_str}.json'
    output_path = RAW_OUTPUT_DIR / filename

    with open(output_path, 'w') as f:
        json.dump(response, f, indent=2)

    if DEBUG:
        print(f'DEBUG: Saved raw usage data to {output_path}', file=sys.stderr)

    return response


def fetch_usage_data(admin_key: str, starting_at: str, ending_at: str, bucket_width: str = '1d') -> dict:
    """
    Fetch usage data for a date range by fetching each day individually.

    Args:
        admin_key: Admin API key
        starting_at: Start timestamp (ISO 8601)
        ending_at: End timestamp (ISO 8601)
        bucket_width: Bucket width (1m, 1h, 1d)

    Returns:
        Merged usage data from all days
    """
    # Parse dates (extract YYYY-MM-DD from ISO 8601 timestamp)
    start_date = starting_at.split('T')[0]
    end_date = ending_at.split('T')[0]

    start_dt = datetime.fromisoformat(start_date)
    end_dt = datetime.fromisoformat(end_date)

    if DEBUG:
        print(f'DEBUG: Fetching usage data from {start_date} to {end_date}', file=sys.stderr)

    # Fetch data day-by-day
    merged_data: dict = {'data': []}
    current_dt = start_dt

    while current_dt <= end_dt:
        date_str = current_dt.strftime('%Y-%m-%d')
        response = fetch_usage_for_date(admin_key, date_str, bucket_width)

        # Merge buckets from this day
        if 'data' in response:
            merged_data['data'].extend(response['data'])

        current_dt += timedelta(days=1)

    if DEBUG:
        print(f'DEBUG: Merged {len(merged_data["data"])} total bucket(s) from all days', file=sys.stderr)

    return merged_data


def fetch_users(admin_key: str) -> dict:
    """
    Fetch all users in the organization.

    Args:
        admin_key: Admin API key

    Returns:
        Dict mapping user_id to user metadata
    """
    endpoint = '/v1/organizations/users'
    params = {}

    if DEBUG:
        print('DEBUG: Fetching users...', file=sys.stderr)

    # Fetch data using shared client (handles pagination automatically)
    response = fetch_api_data(endpoint, params, admin_key, debug=DEBUG)

    # Save raw response to file
    RAW_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    output_path = RAW_OUTPUT_DIR / 'users.json'

    with open(output_path, 'w') as f:
        json.dump(response, f, indent=2)

    if DEBUG:
        print(f'DEBUG: Saved raw users data to {output_path}', file=sys.stderr)

    # Parse into dict
    users = {}
    for user_data in response.get('data', []):
        users[user_data['id']] = {
            'name': user_data.get('name', 'Unknown'),
            'email': user_data.get('email', 'Unknown'),
            'role': user_data.get('role', 'user'),
            'added_at': user_data.get('added_at'),
        }

    return users


def fetch_api_keys_metadata(admin_key: str) -> dict:
    """
    Fetch metadata for all API keys.

    Args:
        admin_key: Admin API key

    Returns:
        Dict mapping api_key_id to API key metadata
    """
    endpoint = '/v1/organizations/api_keys'
    params = {}

    if DEBUG:
        print('DEBUG: Fetching API keys...', file=sys.stderr)

    # Fetch data using shared client (handles pagination automatically)
    response = fetch_api_data(endpoint, params, admin_key, debug=DEBUG)

    # Save raw response to file
    RAW_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    output_path = RAW_OUTPUT_DIR / 'api_keys.json'

    with open(output_path, 'w') as f:
        json.dump(response, f, indent=2)

    if DEBUG:
        print(f'DEBUG: Saved raw API keys data to {output_path}', file=sys.stderr)

    # Parse into dict
    api_keys = {}
    for key_data in response.get('data', []):
        api_keys[key_data['id']] = {
            'name': key_data.get('name', 'Unnamed'),
            'created_at': key_data.get('created_at'),
            'created_by': key_data.get('created_by', {}).get('id', 'Unknown'),
            'workspace_id': key_data.get('workspace_id'),
            'status': key_data.get('status', 'unknown'),
            'partial_key_hint': key_data.get('partial_key_hint', ''),
        }

    return api_keys


def fetch_workspaces(admin_key: str) -> dict:
    """
    Fetch all workspaces in the organization.

    Args:
        admin_key: Admin API key

    Returns:
        Dict mapping workspace_id to workspace metadata
    """
    endpoint = '/v1/organizations/workspaces'
    params = {}

    if DEBUG:
        print('DEBUG: Fetching workspaces...', file=sys.stderr)

    # Fetch data using shared client (handles pagination automatically)
    response = fetch_api_data(endpoint, params, admin_key, debug=DEBUG)

    # Save raw response to file
    RAW_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    output_path = RAW_OUTPUT_DIR / 'workspaces.json'

    with open(output_path, 'w') as f:
        json.dump(response, f, indent=2)

    if DEBUG:
        print(f'DEBUG: Saved raw workspaces data to {output_path}', file=sys.stderr)

    # Parse into dict
    workspaces = {}
    for workspace_data in response.get('data', []):
        workspaces[workspace_data['id']] = {
            'name': workspace_data.get('name', 'Default'),
            'display_color': workspace_data.get('display_color', '#9B87F5'),
            'created_at': workspace_data.get('created_at'),
            'archived_at': workspace_data.get('archived_at'),
            'type': workspace_data.get('type', 'workspace'),
        }

    return workspaces


def enrich_api_keys_with_user_info(api_keys_metadata, users_data):
    """Enrich API key metadata with user name, email, id, role, and added_at from created_by field."""
    for api_key_id, metadata in api_keys_metadata.items():
        created_by_id = metadata.get('created_by', 'Unknown')
        if created_by_id in users_data:
            user_info = users_data[created_by_id]
            metadata['created_by_name'] = user_info.get('name', 'Unknown')
            metadata['created_by_email'] = user_info.get('email', 'Unknown')
            metadata['created_by_user_id'] = created_by_id
            metadata['created_by_role'] = user_info.get('role', 'unknown')
            metadata['created_by_added_at'] = user_info.get('added_at')
        else:
            metadata['created_by_name'] = 'Unknown'
            metadata['created_by_email'] = 'Unknown'
            metadata['created_by_user_id'] = created_by_id
            metadata['created_by_role'] = 'unknown'
            metadata['created_by_added_at'] = None

    return api_keys_metadata


def enrich_api_keys_with_workspace_info(api_keys_metadata, workspaces_data):
    """Enrich API key metadata with workspace name and color from workspace_id field."""
    for api_key_id, metadata in api_keys_metadata.items():
        workspace_id = metadata.get('workspace_id')
        if workspace_id and workspace_id in workspaces_data:
            workspace_info = workspaces_data[workspace_id]
            metadata['workspace_name'] = workspace_info.get('name', 'Default')
            metadata['workspace_color'] = workspace_info.get('display_color', '#9B87F5')
        else:
            metadata['workspace_name'] = 'Default'
            metadata['workspace_color'] = '#9B87F5'

    return api_keys_metadata


def normalize_time_period(timestamp, time_grouping):
    """Normalize a timestamp to a time period string."""
    if not timestamp:
        return 'unknown'

    try:
        dt = datetime.fromisoformat(timestamp.replace('Z', '+00:00'))

        if time_grouping == 'day':
            return dt.strftime('%Y-%m-%d')
        elif time_grouping == 'month':
            return dt.strftime('%Y-%m')
        else:
            return 'total'
    except Exception:
        return 'unknown'


def aggregate_usage(usage_data, time_grouping=None):
    """Aggregate usage data by API key and model, optionally grouped by time period.

    Args:
        usage_data: Raw usage data from the API
        time_grouping: 'day', 'month', or None for total aggregation

    Returns:
        If time_grouping is None: dict[api_key][model] -> usage_dict
        If time_grouping is set: dict[api_key][time_period][model] -> usage_dict
    """
    if time_grouping:
        # Three-level structure: api_key -> time_period -> model -> usage
        api_key_usage = defaultdict(
            lambda: defaultdict(
                lambda: defaultdict(
                    lambda: {
                        'input': 0,
                        'cache_write': 0,
                        'cache_creation_1h': 0,
                        'cache_read': 0,
                        'output': 0,
                        'web_search_requests': 0,
                    }
                )
            )
        )
    else:
        # Two-level structure: api_key -> model -> usage
        api_key_usage = defaultdict(
            lambda: defaultdict(
                lambda: {
                    'input': 0,
                    'cache_write': 0,
                    'cache_creation_1h': 0,
                    'cache_read': 0,
                    'output': 0,
                    'web_search_requests': 0,
                }
            )
        )

    for bucket in usage_data.get('data', []):
        bucket_time = bucket.get('starting_at')
        time_period = normalize_time_period(bucket_time, time_grouping) if time_grouping else None

        for result in bucket.get('results', []):
            api_key_id = result.get('api_key_id')
            model = result.get('model')

            if not api_key_id or not model:
                continue

            if time_grouping:
                usage = api_key_usage[api_key_id][time_period][model]
            else:
                usage = api_key_usage[api_key_id][model]

            usage['input'] += result.get('uncached_input_tokens', 0)
            usage['cache_write'] += result['cache_creation'].get('ephemeral_5m_input_tokens', 0)
            usage['cache_creation_1h'] += result['cache_creation'].get('ephemeral_1h_input_tokens', 0)
            usage['cache_read'] += result.get('cache_read_input_tokens', 0)
            usage['output'] += result.get('output_tokens', 0)
            usage['web_search_requests'] += result.get('server_tool_use', {}).get('web_search_requests', 0)

    return api_key_usage


def calculate_cost(tokens, model):
    """Calculate cost for a given token usage and model."""
    if model not in PRICING:
        return None

    pricing = PRICING[model]

    input_cost = (tokens['input'] / 1_000_000) * pricing['input']
    cache_write_cost = (tokens['cache_write'] / 1_000_000) * pricing['cache_write_5m']
    expensive_cache_creation_cost = (tokens.get('cache_creation_1h', 0) / 1_000_000) * pricing['cache_write_1h']
    cache_read_cost = (tokens['cache_read'] / 1_000_000) * pricing['cache_read']
    output_cost = (tokens['output'] / 1_000_000) * pricing['output']

    return {
        'input': input_cost,
        'cache_write': cache_write_cost,
        'expensive_cache_creation': expensive_cache_creation_cost,
        'cache_read': cache_read_cost,
        'output': output_cost,
        'total': input_cost + cache_write_cost + expensive_cache_creation_cost + cache_read_cost + output_cost,
    }


def format_report(api_key_usage, api_keys_metadata, output_format='text', time_grouping=None):
    """Format the cost report."""
    if output_format == 'json':
        return format_json_report(api_key_usage, api_keys_metadata, time_grouping)
    elif output_format == 'csv':
        return format_csv_report(api_key_usage, api_keys_metadata, time_grouping)
    else:
        return format_text_report(api_key_usage, api_keys_metadata, time_grouping)


def format_json_report(api_key_usage, api_keys_metadata, time_grouping=None):
    """Format report as JSON."""
    report = []
    total_org_cost = 0

    for api_key_id in sorted(api_key_usage.keys()):
        metadata = api_keys_metadata.get(api_key_id, {})

        key_report = {
            'api_key_id': api_key_id,
            'name': metadata.get('name', 'Unknown'),
            'created_by_name': metadata.get('created_by_name', 'Unknown'),
            'created_by_email': metadata.get('created_by_email', 'Unknown'),
            'workspace_id': metadata.get('workspace_id'),
            'workspace_name': metadata.get('workspace_name', 'Default'),
            'workspace_color': metadata.get('workspace_color', '#9B87F5'),
            'status': metadata.get('status'),
            'partial_key_hint': metadata.get('partial_key_hint'),
            'total_cost': 0,
        }

        if time_grouping:
            # Three-level structure with time periods
            key_report['time_periods'] = []

            for time_period in sorted(api_key_usage[api_key_id].keys()):
                period_report = {'time_period': time_period, 'models': [], 'period_cost': 0}

                for model, tokens in sorted(api_key_usage[api_key_id][time_period].items()):
                    cost = calculate_cost(tokens, model)

                    if cost:
                        model_report = {'model': model, 'tokens': tokens, 'cost': cost}
                        period_report['models'].append(model_report)
                        period_report['period_cost'] += cost['total']

                key_report['time_periods'].append(period_report)
                key_report['total_cost'] += period_report['period_cost']
        else:
            # Two-level structure without time periods
            key_report['models'] = []

            for model, tokens in api_key_usage[api_key_id].items():
                cost = calculate_cost(tokens, model)

                if cost:
                    model_report = {'model': model, 'tokens': tokens, 'cost': cost}
                    key_report['models'].append(model_report)
                    key_report['total_cost'] += cost['total']

        total_org_cost += key_report['total_cost']
        report.append(key_report)

    result = {'api_keys': report, 'total_organization_cost': total_org_cost}

    if time_grouping:
        result['time_grouping'] = time_grouping

    return json.dumps(result, indent=2)


def format_csv_report(api_key_usage, api_keys_metadata, time_grouping=None):
    """Format report as CSV.

    Without time_grouping: one row per API key (aggregated across models)
    With time_grouping: one row per API key per time period
    """
    import csv
    from io import StringIO

    output = StringIO()
    writer = csv.writer(output)

    # Write header
    if time_grouping:
        writer.writerow(
            [
                'api_key_id',
                'api_key_name',
                'created_by_name',
                'created_by_email',
                'workspace_id',
                'status',
                'partial_key_hint',
                'time_period',
                'total_input_tokens',
                'total_cache_write_tokens',
                'total_cache_read_tokens',
                'total_output_tokens',
                'total_web_search_requests',
                'total_input_cost_usd',
                'total_cache_write_cost_usd',
                'total_cache_read_cost_usd',
                'total_output_cost_usd',
                'total_cost_usd',
            ]
        )
    else:
        writer.writerow(
            [
                'api_key_id',
                'api_key_name',
                'created_by_name',
                'created_by_email',
                'workspace_id',
                'status',
                'partial_key_hint',
                'total_input_tokens',
                'total_cache_write_tokens',
                'total_cache_read_tokens',
                'total_output_tokens',
                'total_web_search_requests',
                'total_input_cost_usd',
                'total_cache_write_cost_usd',
                'total_cache_read_cost_usd',
                'total_output_cost_usd',
                'total_cost_usd',
            ]
        )

    for api_key_id in sorted(api_key_usage.keys()):
        metadata = api_keys_metadata.get(api_key_id, {})

        if time_grouping:
            # Track totals across all time periods for this API key
            api_key_total_tokens = {
                'input': 0,
                'cache_write': 0,
                'cache_read': 0,
                'output': 0,
                'web_search_requests': 0,
            }
            api_key_total_costs = {
                'input': 0,
                'cache_write': 0,
                'cache_read': 0,
                'output': 0,
                'total': 0,
            }

            # Write one row per time period
            for time_period in sorted(api_key_usage[api_key_id].keys()):
                # Aggregate across all models for this time period
                period_tokens = {
                    'input': 0,
                    'cache_write': 0,
                    'cache_read': 0,
                    'output': 0,
                    'web_search_requests': 0,
                }
                period_costs = {
                    'input': 0,
                    'cache_write': 0,
                    'cache_read': 0,
                    'output': 0,
                    'total': 0,
                }

                for model, tokens in api_key_usage[api_key_id][time_period].items():
                    period_tokens['input'] += tokens['input']
                    period_tokens['cache_write'] += tokens['cache_write']
                    period_tokens['cache_read'] += tokens['cache_read']
                    period_tokens['output'] += tokens['output']
                    period_tokens['web_search_requests'] += tokens['web_search_requests']

                    cost = calculate_cost(tokens, model)
                    if cost:
                        period_costs['input'] += cost['input']
                        period_costs['cache_write'] += cost['cache_write']
                        period_costs['cache_read'] += cost['cache_read']
                        period_costs['output'] += cost['output']
                        period_costs['total'] += cost['total']

                writer.writerow(
                    [
                        api_key_id,
                        metadata.get('name', 'Unknown'),
                        metadata.get('created_by_name', 'Unknown'),
                        metadata.get('created_by_email', 'Unknown'),
                        metadata.get('workspace_id', ''),
                        metadata.get('status', 'unknown'),
                        metadata.get('partial_key_hint', ''),
                        time_period,
                        period_tokens['input'],
                        period_tokens['cache_write'],
                        period_tokens['cache_read'],
                        period_tokens['output'],
                        period_tokens['web_search_requests'],
                        f'{period_costs["input"]:.4f}',
                        f'{period_costs["cache_write"]:.4f}',
                        f'{period_costs["cache_read"]:.4f}',
                        f'{period_costs["output"]:.4f}',
                        f'{period_costs["total"]:.4f}',
                    ]
                )

                # Accumulate totals for this API key
                api_key_total_tokens['input'] += period_tokens['input']
                api_key_total_tokens['cache_write'] += period_tokens['cache_write']
                api_key_total_tokens['cache_read'] += period_tokens['cache_read']
                api_key_total_tokens['output'] += period_tokens['output']
                api_key_total_tokens['web_search_requests'] += period_tokens['web_search_requests']

                api_key_total_costs['input'] += period_costs['input']
                api_key_total_costs['cache_write'] += period_costs['cache_write']
                api_key_total_costs['cache_read'] += period_costs['cache_read']
                api_key_total_costs['output'] += period_costs['output']
                api_key_total_costs['total'] += period_costs['total']

            # Write total row for this API key
            writer.writerow(
                [
                    api_key_id,
                    metadata.get('name', 'Unknown'),
                    metadata.get('created_by_name', 'Unknown'),
                    metadata.get('created_by_email', 'Unknown'),
                    metadata.get('workspace_id', ''),
                    metadata.get('status', 'unknown'),
                    metadata.get('partial_key_hint', ''),
                    'TOTAL',
                    api_key_total_tokens['input'],
                    api_key_total_tokens['cache_write'],
                    api_key_total_tokens['cache_read'],
                    api_key_total_tokens['output'],
                    api_key_total_tokens['web_search_requests'],
                    f'{api_key_total_costs["input"]:.4f}',
                    f'{api_key_total_costs["cache_write"]:.4f}',
                    f'{api_key_total_costs["cache_read"]:.4f}',
                    f'{api_key_total_costs["output"]:.4f}',
                    f'{api_key_total_costs["total"]:.4f}',
                ]
            )
        else:
            # Write one row per API key (aggregate across all models)
            total_tokens = {
                'input': 0,
                'cache_write': 0,
                'cache_read': 0,
                'output': 0,
                'web_search_requests': 0,
            }
            total_costs = {'input': 0, 'cache_write': 0, 'cache_read': 0, 'output': 0, 'total': 0}

            for model, tokens in api_key_usage[api_key_id].items():
                total_tokens['input'] += tokens['input']
                total_tokens['cache_write'] += tokens['cache_write']
                total_tokens['cache_read'] += tokens['cache_read']
                total_tokens['output'] += tokens['output']
                total_tokens['web_search_requests'] += tokens['web_search_requests']

                cost = calculate_cost(tokens, model)
                if cost:
                    total_costs['input'] += cost['input']
                    total_costs['cache_write'] += cost['cache_write']
                    total_costs['cache_read'] += cost['cache_read']
                    total_costs['output'] += cost['output']
                    total_costs['total'] += cost['total']

            writer.writerow(
                [
                    api_key_id,
                    metadata.get('name', 'Unknown'),
                    metadata.get('created_by_name', 'Unknown'),
                    metadata.get('created_by_email', 'Unknown'),
                    metadata.get('workspace_id', ''),
                    metadata.get('status', 'unknown'),
                    metadata.get('partial_key_hint', ''),
                    total_tokens['input'],
                    total_tokens['cache_write'],
                    total_tokens['cache_read'],
                    total_tokens['output'],
                    total_tokens['web_search_requests'],
                    f'{total_costs["input"]:.4f}',
                    f'{total_costs["cache_write"]:.4f}',
                    f'{total_costs["cache_read"]:.4f}',
                    f'{total_costs["output"]:.4f}',
                    f'{total_costs["total"]:.4f}',
                ]
            )

    return output.getvalue()


def format_text_report(api_key_usage, api_keys_metadata, time_grouping=None):
    """Format report as human-readable text."""
    lines = []
    lines.append('=' * 100)
    if time_grouping:
        lines.append(f'CLAUDE API COST REPORT BY API KEY (Grouped by {time_grouping.upper()})')
    else:
        lines.append('CLAUDE API COST REPORT BY API KEY')
    lines.append('=' * 100)
    lines.append('')

    total_org_cost = 0

    for api_key_id in sorted(api_key_usage.keys()):
        metadata = api_keys_metadata.get(api_key_id, {})

        lines.append(f'API Key: {api_key_id}')
        lines.append(f'  Name: {metadata.get("name", "Unknown")}')
        lines.append(f'  Created By: {metadata.get("created_by_name", "Unknown")} ({metadata.get("created_by_email", "Unknown")})')
        lines.append(f'  Status: {metadata.get("status", "unknown")}')
        lines.append(f'  Workspace: {metadata.get("workspace_id", "N/A")}')
        lines.append(f'  Hint: {metadata.get("partial_key_hint", "N/A")}')
        lines.append('-' * 100)

        api_key_total = 0

        if time_grouping:
            # Three-level structure: time_period -> model -> usage
            for time_period in sorted(api_key_usage[api_key_id].keys()):
                lines.append(f'\n  {time_grouping.upper()}: {time_period}')
                lines.append('  ' + '-' * 96)

                period_total = 0

                for model, tokens in sorted(api_key_usage[api_key_id][time_period].items()):
                    cost = calculate_cost(tokens, model)

                    if not cost:
                        lines.append(f'    Model: {model} (pricing not available)')
                        continue

                    lines.append(f'    Model: {model}')
                    lines.append(f'      Input tokens:       {tokens["input"]:>15,} = ${cost["input"]:>10.4f}')
                    lines.append(f'      Cache write:        {tokens["cache_write"]:>15,} = ${cost["cache_write"]:>10.4f}')
                    lines.append(f'      Cache read:         {tokens["cache_read"]:>15,} = ${cost["cache_read"]:>10.4f}')
                    lines.append(f'      Output tokens:      {tokens["output"]:>15,} = ${cost["output"]:>10.4f}')

                    if tokens['web_search_requests'] > 0:
                        lines.append(f'      Web searches:       {tokens["web_search_requests"]:>15,}')

                    lines.append(f'      Model subtotal:                          ${cost["total"]:>10.4f}')
                    lines.append('')

                    period_total += cost['total']

                lines.append(f'    {time_period} Total: ${period_total:>10.4f}')
                lines.append('')
                api_key_total += period_total
        else:
            # Two-level structure: model -> usage
            for model, tokens in sorted(api_key_usage[api_key_id].items()):
                cost = calculate_cost(tokens, model)

                if not cost:
                    lines.append(f'  Model: {model} (pricing not available)')
                    continue

                lines.append(f'  Model: {model}')
                lines.append(f'    Input tokens:       {tokens["input"]:>15,} = ${cost["input"]:>10.4f}')
                lines.append(f'    Cache write:        {tokens["cache_write"]:>15,} = ${cost["cache_write"]:>10.4f}')
                lines.append(f'    Cache read:         {tokens["cache_read"]:>15,} = ${cost["cache_read"]:>10.4f}')
                lines.append(f'    Output tokens:      {tokens["output"]:>15,} = ${cost["output"]:>10.4f}')

                if tokens['web_search_requests'] > 0:
                    lines.append(f'    Web searches:       {tokens["web_search_requests"]:>15,}')

                lines.append(f'    Model subtotal:                          ${cost["total"]:>10.4f}')
                lines.append('')

                api_key_total += cost['total']

        lines.append(f'  API Key Total: ${api_key_total:>10.4f}')
        lines.append('')
        lines.append('')

        total_org_cost += api_key_total

    lines.append('=' * 100)
    lines.append(f'ORGANIZATION TOTAL: ${total_org_cost:>10.4f}')
    lines.append('=' * 100)

    return '\n'.join(lines)


def build_cost_data_structure(api_key_usage: dict, api_keys_metadata: dict, time_grouping: str) -> dict:
    """
    Build cost data structure from aggregated usage.

    Args:
        api_key_usage: Aggregated usage data by API key
        api_keys_metadata: API key metadata
        time_grouping: Time grouping setting

    Returns:
        Cost data dict ready for JSON export
    """
    api_keys_list = []
    total_org_cost = 0

    for api_key_id in sorted(api_key_usage.keys()):
        metadata = api_keys_metadata.get(api_key_id, {})

        key_data = {
            'api_key_id': api_key_id,
            'name': metadata.get('name', 'Unknown'),
            'created_by_name': metadata.get('created_by_name', 'Unknown'),
            'created_by_email': metadata.get('created_by_email', 'Unknown'),
            'created_by_user_id': metadata.get('created_by_user_id', 'Unknown'),
            'created_by_role': metadata.get('created_by_role', 'unknown'),
            'created_by_added_at': metadata.get('created_by_added_at'),
            'workspace_id': metadata.get('workspace_id'),
            'workspace_name': metadata.get('workspace_name', 'Default'),
            'workspace_color': metadata.get('workspace_color', '#9B87F5'),
            'status': metadata.get('status'),
            'partial_key_hint': metadata.get('partial_key_hint'),
            'total_cost': 0,
        }

        if time_grouping:
            # Three-level structure with time periods
            key_data['time_periods'] = []

            for time_period in sorted(api_key_usage[api_key_id].keys()):
                period_data = {'time_period': time_period, 'models': [], 'period_cost': 0}

                for model, tokens in sorted(api_key_usage[api_key_id][time_period].items()):
                    cost = calculate_cost(tokens, model)

                    if cost:
                        model_data = {'model': model, 'tokens': tokens, 'cost': cost}
                        period_data['models'].append(model_data)
                        period_data['period_cost'] += cost['total']

                key_data['time_periods'].append(period_data)
                key_data['total_cost'] += period_data['period_cost']
        else:
            # Two-level structure without time periods
            key_data['models'] = []

            for model, tokens in api_key_usage[api_key_id].items():
                cost = calculate_cost(tokens, model)

                if cost:
                    model_data = {'model': model, 'tokens': tokens, 'cost': cost}
                    key_data['models'].append(model_data)
                    key_data['total_cost'] += cost['total']

        total_org_cost += key_data['total_cost']
        api_keys_list.append(key_data)

    result = {'api_keys': api_keys_list, 'total_organization_cost': total_org_cost}

    if time_grouping:
        result['time_grouping'] = time_grouping

    return result


def main():
    parser = argparse.ArgumentParser(
        description='Calculate Claude API costs grouped by API key',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Last 7 days
  %(prog)s --days 7

  # Specific date range
  %(prog)s --start 2025-09-01T00:00:00Z --end 2025-09-30T23:59:59Z

  # Group by day
  %(prog)s --days 30 --time-grouping day

  # Group by month
  %(prog)s --days 90 --time-grouping month

  # Output as JSON with daily breakdown
  %(prog)s --days 30 --time-grouping day --format json

  # Output as CSV with monthly breakdown
  %(prog)s --days 90 --time-grouping month --format csv --output report.csv

  # Save to file
  %(prog)s --days 30 --output report.txt
        """,
    )

    parser.add_argument('--admin-key', help='Admin API key (or use CLAUDE_ADMIN_KEY env var)')
    parser.add_argument('--days', type=int, default=7, help='Days back from now (default: 7)')
    parser.add_argument('--start', help='Start timestamp (ISO 8601, e.g., 2025-09-01T00:00:00Z)')
    parser.add_argument('--end', help='End timestamp (ISO 8601, e.g., 2025-09-30T23:59:59Z)')
    parser.add_argument(
        '--bucket-width',
        choices=['1m', '1h', '1d'],
        default='1d',
        help='Time bucket width (default: 1d)',
    )
    parser.add_argument(
        '--time-grouping',
        choices=['day', 'month'],
        default=None,
        help='Group usage by time period (day or month). Default: no time grouping',
    )
    parser.add_argument(
        '--format',
        choices=['text', 'json', 'csv'],
        default='text',
        help='Output format (default: text)',
    )
    parser.add_argument('--output', help='Output file path (default: stdout)')
    parser.add_argument(
        '--debug',
        action='store_true',
        help='Enable debug output (shows detailed information about data fetching and processing)',
    )

    args = parser.parse_args()

    # Set global DEBUG flag
    global DEBUG
    DEBUG = args.debug

    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:
        # Calculate date range for the last N days (inclusive)
        # Use current time as ending_at to include partial data for today
        now_dt = datetime.now(timezone.utc)
        ending_at_dt = now_dt  # Use current time to get most recent data
        starting_at_dt = now_dt.replace(hour=0, minute=0, second=0, microsecond=0) - timedelta(days=args.days - 1)  # Start of N days ago
        starting_at = starting_at_dt.isoformat().replace('+00:00', 'Z')
        ending_at = ending_at_dt.isoformat().replace('+00:00', 'Z')

    try:
        # Step 1: Fetch new data
        print(f'Fetching usage data from {starting_at} to {ending_at}...', file=sys.stderr)
        if DEBUG:
            if args.start and args.end:
                print(f'DEBUG: Using explicit date range: {starting_at} to {ending_at}', file=sys.stderr)
            else:
                print(
                    f'DEBUG: Requested {args.days} days. Date range covers: {starting_at} to {ending_at}',
                    file=sys.stderr,
                )
        usage_data = fetch_usage_data(admin_key, starting_at, ending_at, args.bucket_width)

        # Debug: Show what buckets we actually received
        if DEBUG and usage_data and 'data' in usage_data:
            bucket_count = len(usage_data['data'])
            print(f'DEBUG: Received {bucket_count} bucket(s) from API', file=sys.stderr)
            if bucket_count > 0:
                first_bucket = usage_data['data'][0]
                last_bucket = usage_data['data'][-1]
                print(
                    f'DEBUG: First bucket: {first_bucket.get("starting_at")} → {first_bucket.get("ending_at")}',
                    file=sys.stderr,
                )
                print(
                    f'DEBUG: Last bucket: {last_bucket.get("starting_at")} → {last_bucket.get("ending_at")}',
                    file=sys.stderr,
                )

        print('Fetching users...', file=sys.stderr)
        users_data = fetch_users(admin_key)

        print('Fetching API key metadata...', file=sys.stderr)
        api_keys_metadata = fetch_api_keys_metadata(admin_key)
        if DEBUG:
            print(f'DEBUG: Found {len(api_keys_metadata)} API keys in organization', file=sys.stderr)
            for key_id, metadata in api_keys_metadata.items():
                print(
                    f'  - {key_id}: {metadata.get("name")} (status: {metadata.get("status")})',
                    file=sys.stderr,
                )

        print('Fetching workspaces...', file=sys.stderr)
        workspaces_data = fetch_workspaces(admin_key)
        if DEBUG:
            print(f'DEBUG: Found {len(workspaces_data)} workspaces', file=sys.stderr)

        print('Enriching API keys with user information...', file=sys.stderr)
        api_keys_metadata = enrich_api_keys_with_user_info(api_keys_metadata, users_data)

        print('Enriching API keys with workspace information...', file=sys.stderr)
        api_keys_metadata = enrich_api_keys_with_workspace_info(api_keys_metadata, workspaces_data)

        print('Aggregating usage by API key and model...', file=sys.stderr)
        api_key_usage = aggregate_usage(usage_data, args.time_grouping)
        if DEBUG:
            print(f'DEBUG: Found usage data for {len(api_key_usage)} API keys', file=sys.stderr)

        # Show which keys have usage vs. which don't
        if DEBUG:
            keys_with_usage = set(api_key_usage.keys())
            keys_without_usage = set(api_keys_metadata.keys()) - keys_with_usage
            if keys_without_usage:
                print(
                    f'DEBUG: {len(keys_without_usage)} API key(s) have no usage in this time period:',
                    file=sys.stderr,
                )
                for key_id in sorted(keys_without_usage):
                    metadata = api_keys_metadata[key_id]
                    print(
                        f'  - {key_id}: {metadata.get("name")} (status: {metadata.get("status")})',
                        file=sys.stderr,
                    )

        # Format and output report
        print('Calculating costs...', file=sys.stderr)
        report = format_report(api_key_usage, api_keys_metadata, args.format, args.time_grouping)

        if args.output:
            with open(args.output, 'w') as f:
                f.write(report)
            print(f'✅ Report written to {args.output}', file=sys.stderr)
        else:
            print(report)

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


if __name__ == '__main__':
    main()