from django import forms
from django.forms import (
    ModelForm,
    ValidationError,
    CheckboxInput,
    TextInput,
)
from crispy_forms.helper import FormHelper
from crispy_forms.layout import (
    Layout,
    Fieldset,
    Div,
    Row,
    Column,
    Submit,
    HTML,
)

from research.models import ResearchStudy, ResearchStudyArm, ResearchParticipant


class ResearchStudyForm(ModelForm):
    """Form for creating/editing research studies."""
    
    class Meta:
        model = ResearchStudy
        fields = [
            "name",
            "description",
            "active",
            "randomisation_mode",
            "generation_mode",
            "balance_within_candidate_group",
            "collect_demographics",
            "allocation_targets",
        ]
        widgets = {
            "description": forms.Textarea(attrs={"rows": 4}),
            "allocation_targets": forms.Textarea(attrs={"rows": 4, "placeholder": '{"arm_1": 30, "arm_2": 30}'}),
        }

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.helper = FormHelper()
        self.helper.form_method = 'post'
        self.helper.layout = Layout(
            Fieldset(
                'Study Details',
                'name',
                'description',
                'active',
            ),
            Fieldset(
                'Participant Creation',
                'generation_mode',
                'collect_demographics',
            ),
            Fieldset(
                'Randomisation Strategy',
                'randomisation_mode',
                'balance_within_candidate_group',
                Row(
                    Column('allocation_targets', css_class='col-md-12'),
                    HTML('<small class="form-text text-muted">For target-based allocation, provide JSON like: {"arm_id": count, ...}</small>'),
                ),
            ),
            Submit('submit', 'Save Study'),
        )


class ResearchStudyArmForm(ModelForm):
    """Form for creating/editing study arms (packets)."""
    
    class Meta:
        model = ResearchStudyArm
        fields = [
            "name",
            "packet",
            "active",
            "allocation_weight",
            "sort_order",
            "order_sequence",
            "required_previous_arm",
        ]

    def __init__(self, *args, study=None, **kwargs):
        super().__init__(*args, **kwargs)
        self.helper = FormHelper()
        self.helper.form_method = 'post'
        
        # Filter previous_arm choices to arms in the same study
        if study:
            self.fields['required_previous_arm'].queryset = ResearchStudyArm.objects.filter(
                study=study
            ).exclude(pk=self.instance.pk if self.instance.pk else None)
        else:
            self.fields['required_previous_arm'].queryset = ResearchStudyArm.objects.none()
        
        self.helper.layout = Layout(
            Fieldset(
                'Arm Configuration',
                'name',
                'packet',
                'active',
                'sort_order',
            ),
            Fieldset(
                'Randomisation',
                'allocation_weight',
            ),
            Fieldset(
                'Ordering & Prerequisites',
                'order_sequence',
                'required_previous_arm',
                HTML('<small class="form-text text-muted">Set order_sequence to enforce ordering. Set required_previous_arm to enforce prerequisites.</small>'),
            ),
            Submit('submit', 'Save Arm'),
        )


class ResearchParticipantIntakeForm(ModelForm):
    """Form for participant intake (demographics and consent)."""
    
    class Meta:
        model = ResearchParticipant
        fields = [
            "candidate_group",
            "age_band",
            "sex",
            "training_grade",
            "years_experience",
            "email",
            "consented",
        ]
        widgets = {
            "candidate_group": TextInput(attrs={"placeholder": "e.g. Group A"}),
            "age_band": TextInput(attrs={"placeholder": "e.g. 30-39"}),
            "sex": TextInput(attrs={"placeholder": "Optional"}),
            "training_grade": TextInput(attrs={"placeholder": "e.g. ST4"}),
            "years_experience": TextInput(attrs={"placeholder": "e.g. 5"}),
            "email": TextInput(attrs={"placeholder": "optional@example.com", "type": "email"}),
            "consented": CheckboxInput(),
        }

    def __init__(self, *args, collect_demographics=True, **kwargs):
        super().__init__(*args, **kwargs)
        self.collect_demographics = collect_demographics
        
        if not collect_demographics:
            # Make demographic fields optional if not collecting
            for f in [
                "candidate_group",
                "age_band",
                "sex",
                "training_grade",
                "years_experience",
            ]:
                self.fields[f].required = False
        else:
            # Make candidate_group required if collecting (useful for balancing)
            self.fields["candidate_group"].required = True
        
        self.helper = FormHelper()
        self.helper.form_method = 'post'
        self.helper.layout = Layout(
            Fieldset(
                'Demographics',
                'candidate_group',
                'age_band',
                'sex',
                'training_grade',
                'years_experience',
            ),
            Fieldset(
                'Contact & Consent',
                'email',
                HTML('<div class="alert alert-info">Provide an email address to reaccess your results later if you forget your passcode.</div>'),
                'consented',
                HTML('<div class="alert alert-warning"><strong>Consent Required:</strong> You must consent to participate to continue.</div>'),
            ),
            Submit('submit', 'Start Study'),
        )

    def clean_consented(self):
        """Require explicit consent."""
        consented = self.cleaned_data.get("consented")
        if not consented:
            raise ValidationError("Consent is required to continue.")
        return consented

    def clean_email(self):
        """Validate email format if provided."""
        email = self.cleaned_data.get("email", "").strip()
        if email:
            # Basic email validation is done by EmailField, but we can add custom logic here
            pass
        return email


class BulkParticipantGenerationForm(forms.Form):
    """Form for bulk generating participants from CSV/JSON."""
    
    UPLOAD_FORMAT_CHOICES = [
        ('csv', 'CSV (pseudo_id, candidate_group, email)'),
        ('json', 'JSON (array of objects)'),
    ]
    
    format = forms.ChoiceField(
        choices=UPLOAD_FORMAT_CHOICES,
        initial='csv',
        help_text="Choose upload format",
    )
    
    file = forms.FileField(
        help_text="CSV or JSON file with participant data",
    )

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.helper = FormHelper()
        self.helper.form_method = 'post'
        self.helper.layout = Layout(
            Fieldset(
                'Bulk Generate Participants',
                'format',
                'file',
                HTML('<small class="form-text text-muted">CSV format: pseudo_id,candidate_group,email<br>JSON format: [{"pseudo_id": "P001", "candidate_group": "A", "email": "..."}]</small>'),
            ),
            Submit('submit', 'Generate Participants'),
        )