#!/usr/bin/env python3
"""
Claude Code Analytics Analyzer
Fetches and analyzes Claude Code usage metrics per user.
"""

import argparse
import json
import os
import sys
import time
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.analytics.analytics_types import ClaudeCodeUsageRecord, parse_usage_record
from src.shared.api_client import fetch_api_data

# Global debug flag
DEBUG = False


def fetch_analytics_for_date(admin_key: str, date_str: str) -> list[ClaudeCodeUsageRecord]:
    """
    Fetch analytics data for a single date.

    Args:
        admin_key: Admin API key
        date_str: Date in YYYY-MM-DD format

    Returns:
        List of parsed usage records for that date
    """
    endpoint = '/v1/organizations/usage_report/claude_code'
    params = {'starting_at': date_str}

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

    response = fetch_api_data(endpoint, params, admin_key, debug=DEBUG)
    with open(f'output/analytics/raw/{date_str}.json', 'w') as f:
        json.dump(response, f, indent=2)
    # Parse response into dataclasses
    records = []
    for record_dict in response.get('data', []):
        record_dict['date'] = date_str
        record = parse_usage_record(record_dict)
        records.append(record)

    return records


def fetch_analytics_data(admin_key: str, start_date: str, end_date: str) -> list[ClaudeCodeUsageRecord]:
    """
    Fetch analytics data for a date range (inclusive).

    Args:
        admin_key: Admin API key
        start_date: Start date in YYYY-MM-DD format
        end_date: End date in YYYY-MM-DD format

    Returns:
        List of parsed usage records
    """
    start_dt = datetime.fromisoformat(start_date)
    end_dt = datetime.fromisoformat(end_date)

    all_records: list[ClaudeCodeUsageRecord] = []
    current_dt = start_dt

    while current_dt <= end_dt:
        date_str = current_dt.strftime('%Y-%m-%d')
        records = fetch_analytics_for_date(admin_key, date_str)
        time.sleep(1)
        all_records.extend(records)
        current_dt += timedelta(days=1)

    return all_records


def load_api_key_to_email_map(
    api_keys_path: Path = Path('output/usage/raw/api_keys.json'),
    users_path: Path = Path('output/usage/raw/users.json'),
) -> dict[str, str]:
    """
    Build api_key_name -> user email map from usage metadata.

    Used to attribute analytics records keyed by api_key_name (workspace API keys)
    back to the human user who owns the key. Returns empty dict if files missing.
    """
    if not api_keys_path.exists() or not users_path.exists():
        print(
            f'Warning: api_keys/users metadata not found at {api_keys_path}/{users_path}; '
            'api_actor records will not be remapped to email',
            file=sys.stderr,
        )
        return {}

    with open(users_path) as f:
        users_by_id = {u['id']: u for u in json.load(f).get('data', [])}

    mapping: dict[str, str] = {}
    with open(api_keys_path) as f:
        for key in json.load(f).get('data', []):
            name = key.get('name')
            created_by = key.get('created_by') or {}
            user_id = created_by.get('id') if isinstance(created_by, dict) else created_by
            email = users_by_id.get(user_id, {}).get('email')
            if name and email:
                mapping[name] = email

    return mapping


def _create_metrics_dict() -> dict:
    """Factory function to create a new metrics dict with models defaultdict."""
    return {
        'num_sessions': 0,
        'lines_added': 0,
        'lines_removed': 0,
        'commits': 0,
        'pull_requests': 0,
        'edit_tool_accepted': 0,
        'edit_tool_rejected': 0,
        'write_tool_accepted': 0,
        'write_tool_rejected': 0,
        'notebook_edit_accepted': 0,
        'notebook_edit_rejected': 0,
        'multi_edit_accepted': 0,
        'multi_edit_rejected': 0,
        'terminal_types': set(),  # Track unique terminal types used
        'models': defaultdict(
            lambda: {
                'input_tokens': 0,
                'output_tokens': 0,
                'cache_creation_tokens': 0,
                'cache_read_tokens': 0,
                'cost_cents': 0,
            }
        ),
    }


