# backend/apps/hadis/management/commands/seed_hadis_chains.py

import random
import re
from django.core.management.base import BaseCommand
from django.core.cache import cache
from django.db import transaction
from apps.hadis.models import (
    Hadis,
    Transmitters,
    HadisTransmitter,
    NarratorLayer,
    TransmitterReliability,
)

DEFAULT_LAYERS = [
    {
        "number": 1,
        "slug_suffix": "l1-intermediate-regional-transmitters",
        "name": [
            {"language_code": "en", "text": "Intermediate Regional Transmitters"},
            {"language_code": "ru", "text": "Промежуточные региональные передатчики"},
        ],
        "description": [
            {
                "language_code": "en",
                "text": "Scholars and compilers transmitting across regional centers of learning.",
            },
            {
                "language_code": "ru",
                "text": "Ученые и составители, передававшие хадисы между региональными центрами знаний.",
            },
        ],
    },
    {
        "number": 2,
        "slug_suffix": "l2-primary-chain-early-narrators",
        "name": [
            {"language_code": "en", "text": "Primary Chain & Early Narrators"},
            {"language_code": "ru", "text": "Первичная цепочка и ранние передатчики"},
        ],
        "description": [
            {
                "language_code": "en",
                "text": "Early generation transmitters, companions, and direct students.",
            },
            {
                "language_code": "ru",
                "text": "Передатчики раннего поколения, сподвижники и прямые ученики.",
            },
        ],
    },
]


