from django.db import models
from django.conf import settings
from django.urls import reverse
from django.utils.translation import gettext_lazy as _
from generic.mixins import AuthorMixin
from generic.models import CidUser
from atlas.models import CaseCollection


class ResearchStudy(models.Model, AuthorMixin):
    """Container for a research trial that assigns participants to collection packets."""

    class RandomisationMode(models.TextChoices):
        COMPLETELY_RANDOM = "RANDOM", _("Completely random")
        BALANCED = "BALANCED", _("Balanced (equal-fill)")
        BALANCED_WITHIN_GROUP = "BALANCED_GROUP", _("Balanced within candidate group")
        TARGET_BASED = "TARGET_BASED", _("Target-based allocation")

    class GenerationMode(models.TextChoices):
        AUTO_ON_ACCESS = "AUTO", _("Auto-create on first access")
        BULK_ONLY = "BULK", _("Bulk generation only")

    name = models.CharField(max_length=255)
    description = models.TextField(blank=True)

    active = models.BooleanField(default=True)
    
    randomisation_mode = models.CharField(
        max_length=20,
        choices=RandomisationMode.choices,
        default=RandomisationMode.BALANCED,
        help_text="How participants are assigned to packets.",
    )
    
    generation_mode = models.CharField(
        max_length=20,
        choices=GenerationMode.choices,
        default=GenerationMode.AUTO_ON_ACCESS,
        help_text="How participants are created for this study.",
    )

    balance_within_candidate_group = models.BooleanField(
        default=True,
        help_text="If enabled, balanced randomisation is performed independently per candidate group.",
    )
    
    collect_demographics = models.BooleanField(
        default=True,
        help_text="If enabled, participants are asked for simple demographic fields before starting.",
    )

    allocation_targets = models.JSONField(
        default=dict,
        blank=True,
        help_text="For target-based allocation: maps arm_id or group_name to target count. E.g. {'arm_1': 30, 'arm_2': 30}",
    )

    author = models.ManyToManyField(
        settings.AUTH_USER_MODEL,
        blank=True,
        help_text="Users allowed to manage this study.",
        related_name="research_studies",
    )

    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)

    class Meta:
        ordering = ("-created_at",)
        verbose_name = "Research Study"
        verbose_name_plural = "Research Studies"

    def __str__(self):
        return self.name

    def get_absolute_url(self):
        return reverse("research:study_detail", kwargs={"pk": self.pk})


class ResearchStudyArm(models.Model):
    """One packet/arm in a study linked to a CaseCollection."""

    study = models.ForeignKey(
        ResearchStudy,
        on_delete=models.CASCADE,
        related_name="arms",
    )
    name = models.CharField(max_length=100)
    packet = models.ForeignKey(
        CaseCollection,
        on_delete=models.PROTECT,
        related_name="research_arms",
        help_text="CaseCollection used as this packet.",
    )
    active = models.BooleanField(default=True)
    allocation_weight = models.PositiveIntegerField(
        default=1,
        help_text="Relative weighting used for completely-random assignment.",
    )
    sort_order = models.PositiveIntegerField(default=100)

    # Support for ordered packets with prerequisites
    order_sequence = models.PositiveIntegerField(
        default=0,
        help_text="Order in which this arm should be completed (0 = no ordering). Used with prerequisite support.",
    )
    required_previous_arm = models.ForeignKey(
        "self",
        on_delete=models.SET_NULL,
        null=True,
        blank=True,
        related_name="dependent_arms",
        help_text="If set, participant must complete this arm before the dependent arm.",
    )

    class Meta:
        ordering = ("sort_order", "id")
        unique_together = ("study", "name")
        verbose_name = "Research Study Arm"
        verbose_name_plural = "Research Study Arms"

    def __str__(self):
        return f"{self.study.name}: {self.name}"


class ResearchParticipant(models.Model):
    """Pseudo-anonymised participant record for one study with 1-1 CidUser mapping."""

    class AssignmentMethod(models.TextChoices):
        COMPLETELY_RANDOM = "RANDOM", _("Completely random")
        BALANCED = "BALANCED", _("Balanced")
        BALANCED_WITHIN_GROUP = "BALANCED_GROUP", _("Balanced within group")
        TARGET_BASED = "TARGET_BASED", _("Target-based")
        MANUAL = "MANUAL", _("Manual")

    study = models.ForeignKey(
        ResearchStudy,
        on_delete=models.CASCADE,
        related_name="participants",
    )
    pseudo_id = models.CharField(
        max_length=100,
        help_text="Pseudo-anonymised ID used in participant URL.",
    )

    # 1-1 mapping to CidUser per study. Nullable until participant completes intake.
    cid_user = models.OneToOneField(
        CidUser,
        on_delete=models.SET_NULL,
        null=True,
        blank=True,
        related_name="research_participant",
        help_text="CidUser for this participant (1-1 mapping; assigned at intake).",
    )

    # Optional link to authenticated user (if participant logged in)
    user_user = models.ForeignKey(
        settings.AUTH_USER_MODEL,
        on_delete=models.SET_NULL,
        null=True,
        blank=True,
        related_name="research_participants",
    )

    # Demographics captured at intake
    candidate_group = models.CharField(max_length=100, blank=True)
    age_band = models.CharField(max_length=100, blank=True)
    sex = models.CharField(max_length=50, blank=True)
    training_grade = models.CharField(max_length=100, blank=True)
    years_experience = models.CharField(max_length=50, blank=True)
    demographics_extra = models.JSONField(default=dict, blank=True)

    # Email for result reaccess (stored on CidUser, but captured here for convenience)
    email = models.EmailField(
        blank=True,
        help_text="Email address for participant result reaccess.",
    )

    consented = models.BooleanField(default=False)

    # Arm assignment
    assigned_arm = models.ForeignKey(
        ResearchStudyArm,
        on_delete=models.SET_NULL,
        null=True,
        blank=True,
        related_name="participants",
    )
    assignment_method = models.CharField(
        max_length=20,
        choices=AssignmentMethod.choices,
        blank=True,
    )
    assigned_at = models.DateTimeField(null=True, blank=True)

    # Completion tracking
    completed_arms = models.ManyToManyField(
        ResearchStudyArm,
        blank=True,
        related_name="completed_by_participants",
        help_text="Arms this participant has completed.",
    )

    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)

    class Meta:
        unique_together = ("study", "pseudo_id")
        # Ensure CidUser is unique per study (1-1 mapping)
        constraints = [
            models.UniqueConstraint(
                fields=['study', 'cid_user'],
                name='unique_cid_per_study'
            )
        ]
        ordering = ("-created_at",)
        verbose_name = "Research Participant"
        verbose_name_plural = "Research Participants"

    def __str__(self):
        return f"{self.study.pk}:{self.pseudo_id}"

    def is_prerequisite_satisfied(self):
        """Check if participant has completed prerequisite arm(s) for current assignment."""
        if not self.assigned_arm or not self.assigned_arm.required_previous_arm:
            return True
        return self.completed_arms.filter(pk=self.assigned_arm.required_previous_arm.pk).exists()