#!/usr/bin/env python3
"""
Transform cost.csv to add TOTAL rows for each API key.
"""

import csv
import sys
from collections import defaultdict

# Default paths
input_file = 'output/usage/cost.csv'
output_file = 'output/usage/cost_with_totals.csv'

# Allow command line arguments
if len(sys.argv) >= 2:
    input_file = sys.argv[1]
if len(sys.argv) >= 3:
    output_file = sys.argv[2]

# Read the CSV
with open(input_file, 'r') as f:
    reader = csv.DictReader(f)
    rows = list(reader)
    fieldnames = reader.fieldnames

# Group rows by api_key_id
api_key_groups = defaultdict(list)
for row in rows:
    api_key_groups[row['api_key_id']].append(row)

# Write output with totals
with open(output_file, 'w', newline='') as f:
    writer = csv.DictWriter(f, fieldnames=fieldnames)
    writer.writeheader()

    for api_key_id in sorted(api_key_groups.keys()):
        group_rows = api_key_groups[api_key_id]

        # Initialize totals
        totals = {
            'total_input_tokens': 0,
            'total_cache_write_tokens': 0,
            'total_cache_read_tokens': 0,
            'total_output_tokens': 0,
            'total_web_search_requests': 0,
            'total_input_cost_usd': 0.0,
            'total_cache_write_cost_usd': 0.0,
            'total_cache_read_cost_usd': 0.0,
            'total_output_cost_usd': 0.0,
            'total_cost_usd': 0.0,
        }

        # Write original rows and accumulate totals
        for row in group_rows:
            writer.writerow(row)

            # Accumulate totals
            totals['total_input_tokens'] += int(row['total_input_tokens'])
            totals['total_cache_write_tokens'] += int(row['total_cache_write_tokens'])
            totals['total_cache_read_tokens'] += int(row['total_cache_read_tokens'])
            totals['total_output_tokens'] += int(row['total_output_tokens'])
            totals['total_web_search_requests'] += int(row['total_web_search_requests'])
            totals['total_input_cost_usd'] += float(row['total_input_cost_usd'])
            totals['total_cache_write_cost_usd'] += float(row['total_cache_write_cost_usd'])
            totals['total_cache_read_cost_usd'] += float(row['total_cache_read_cost_usd'])
            totals['total_output_cost_usd'] += float(row['total_output_cost_usd'])
            totals['total_cost_usd'] += float(row['total_cost_usd'])

        # Write TOTAL row
        total_row = {
            'api_key_id': row['api_key_id'],
            'api_key_name': row['api_key_name'],
            'workspace_id': row['workspace_id'],
            'status': row['status'],
            'partial_key_hint': row['partial_key_hint'],
            'time_period': 'TOTAL',
            'total_input_tokens': totals['total_input_tokens'],
            'total_cache_write_tokens': totals['total_cache_write_tokens'],
            'total_cache_read_tokens': totals['total_cache_read_tokens'],
            'total_output_tokens': totals['total_output_tokens'],
            'total_web_search_requests': totals['total_web_search_requests'],
            'total_input_cost_usd': f'{totals["total_input_cost_usd"]:.4f}',
            'total_cache_write_cost_usd': f'{totals["total_cache_write_cost_usd"]:.4f}',
            'total_cache_read_cost_usd': f'{totals["total_cache_read_cost_usd"]:.4f}',
            'total_output_cost_usd': f'{totals["total_output_cost_usd"]:.4f}',
            'total_cost_usd': f'{totals["total_cost_usd"]:.4f}',
        }
        writer.writerow(total_row)

print(f'Transformed {input_file} -> {output_file}')
print(f'Processed {len(api_key_groups)} API keys')