def aggregate_analytics(
    records: list[ClaudeCodeUsageRecord],
    time_grouping: str | None = None,
    api_key_to_email: dict[str, str] | None = None,
) -> dict:
    """
    Aggregate analytics by actor (user).

    Args:
        records: List of parsed usage records
        time_grouping: 'day', 'month', or None for total
        api_key_to_email: Optional api_key_name -> email map. When set, api_actor
            records resolve to the owning user's email so workspace-key activity
            merges with that user's user_actor records.

    Returns:
        Dict structure: actor -> [time_period] -> metrics
    """
    api_key_to_email = api_key_to_email or {}
    if time_grouping:
        # Three-level: actor -> time_period -> metrics
        actor_data = defaultdict(lambda: defaultdict(_create_metrics_dict))
    else:
        # Two-level: actor -> metrics
        actor_data = defaultdict(_create_metrics_dict)

    for idx, record in enumerate(records):
        try:
            if DEBUG:
                print(f'DEBUG: Processing record {idx + 1}/{len(records)}', file=sys.stderr)

            # Extract actor identifier. api_actor records (workspace keys) resolve
            # to the owning user's email when we have a mapping; otherwise fall
            # back to the key name so data still surfaces.
            api_key_name = record.actor.api_key_name
            actor = (
                record.actor.email_address
                or (api_key_to_email.get(api_key_name) if api_key_name else None)
                or api_key_name
                or 'Unknown'
            )
            date = record.date

            if DEBUG:
                print(f'DEBUG:   Actor: {actor}, Date: {date}', file=sys.stderr)

            # Determine time period
            if time_grouping:
                if time_grouping == 'day':
                    period = date
                elif time_grouping == 'month':
                    period = date[:7]  # YYYY-MM
                else:
                    period = 'total'

                if DEBUG:
                    print(f'DEBUG:   Time grouping: {time_grouping}, Period: {period}', file=sys.stderr)

                metrics = actor_data[actor][period]
            else:
                if DEBUG:
                    print('DEBUG:   No time grouping, accessing actor', file=sys.stderr)

                metrics = actor_data[actor]

            if DEBUG:
                print('DEBUG:   Got metrics dict', file=sys.stderr)

            # Aggregate core metrics
            metrics['num_sessions'] += record.core_metrics.num_sessions
            metrics['lines_added'] += record.core_metrics.lines_of_code.added
            metrics['lines_removed'] += record.core_metrics.lines_of_code.removed
            metrics['commits'] += record.core_metrics.commits_by_claude_code
            metrics['pull_requests'] += record.core_metrics.pull_requests_by_claude_code

            # Aggregate tool metrics
            metrics['edit_tool_accepted'] += record.tool_actions.edit_tool.accepted
            metrics['edit_tool_rejected'] += record.tool_actions.edit_tool.rejected
            metrics['write_tool_accepted'] += record.tool_actions.write_tool.accepted
            metrics['write_tool_rejected'] += record.tool_actions.write_tool.rejected
            metrics['notebook_edit_accepted'] += record.tool_actions.notebook_edit_tool.accepted
            metrics['notebook_edit_rejected'] += record.tool_actions.notebook_edit_tool.rejected
            metrics['multi_edit_accepted'] += record.tool_actions.multi_edit_tool.accepted
            metrics['multi_edit_rejected'] += record.tool_actions.multi_edit_tool.rejected

            # Track terminal type if present
            if record.terminal_type:
                metrics['terminal_types'].add(record.terminal_type)

            if DEBUG:
                print('DEBUG:   Aggregated core metrics', file=sys.stderr)

            # Aggregate model usage
            if DEBUG:
                print(
                    f'DEBUG:   Model breakdown has {len(record.model_breakdown)} models',
                    file=sys.stderr,
                )

            for model_idx, model_data in enumerate(record.model_breakdown):
                model_name = model_data.model

                if DEBUG:
                    print(f'DEBUG:     Model {model_idx + 1}: {model_name}', file=sys.stderr)

                model_metrics = metrics['models'][model_name]

                model_metrics['input_tokens'] += model_data.tokens.input
                model_metrics['output_tokens'] += model_data.tokens.output
                model_metrics['cache_creation_tokens'] += model_data.tokens.cache_creation
                model_metrics['cache_read_tokens'] += model_data.tokens.cache_read
                model_metrics['cost_cents'] += model_data.estimated_cost.amount

                if DEBUG:
                    print('DEBUG:     Updated model metrics', file=sys.stderr)

        except Exception as e:
            print(f'ERROR: Failed processing record {idx + 1}', file=sys.stderr)
            print(f'ERROR: Actor: {actor if "actor" in locals() else "N/A"}', file=sys.stderr)
            print(f'ERROR: Date: {date if "date" in locals() else "N/A"}', file=sys.stderr)
            print(f'ERROR: Exception: {e}', file=sys.stderr)
            import traceback

            traceback.print_exc()
            raise

    return actor_data