class Command(BaseCommand):
    help = (
        "Reassigns non-hadith transmission chains to chainless hadith arguments, "
        "creates chains for any remaining chainless hadiths, removes leftover non-hadith chains, "
        "and normalizes each hadith chain to have 1 uncertain node and 2-3 gaps."
    )

    def handle(self, *args, **options):
        rng = random.Random(42)

        narrators = list(Transmitters.objects.all().order_by("id"))
        reliabilities = list(TransmitterReliability.objects.all().order_by("id"))

        if len(narrators) < 6:
            self.stdout.write(
                self.style.ERROR("Not enough Transmitters found. Need at least 6 narrators.")
            )
            return

        # 1. Identify hadith vs non-hadith arguments
        hadith_qs = Hadis.objects.filter(category__source_type="hadith").order_by("id")
        non_hadith_with_chains = list(
            Hadis.objects.exclude(category__source_type="hadith")
            .filter(transmitters__isnull=False)
            .distinct()
            .order_by("id")
        )
        chainless_hadiths = list(
            hadith_qs.filter(transmitters__isnull=True).distinct().order_by("id")
        )

        self.stdout.write(
            f"Total hadith arguments: {hadith_qs.count()} | "
            f"Chainless hadiths: {len(chainless_hadiths)} | "
            f"Non-hadith arguments with chains: {len(non_hadith_with_chains)}"
        )

        with transaction.atomic():
            # Step 1: Move chains from non-hadith arguments to chainless hadith arguments (1-to-1) in bulk
            pairs_count = min(len(non_hadith_with_chains), len(chainless_hadiths))
            src_to_dst = {
                non_hadith_with_chains[i].id: chainless_hadiths[i].id
                for i in range(pairs_count)
            }

            if src_to_dst:
                dst_ids = list(src_to_dst.values())
                src_ids = list(src_to_dst.keys())

                # Remove any stray empty layers on dst_hadis if present
                NarratorLayer.objects.filter(hadis_id__in=dst_ids).delete()

                # Bulk update layers
                layers_to_move = list(
                    NarratorLayer.objects.filter(hadis_id__in=src_ids).order_by("id")
                )
                for layer in layers_to_move:
                    old_hadis_id = layer.hadis_id
                    new_hadis_id = src_to_dst[old_hadis_id]
                    old_slug = layer.slug or f"hadis-{old_hadis_id}-l{layer.number}"
                    new_slug = re.sub(
                        rf"^hadis-{old_hadis_id}-",
                        f"hadis-{new_hadis_id}-",
                        old_slug,
                    )
                    if new_slug == old_slug:
                        new_slug = f"hadis-{new_hadis_id}-l{layer.number}-{layer.id}"
                    layer.hadis_id = new_hadis_id
                    layer.slug = new_slug

                NarratorLayer.objects.bulk_update(
                    layers_to_move, ["hadis_id", "slug"], batch_size=500
                )

                # Bulk update transmitters
                txs_to_move = list(
                    HadisTransmitter.objects.filter(hadis_id__in=src_ids).order_by("id")
                )
                for tx in txs_to_move:
                    tx.hadis_id = src_to_dst[tx.hadis_id]

                HadisTransmitter.objects.bulk_update(
                    txs_to_move, ["hadis_id"], batch_size=500
                )

                self.stdout.write(
                    self.style.SUCCESS(
                        f"Moved {len(src_to_dst)} chains ({len(layers_to_move)} layers, {len(txs_to_move)} transmitters) "
                        "from non-hadith arguments to chainless hadith arguments."
                    )
                )

            # Step 2: Delete any remaining chains/layers on non-hadith arguments
            deleted_tx, _ = (
                HadisTransmitter.objects.exclude(hadis__category__source_type="hadith").delete()
            )
            deleted_layers, _ = (
                NarratorLayer.objects.filter(hadis__isnull=False)
                .exclude(hadis__category__source_type="hadith")
                .delete()
            )
            if deleted_tx or deleted_layers:
                self.stdout.write(
                    self.style.WARNING(
                        f"Deleted leftover non-hadith records: {deleted_tx} transmitters, {deleted_layers} layers."
                    )
                )

            # Step 3: Seed chains for any still-chainless hadith arguments in bulk
            remaining_chainless = list(
                hadith_qs.filter(transmitters__isnull=True).distinct().order_by("id")
            )
            if remaining_chainless:
                rem_ids = [h.id for h in remaining_chainless]
                NarratorLayer.objects.filter(hadis_id__in=rem_ids).delete()

                new_layers = []
                for hadis in remaining_chainless:
                    for ldef in DEFAULT_LAYERS:
                        new_layers.append(
                            NarratorLayer(
                                hadis_id=hadis.id,
                                number=ldef["number"],
                                name=ldef["name"],
                                description=ldef["description"],
                                slug=f"hadis-{hadis.id}-{ldef['slug_suffix']}",
                            )
                        )
                created_layers = NarratorLayer.objects.bulk_create(new_layers, batch_size=500)
                layers_by_hadis = {}
                for l in created_layers:
                    layers_by_hadis.setdefault(l.hadis_id, []).append(l)

                new_txs = []
                for hadis in remaining_chainless:
                    h_layers = sorted(layers_by_hadis[hadis.id], key=lambda x: x.number)
                    num_nodes = rng.randint(8, 10)
                    chosen_narrators = rng.sample(narrators, num_nodes)
                    split_idx = num_nodes // 2

                    for order_idx, narrator in enumerate(chosen_narrators):
                        layer_obj = h_layers[0] if order_idx < split_idx else h_layers[1]
                        status_obj = rng.choice(reliabilities) if reliabilities else None
                        new_txs.append(
                            HadisTransmitter(
                                hadis_id=hadis.id,
                                transmitter=narrator,
                                chain_index=0,
                                order=order_idx,
                                narrator_layer=layer_obj,
                                status=status_obj,
                                is_gap=False,
                                is_uncertain=False,
                            )
                        )
                HadisTransmitter.objects.bulk_create(new_txs, batch_size=500)
                self.stdout.write(
                    self.style.SUCCESS(
                        f"Created new chains for {len(remaining_chainless)} remaining chainless hadiths."
                    )
                )

            # Step 4: Normalize all hadith chains so every hadith has:
            # - 1 uncertain node (is_uncertain=True with 2 candidate transmitters in uncertain_transmitters) per chain
            # - 2 to 3 gaps (is_gap=True) per hadith
            ThroughModel = HadisTransmitter.uncertain_transmitters.through
            all_hadiths = list(
                hadith_qs.prefetch_related("transmitters", "narrator_layers").order_by("id")
            )

            nodes_to_update = []
            hadis_ids = [h.id for h in all_hadiths]
            ThroughModel.objects.filter(hadistransmitter__hadis_id__in=hadis_ids).delete()
            new_m2m_relations = []

            normalized_count = 0
            for hadis in all_hadiths:
                nodes = sorted(
                    list(hadis.transmitters.all()),
                    key=lambda x: (x.chain_index, x.order, x.id),
                )
                if not nodes:
                    continue

                chains = {}
                for n in nodes:
                    chains.setdefault(n.chain_index, []).append(n)

                for n in nodes:
                    n.is_gap = False
                    n.is_uncertain = False

                uncertain_node_ids = set()
                for c_idx, c_nodes in chains.items():
                    candidate_pool = c_nodes[1:] if len(c_nodes) > 2 else c_nodes
                    unc_node = rng.choice(candidate_pool)
                    unc_node.is_uncertain = True
                    uncertain_node_ids.add(unc_node.id)

                    used_tx_ids = {n.transmitter_id for n in c_nodes if n.transmitter_id}
                    available_candidates = [t for t in narrators if t.id not in used_tx_ids]
                    if len(available_candidates) < 2:
                        available_candidates = [
                            t for t in narrators if t.id != unc_node.transmitter_id
                        ]
                    picked_candidates = rng.sample(
                        available_candidates, min(2, len(available_candidates))
                    )
                    for cand in picked_candidates:
                        new_m2m_relations.append(
                            ThroughModel(
                                hadistransmitter_id=unc_node.id,
                                transmitters_id=cand.id,
                            )
                        )

                target_gaps = rng.choice([2, 3])
                gap_pool = [
                    n for n in nodes if n.id not in uncertain_node_ids and n.order > 0
                ]
                if len(gap_pool) < target_gaps:
                    gap_pool = [n for n in nodes if n.id not in uncertain_node_ids]
                if len(gap_pool) < target_gaps:
                    gap_pool = list(nodes)

                chosen_gaps = rng.sample(gap_pool, min(target_gaps, len(gap_pool)))
                for g_node in chosen_gaps:
                    g_node.is_gap = True

                nodes_to_update.extend(nodes)
                normalized_count += 1

            HadisTransmitter.objects.bulk_update(
                nodes_to_update, ["is_gap", "is_uncertain"], batch_size=500
            )
            ThroughModel.objects.bulk_create(new_m2m_relations, batch_size=500)

            self.stdout.write(
                self.style.SUCCESS(
                    f"Normalized {normalized_count} hadith chains (1 uncertain node per chain, 2-3 gaps per hadith)."
                )
            )

        try:
            cache.clear()
            self.stdout.write(self.style.SUCCESS("Cleared Django cache."))
        except Exception as e:
            self.stdout.write(self.style.WARNING(f"Cache clear warning: {e}"))