# your_app/management/commands/link_referrals_from_excel.py
from pathlib import Path
from typing import Dict, List

from django.core.management.base import BaseCommand, CommandError
from django.db import transaction

from openpyxl import load_workbook

from users.models import CustomUser  # adjust import

def clean_mobile(m):
    import re
    if m is None:
        return None
    digits = re.sub(r"\D+", "", str(m))
    return digits or None

class Command(BaseCommand):
    help = "Set CustomUser.referred_by from Excel 'Data' sheet, using referred_by column = referrer referral_code."

    def add_arguments(self, parser):
        parser.add_argument("excel_path", type=str, help="Path to the Excel file")
        parser.add_argument(
            "--batch-size",
            type=int,
            default=1000,
            help="Bulk update batch size (default 1000)",
        )
        parser.add_argument(
            "--dry-run",
            action="store_true",
            help="Parse and match but do not write to DB",
        )

    def handle(self, *args, **opts):
        excel_path = Path(opts["excel_path"])
        batch_size = opts["batch_size"]
        dry_run = opts["dry_run"]

        if not excel_path.exists():
            raise CommandError(f"Excel not found: {excel_path}")

        wb = load_workbook(excel_path, data_only=True, read_only=True)
        if "Data" not in wb.sheetnames:
            raise CommandError("Sheet 'Data' not found.")
        ws = wb["Data"]

        header = [str(c.value).strip() if c.value is not None else "" for c in next(ws.iter_rows(min_row=1, max_row=1))[0:]]
        need_cols = {"mobile_number", "referred_by"}
        if not need_cols.issubset(set(header)):
            missing = need_cols - set(header)
            raise CommandError(f"Missing columns in sheet: {', '.join(missing)}")

        # 1) Build referral_code → user_id map (only those who actually have a code)
        code_to_userid: Dict[str, int] = {
            rc: uid for rc, uid in
            CustomUser.objects.exclude(referral_code__isnull=True)
                              .exclude(referral_code__exact="")
                              .values_list("referral_code", "id")
        }

        # 2) Build mobile → id map for quick target lookup
        mobile_to_userid: Dict[str, int] = {
            m: uid for m, uid in
            CustomUser.objects.values_list("mobile_number", "id")
        }

        to_update: List[CustomUser] = []
        matched, not_found_target, no_ref_code_match, already_set = 0, 0, 0, 0

        # 3) Iterate rows and prepare updates
        for row in ws.iter_rows(min_row=2, values_only=True):
            row_data = {col: row[header.index(col)] if col in header else None for col in header}
            mobile = clean_mobile(row_data.get("mobile_number"))
            referred_by_code = (str(row_data.get("referred_by")).strip()
                                if row_data.get("referred_by") is not None else "")

            if not mobile:
                continue
            target_id = mobile_to_userid.get(mobile)
            if not target_id:
                not_found_target += 1
                continue
            if not referred_by_code:
                # No referral code in the sheet for this user
                continue

            referrer_id = code_to_userid.get(referred_by_code)
            if not referrer_id:
                no_ref_code_match += 1
                continue
            if referrer_id == target_id:
                # Prevent self-referral
                continue

            # Fetch only the id + referred_by to minimize memory
            user = CustomUser.objects.only("id", "referred_by_id").get(id=target_id)
            if user.referred_by_id == referrer_id:
                already_set += 1
                continue

            user.referred_by_id = referrer_id
            to_update.append(user)

            if len(to_update) >= batch_size:
                if not dry_run:
                    with transaction.atomic():
                        CustomUser.objects.bulk_update(to_update, ["referred_by"])
                matched += len(to_update)
                to_update.clear()

        # Final flush
        if to_update:
            if not dry_run:
                with transaction.atomic():
                    CustomUser.objects.bulk_update(to_update, ["referred_by"])
            matched += len(to_update)

        self.stdout.write(self.style.SUCCESS(
            f"Linked referred_by for {matched} users. "
            f"Targets not found: {not_found_target}, "
            f"No referral_code match: {no_ref_code_match}, "
            f"Already correct: {already_set}. "
            f"{'(dry-run)' if dry_run else ''}"
        ))