def format_json_report(actor_data: dict, time_grouping: str | None) -> str:
    """Format report as JSON."""
    actors_list = []
    total_sessions = 0
    total_commits = 0
    total_prs = 0
    total_cost_cents = 0

    for actor in sorted(actor_data.keys()):
        actor_report = {
            'actor': actor,
            'total_sessions': 0,
            'total_commits': 0,
            'total_prs': 0,
            'total_cost_cents': 0,
            'terminal_types': set(),  # Will aggregate all terminal types for this actor
        }

        if time_grouping:
            # Time periods structure
            actor_report['time_periods'] = []

            for period in sorted(actor_data[actor].keys()):
                metrics = actor_data[actor][period]

                # Convert models to list
                models_list = []
                period_cost = 0
                for model_name, model_metrics in metrics['models'].items():
                    models_list.append(
                        {
                            'model': model_name,
                            'input_tokens': model_metrics['input_tokens'],
                            'output_tokens': model_metrics['output_tokens'],
                            'cache_creation_tokens': model_metrics['cache_creation_tokens'],
                            'cache_read_tokens': model_metrics['cache_read_tokens'],
                            'cost_cents': model_metrics['cost_cents'],
                        }
                    )
                    period_cost += model_metrics['cost_cents']

                period_data = {
                    'time_period': period,
                    'num_sessions': metrics['num_sessions'],
                    'lines_added': metrics['lines_added'],
                    'lines_removed': metrics['lines_removed'],
                    'commits': metrics['commits'],
                    'pull_requests': metrics['pull_requests'],
                    'terminal_types': sorted(list(metrics['terminal_types'])),  # Convert set to sorted list
                    'tool_actions': {
                        'edit': {'accepted': metrics['edit_tool_accepted'], 'rejected': metrics['edit_tool_rejected']},
                        'write': {'accepted': metrics['write_tool_accepted'], 'rejected': metrics['write_tool_rejected']},
                        'notebook_edit': {'accepted': metrics['notebook_edit_accepted'], 'rejected': metrics['notebook_edit_rejected']},
                        'multi_edit': {'accepted': metrics['multi_edit_accepted'], 'rejected': metrics['multi_edit_rejected']},
                    },
                    'models': models_list,
                    'period_cost_cents': period_cost,
                }

                actor_report['time_periods'].append(period_data)
                actor_report['total_sessions'] += metrics['num_sessions']
                actor_report['total_commits'] += metrics['commits']
                actor_report['total_prs'] += metrics['pull_requests']
                actor_report['total_cost_cents'] += period_cost
                actor_report['terminal_types'].update(metrics['terminal_types'])  # Aggregate terminal types
        else:
            # No time periods
            metrics = actor_data[actor]

            # Convert models to list
            models_list = []
            for model_name, model_metrics in metrics['models'].items():
                models_list.append(
                    {
                        'model': model_name,
                        'input_tokens': model_metrics['input_tokens'],
                        'output_tokens': model_metrics['output_tokens'],
                        'cache_creation_tokens': model_metrics['cache_creation_tokens'],
                        'cache_read_tokens': model_metrics['cache_read_tokens'],
                        'cost_cents': model_metrics['cost_cents'],
                    }
                )
                actor_report['total_cost_cents'] += model_metrics['cost_cents']

            actor_report['num_sessions'] = metrics['num_sessions']
            actor_report['lines_added'] = metrics['lines_added']
            actor_report['lines_removed'] = metrics['lines_removed']
            actor_report['commits'] = metrics['commits']
            actor_report['pull_requests'] = metrics['pull_requests']
            actor_report['terminal_types'] = sorted(list(metrics['terminal_types']))  # Convert set to sorted list
            actor_report['tool_actions'] = {
                'edit': {'accepted': metrics['edit_tool_accepted'], 'rejected': metrics['edit_tool_rejected']},
                'write': {'accepted': metrics['write_tool_accepted'], 'rejected': metrics['write_tool_rejected']},
                'notebook_edit': {'accepted': metrics['notebook_edit_accepted'], 'rejected': metrics['notebook_edit_rejected']},
                'multi_edit': {'accepted': metrics['multi_edit_accepted'], 'rejected': metrics['multi_edit_rejected']},
            }
            actor_report['models'] = models_list
            actor_report['total_sessions'] = metrics['num_sessions']
            actor_report['total_commits'] = metrics['commits']
            actor_report['total_prs'] = metrics['pull_requests']
            actor_report['terminal_types'].update(metrics['terminal_types'])  # Aggregate terminal types

        # Convert terminal_types set to sorted list before adding to actors_list
        actor_report['terminal_types'] = sorted(list(actor_report['terminal_types']))

        total_sessions += actor_report['total_sessions']
        total_commits += actor_report['total_commits']
        total_prs += actor_report['total_prs']
        total_cost_cents += actor_report['total_cost_cents']

        actors_list.append(actor_report)

    result = {
        'actors': actors_list,
        'totals': {
            'total_sessions': total_sessions,
            'total_commits': total_commits,
            'total_prs': total_prs,
            'total_cost_cents': total_cost_cents,
        },
    }

    if time_grouping:
        result['time_grouping'] = time_grouping

    return json.dumps(result, indent=2)


