#!/usr/bin/env python3
"""
Convert cost CSV (with time periods) to JSON format for visualization.
"""

import csv
import json
from collections import defaultdict
from pathlib import Path


def csv_to_json(csv_path: Path, output_path: Path) -> None:
    """Convert CSV to JSON format compatible with view_api_costs.py"""

    with csv_path.open('r') as f:
        reader = csv.DictReader(f)
        rows = list(reader)

    # Group by api_key_id
    api_keys_data = defaultdict(lambda: {'time_periods': [], 'metadata': {}})

    for row in rows:
        api_key_id = row['api_key_id']

        # Skip TOTAL rows
        if row['time_period'] == 'TOTAL':
            continue

        # Store metadata (first occurrence)
        if not api_keys_data[api_key_id]['metadata']:
            api_keys_data[api_key_id]['metadata'] = {
                'api_key_id': api_key_id,
                'name': row['api_key_name'],
                'workspace_id': row['workspace_id'],
                'status': row['status'],
                'partial_key_hint': row['partial_key_hint'],
            }

        # Add time period data
        api_keys_data[api_key_id]['time_periods'].append(
            {
                'time_period': row['time_period'],
                'period_cost': float(row['total_cost_usd']),
                'models': [
                    {
                        'model': 'aggregated',  # CSV doesn't break down by model
                        'tokens': {
                            'input': int(row['total_input_tokens']),
                            'cache_write': int(row['total_cache_write_tokens']),
                            'cache_read': int(row['total_cache_read_tokens']),
                            'output': int(row['total_output_tokens']),
                            'web_search_requests': int(row['total_web_search_requests']),
                        },
                        'cost': {
                            'input': float(row['total_input_cost_usd']),
                            'cache_write': float(row['total_cache_write_cost_usd']),
                            'cache_read': float(row['total_cache_read_cost_usd']),
                            'output': float(row['total_output_cost_usd']),
                            'total': float(row['total_cost_usd']),
                        },
                    }
                ],
            }
        )

    # Build final JSON structure
    api_keys = []
    total_org_cost = 0

    for api_key_id, data in api_keys_data.items():
        metadata = data['metadata']
        time_periods = data['time_periods']

        # Calculate total cost
        total_cost = sum(p['period_cost'] for p in time_periods)
        total_org_cost += total_cost

        api_keys.append({**metadata, 'total_cost': total_cost, 'time_periods': time_periods})

    # Detect time grouping from data
    if api_keys and api_keys[0]['time_periods']:
        first_period = api_keys[0]['time_periods'][0]['time_period']
        if len(first_period) == 10:  # YYYY-MM-DD
            time_grouping = 'day'
        elif len(first_period) == 7:  # YYYY-MM
            time_grouping = 'month'
        else:
            time_grouping = 'unknown'
    else:
        time_grouping = 'unknown'

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

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

    print(f'Converted {csv_path} -> {output_path}')
    print(f'Found {len(api_keys)} API keys')
    print(f'Time grouping: {time_grouping}')
    print(f'Total Spend: ${total_org_cost:.2f}')


if __name__ == '__main__':
    import sys

    if len(sys.argv) < 2:
        print('Usage: python csv_to_json.py <csv_file> [output.json]')
        print('\nExample:')
        print('  python csv_to_json.py output/usage/cost.csv output/usage/costs.json')
        print('\nDefaults:')
        print('  python csv_to_json.py cost.csv')
        print('  -> Outputs to output/usage/costs.json')
        sys.exit(1)

    csv_path = Path(sys.argv[1])
    output_path = Path(sys.argv[2]) if len(sys.argv) > 2 else Path('output/usage/costs.json')

    if not csv_path.exists():
        print(f'Error: CSV file not found: {csv_path}')
        sys.exit(1)

    csv_to_json(csv_path, output_path)