#!/usr/bin/env python3
import argparse
import csv
import sys
from pathlib import Path

PROJECT_ROOT = Path('/Users/rating/workspace/neta-bi-bot')
sys.path.insert(0, str(PROJECT_ROOT))

from dotenv import load_dotenv
load_dotenv(PROJECT_ROOT / '.env')

from scripts.lib.db import get_mysql_connection


def chunked(seq, size):
    for i in range(0, len(seq), size):
        yield seq[i:i + size]


def load_uuid_mapping(uuids):
    mapping = {}
    if not uuids:
        return mapping

    conn = get_mysql_connection(database='talesofai')
    try:
        with conn.cursor() as cursor:
            for batch in chunked(uuids, 500):
                placeholders = ','.join(['%s'] * len(batch))
                sql = f"SELECT id, uuid, nick_name FROM user WHERE uuid IN ({placeholders})"
                cursor.execute(sql, batch)
                for row in cursor.fetchall():
                    mapping[str(row['uuid'])] = {
                        'user_id': row['id'],
                        'nick_name': row.get('nick_name') or ''
                    }
    finally:
        conn.close()

    return mapping


def main():
    parser = argparse.ArgumentParser(description='Enrich trigger userlist csv with MySQL user_id/nick_name by user_uuid')
    parser.add_argument('--input', required=True, help='Input CSV path from ByteHouse query')
    parser.add_argument('--output', required=True, help='Output CSV path')
    args = parser.parse_args()

    input_path = Path(args.input)
    output_path = Path(args.output)

    if not input_path.exists():
        raise FileNotFoundError(f'Input file not found: {input_path}')

    with input_path.open('r', encoding='utf-8') as f:
        reader = csv.DictReader(f)
        rows = list(reader)

    uuid_order = []
    uuid_seen = set()
    for row in rows:
        u = str(row.get('user_uuid', '')).strip()
        if u and u not in uuid_seen:
            uuid_seen.add(u)
            uuid_order.append(u)

    mapping = load_uuid_mapping(uuid_order)

    # Insert user_id/nick_name right after user_uuid for readability.
    in_fields = list(rows[0].keys()) if rows else []
    out_fields = []
    inserted = False
    for f in in_fields:
        out_fields.append(f)
        if f == 'user_uuid':
            out_fields.extend(['user_id', 'nick_name'])
            inserted = True
    if not inserted:
        out_fields = ['user_id', 'nick_name'] + in_fields

    output_path.parent.mkdir(parents=True, exist_ok=True)
    with output_path.open('w', encoding='utf-8', newline='') as f:
        writer = csv.DictWriter(f, fieldnames=out_fields)
        writer.writeheader()
        for row in rows:
            u = str(row.get('user_uuid', '')).strip()
            m = mapping.get(u, {'user_id': '', 'nick_name': ''})
            row_out = dict(row)
            row_out['user_id'] = m['user_id']
            row_out['nick_name'] = m['nick_name']
            writer.writerow(row_out)

    print(f'Enriched rows: {len(rows)}')
    print(f'Mapped uuids: {len(mapping)} / {len(uuid_order)}')
    print(f'Output: {output_path}')


if __name__ == '__main__':
    main()