def main() -> None:
    parser = argparse.ArgumentParser(
        description='Analyze Claude Code usage metrics',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Last 7 days
  %(prog)s --days 7

  # Specific date range
  %(prog)s --start 2025-10-01 --end 2025-10-07

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

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

    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 date (YYYY-MM-DD)')
    parser.add_argument('--end', help='End date (YYYY-MM-DD)')
    parser.add_argument(
        '--time-grouping',
        choices=['day', 'month'],
        default=None,
        help='Group by time period (day or month). Default: no grouping',
    )
    parser.add_argument(
        '--format',
        choices=['json'],
        default='json',
        help='Output format (default: json)',
    )
    parser.add_argument('--output', help='Output file path (default: stdout)')
    parser.add_argument(
        '--debug',
        action='store_true',
        help='Enable debug output',
    )

    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 date range
    if args.start and args.end:
        start_date = args.start
        end_date = args.end
    else:
        today = datetime.now(timezone.utc).date()
        end_date = today.isoformat()
        start_date = (today - timedelta(days=args.days - 1)).isoformat()

    try:
        # Fetch analytics data
        print(f'Fetching analytics from {start_date} to {end_date}...', file=sys.stderr)
        records = fetch_analytics_data(admin_key, start_date, end_date)
        print(f'Fetched {len(records)} records', file=sys.stderr)

        # Aggregate data
        print('Aggregating analytics...', file=sys.stderr)
        api_key_to_email = load_api_key_to_email_map()
        actor_data = aggregate_analytics(records, args.time_grouping, api_key_to_email)
        print(f'Found data for {len(actor_data)} actors', file=sys.stderr)

        # Format and output report
        report = format_json_report(actor_data, 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)
        if DEBUG:
            import traceback

            traceback.print_exc()
        sys.exit(1)


if __name__ == '__main__':
    main()