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

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

import argparse
import json
import sys
from pathlib import Path
from typing import Any

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

from src.cost_report.cost_report_analyzer import aggregate_cost_buckets, load_workspaces_metadata
from src.cost_report.cost_report_types import parse_cost_bucket

RAW_OUTPUT_DIR = Path('output/cost_report/raw')


def load_raw_cost_report_data() -> list:
    """
    Load and merge all cost report data from raw files.

    Deduplicates cost buckets by (starting_at, ending_at, workspace_id, model, cost_type)
    to avoid counting the same day multiple times when raw files have overlapping date ranges.

    Returns:
        List of deduplicated cost buckets from all raw files
    """
    # Find all daily cost report files (YYYY-MM-DD.json format, like analytics app)
    cost_report_files = sorted(RAW_OUTPUT_DIR.glob('????-??-??.json'))

    if not cost_report_files:
        print(f'Error: No cost report files found in {RAW_OUTPUT_DIR}', file=sys.stderr)
        print('Expected files in YYYY-MM-DD.json format', file=sys.stderr)
        sys.exit(1)

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

    # Collect all cost buckets from all files with deduplication
    # Use dict to deduplicate: key = (starting_at, ending_at, workspace_id, model, cost_type)
    cost_buckets_dict: dict[tuple, Any] = {}

    for cost_report_file in cost_report_files:
        with open(cost_report_file, 'r') as f:
            response = json.load(f)

        # Parse response into CostBucket objects
        # Response structure: data array with time buckets, each bucket has results array
        for time_bucket in response.get('data', []):
            for result in time_bucket.get('results', []):
                # Add time period info from parent bucket
                result['starting_at'] = time_bucket.get('starting_at', '')
                result['ending_at'] = time_bucket.get('ending_at', '')
                cost_bucket = parse_cost_bucket(result)

                # Create deduplication key (includes all fields populated by description grouping)
                dedup_key = (
                    cost_bucket.starting_at,
                    cost_bucket.ending_at,
                    cost_bucket.workspace_id or 'default',
                    cost_bucket.model or 'unknown',
                    cost_bucket.cost_type or 'unknown',
                    cost_bucket.context_window or '',
                    cost_bucket.service_tier or 'standard',
                    cost_bucket.token_type or '',
                )

                # Keep the cost bucket (later files overwrite earlier ones if same key)
                cost_buckets_dict[dedup_key] = cost_bucket

    all_cost_buckets = list(cost_buckets_dict.values())

    print(f'Total cost buckets loaded: {len(all_cost_buckets)} (after deduplication)', file=sys.stderr)

    return all_cost_buckets


def main() -> None:
    parser = argparse.ArgumentParser(
        description='Rebuild cost_report.json from raw API response files',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Rebuild cost_report.json from raw files
  %(prog)s --output output/cost_report/cost_report.json
        """,
    )

    parser.add_argument('--output', required=True, help='Output file path for cost_report.json')

    args = parser.parse_args()

    try:
        # Step 1: Load raw cost report data
        print('Loading raw cost report files...', file=sys.stderr)
        cost_buckets = load_raw_cost_report_data()

        # Step 2: Load workspace metadata for enrichment
        print('Loading workspace metadata from output/usage/raw/workspaces.json...', file=sys.stderr)
        workspaces_metadata = load_workspaces_metadata()
        if workspaces_metadata:
            print(f'Loaded metadata for {len(workspaces_metadata)} workspaces', file=sys.stderr)
        else:
            print('⚠️  Warning: No workspace metadata found, continuing without enrichment', file=sys.stderr)

        # Step 3: Aggregate cost buckets with workspace metadata
        print('Aggregating cost data...', file=sys.stderr)
        aggregated_data = aggregate_cost_buckets(cost_buckets, workspaces_metadata)

        # Step 4: 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(aggregated_data, f, indent=2)

        print(f'✅ Cost report data written to {output_path}', file=sys.stderr)
        print(f'Total cost: ${aggregated_data["total_cost"] / 100:.2f} (from {aggregated_data["total_cost"]:.2f} cents)', file=sys.stderr)
        print(f'Total workspaces: {aggregated_data["metadata"]["workspace_count"]}', file=sys.stderr)
        print(f'Date range: {aggregated_data["date_range"]["start"]} to {aggregated_data["date_range"]["end"]}', 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()