#!/usr/bin/env python3
"""
Extract unique tenant_uuid values from inputs.resource.attr.tenant.tenant_uuid
in the Datadog logs CSV export.
"""

import argparse
import csv
import json
from pathlib import Path
from collections import Counter, defaultdict


def extract_tenant_uuids(csv_file_path):
    """Extract tenant UUIDs, their types, and counts grouped by principal ID."""

    # Structure: principal_id -> list of tenant_uuids
    principal_tenants = defaultdict(list)
    # Structure: principal_id -> tenant_uuid -> set of types
    principal_tenant_types = defaultdict(lambda: defaultdict(set))
    # Overall tenant tracking
    all_tenant_uuids = []
    all_tenant_types = defaultdict(set)

    with open(csv_file_path, 'r', encoding='utf-8') as csvfile:
        reader = csv.DictReader(csvfile)

        for row_num, row in enumerate(reader, start=2):  # start=2 because of header
            content = row.get('Content', '')

            if not content:
                continue

            try:
                # Parse the JSON content
                data = json.loads(content)

                # Navigate to the nested structure
                check_resources = data.get('checkResources', {})
                inputs = check_resources.get('inputs', [])

                # Track unique combinations per log entry to avoid duplicates
                log_entries = set()

                for input_item in inputs:
                    # Extract principal ID
                    principal = input_item.get('principal', {})
                    principal_id = principal.get('id')

                    # Extract tenant info
                    resource = input_item.get('resource', {})
                    attr = resource.get('attr', {})
                    tenant = attr.get('tenant', {})
                    tenant_uuid = tenant.get('tenant_uuid')
                    tenant_type = tenant.get('tenant_type')

                    if tenant_uuid and principal_id:
                        # Create unique key for this log entry
                        log_key = (principal_id, tenant_uuid, tenant_type or 'N/A')

                        # Only process if we haven't seen this combination in this log
                        if log_key not in log_entries:
                            log_entries.add(log_key)

                            # Track by principal
                            principal_tenants[principal_id].append(tenant_uuid)
                            if tenant_type:
                                principal_tenant_types[principal_id][tenant_uuid].add(tenant_type)

                            # Track overall
                            all_tenant_uuids.append(tenant_uuid)
                            if tenant_type:
                                all_tenant_types[tenant_uuid].add(tenant_type)

            except json.JSONDecodeError as e:
                print(f"Warning: Failed to parse JSON on row {row_num}: {e}")
                continue
            except Exception as e:
                print(f"Warning: Error processing row {row_num}: {e}")
                continue

    # Convert lists to Counters for each principal
    principal_tenant_counts = {
        principal_id: Counter(tenants)
        for principal_id, tenants in principal_tenants.items()
    }

    # Overall counts
    overall_counts = Counter(all_tenant_uuids)

    return principal_tenant_counts, principal_tenant_types, overall_counts, all_tenant_types

def main():
    # Setup argument parser
    parser = argparse.ArgumentParser(
        description='Extract tenant UUIDs from Datadog logs CSV export'
    )
    parser.add_argument(
        'csv_file',
        nargs='?',
        type=Path,
        help='Path to the CSV file to process (defaults to most recent extract-*.csv in current directory)'
    )

    args = parser.parse_args()

    # Determine which CSV file to use
    if args.csv_file:
        csv_file = args.csv_file
        if not csv_file.exists():
            print(f"Error: File not found: {csv_file}")
            return
    else:
        # Find the CSV file in the current directory (use the most recent one)
        csv_files = sorted(Path('.').glob('extract-*.csv'), reverse=True)

        if not csv_files:
            print("Error: No CSV file found matching pattern 'extract-*.csv'")
            print("Usage: python extract_tenant_uuids.py <path_to_csv_file>")
            return

        csv_file = csv_files[0]

    print(f"Processing file: {csv_file}")
    print("Writing results to output.md...")

    # Extract tenant UUIDs grouped by principal
    principal_tenant_counts, principal_tenant_types, overall_counts, all_tenant_types = extract_tenant_uuids(csv_file)

    # Write results to output.md
    with open('output.md', 'w', encoding='utf-8') as f:
        # Header
        f.write(f"# Datadog Logs Analysis - Tenant Access by Principal\n\n")
        f.write(f"**Source File:** `{csv_file.name}`\n\n")
        f.write(f"**Total Principals:** {len(principal_tenant_counts)}\n\n")
        f.write("---\n\n")

        # Results by principal
        for idx, (principal_id, tenant_counts) in enumerate(sorted(principal_tenant_counts.items()), 1):
            f.write(f"## Principal {idx}: `{principal_id}`\n\n")

            total_for_principal = sum(tenant_counts.values())

            f.write("| Tenant UUID | Type | Count |\n")
            f.write("|-------------|------|-------|\n")

            for tenant_uuid, count in tenant_counts.most_common():
                types = ', '.join(sorted(principal_tenant_types[principal_id][tenant_uuid]))
                if not types:
                    types = 'N/A'
                f.write(f"| `{tenant_uuid}` | {types} | {count} |\n")

            f.write(f"\n**Total for this principal:** {total_for_principal}\n\n")
            f.write("---\n\n")

        # Overall summary
        f.write("# Overall Summary - All Principals Combined\n\n")
        total_occurrences = sum(overall_counts.values())
        f.write(f"**Unique Tenant UUIDs:** {len(overall_counts)}\n\n")
        f.write(f"**Total Occurrences:** {total_occurrences}\n\n")

        f.write("| Tenant UUID | Type | Count |\n")
        f.write("|-------------|------|-------|\n")

        for uuid, count in overall_counts.most_common():
            types = ', '.join(sorted(all_tenant_types.get(uuid, ['N/A'])))
            f.write(f"| `{uuid}` | {types} | {count} |\n")

    print(f"✓ Results written to output.md")
    print(f"  - {len(principal_tenant_counts)} principals analyzed")
    print(f"  - {len(overall_counts)} unique tenant UUIDs found")
    print(f"  - {total_occurrences} total tenant access events")

if __name__ == '__main__':
    main()