#!/usr/bin/env python3
"""
Rebuild cost.json from raw API response files.

This script reads raw API responses from output/usage/raw/ and regenerates
the complete cost.json file with proper aggregation and enrichment.
"""

import argparse
import json
import sys
from pathlib import Path

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

from src.usage.cost_analyzer import (
    RAW_OUTPUT_DIR,
    aggregate_usage,
    build_cost_data_structure,
    enrich_api_keys_with_user_info,
    enrich_api_keys_with_workspace_info,
)


def load_raw_users() -> dict:
    """Load users from raw response file."""
    users_file = RAW_OUTPUT_DIR / 'users.json'
    if not users_file.exists():
        print(f'Warning: {users_file} not found, using empty users dict', file=sys.stderr)
        return {}

    with open(users_file, 'r') as f:
        response = json.load(f)

    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 load_raw_api_keys() -> dict:
    """Load API keys from raw response file."""
    api_keys_file = RAW_OUTPUT_DIR / 'api_keys.json'
    if not api_keys_file.exists():
        print(f'Warning: {api_keys_file} not found, using empty API keys dict', file=sys.stderr)
        return {}

    with open(api_keys_file, 'r') as f:
        response = json.load(f)

    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 load_raw_workspaces() -> dict:
    """Load workspaces from raw response file."""
    workspaces_file = RAW_OUTPUT_DIR / 'workspaces.json'
    if not workspaces_file.exists():
        print(f'Warning: {workspaces_file} not found, using empty workspaces dict', file=sys.stderr)
        return {}

    with open(workspaces_file, 'r') as f:
        response = json.load(f)

    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 load_raw_usage_data() -> dict:
    """
    Load and merge all usage data from raw message files.

    Returns:
        Merged usage data dict with all buckets
    """
    # Find all messages files (messages_YYYY-MM-DD.json - one file per day)
    messages_files = sorted(RAW_OUTPUT_DIR.glob('messages_*.json'))

    if not messages_files:
        print(f'Error: No messages files found in {RAW_OUTPUT_DIR}', file=sys.stderr)
        sys.exit(1)

    print(f'Found {len(messages_files)} message file(s):', file=sys.stderr)
    for f in messages_files:
        print(f'  - {f.name}', file=sys.stderr)

    # Merge all usage data from individual day files
    merged_data = {'data': []}

    for messages_file in messages_files:
        with open(messages_file, 'r') as f:
            data = json.load(f)

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

    print(f'Total buckets loaded: {len(merged_data["data"])}', file=sys.stderr)

    return merged_data


def main() -> None:
    parser = argparse.ArgumentParser(
        description='Rebuild cost.json from raw API response files',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Rebuild with daily grouping
  %(prog)s --time-grouping day --output output/usage/costs.json

  # Rebuild with monthly grouping
  %(prog)s --time-grouping month --output output/usage/costs_monthly.json
        """,
    )

    parser.add_argument(
        '--time-grouping',
        choices=['day', 'month'],
        required=True,
        help='Time grouping (day or month)',
    )
    parser.add_argument('--output', required=True, help='Output file path for cost.json')

    args = parser.parse_args()

    try:
        # Step 1: Load raw data
        print('Loading raw data files...', file=sys.stderr)
        users_data = load_raw_users()
        api_keys_metadata = load_raw_api_keys()
        workspaces_data = load_raw_workspaces()
        usage_data = load_raw_usage_data()

        print(f'Loaded {len(users_data)} users', file=sys.stderr)
        print(f'Loaded {len(api_keys_metadata)} API keys', file=sys.stderr)
        print(f'Loaded {len(workspaces_data)} workspaces', file=sys.stderr)

        # Step 2: Enrich API keys
        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)

        # Step 3: Aggregate usage
        print(f'Aggregating usage with time grouping: {args.time_grouping}...', file=sys.stderr)
        api_key_usage = aggregate_usage(usage_data, args.time_grouping)
        print(f'Found usage data for {len(api_key_usage)} API keys', file=sys.stderr)

        # Step 4: Build cost data structure
        print('Building cost data structure...', file=sys.stderr)
        cost_data = build_cost_data_structure(api_key_usage, api_keys_metadata, args.time_grouping)

        # Step 5: Write output
        output_path = Path(args.output)
        output_path.parent.mkdir(parents=True, exist_ok=True)

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

        print(f'✅ Cost data written to {output_path}', file=sys.stderr)
        print(f'Total organization cost: ${cost_data["total_organization_cost"]:.2f}', file=sys.stderr)
        print(f'Total API keys: {len(cost_data["api_keys"])}', file=sys.stderr)

        # Show time periods stats
        if args.time_grouping:
            total_periods = sum(len(k.get('time_periods', [])) for k in cost_data['api_keys'])
            print(f'Total time periods: {total_periods}', file=sys.stderr)

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

        traceback.print_exc()
        sys.exit(1)


if __name__ == '__main__':
    main()