import csv
import json
import random
from django.shortcuts import render, redirect, get_object_or_404
from django.contrib.auth.decorators import login_required
from django.core.exceptions import PermissionDenied
from django.db import transaction
from django.http import HttpResponse, JsonResponse, Http404
from django.urls import reverse
from django.utils import timezone
from django.utils.crypto import get_random_string
from django.contrib import messages

from atlas.models import CaseCollection, CaseDetail, CidReportAnswer
from generic.models import CidUser

from research.models import ResearchStudy, ResearchStudyArm, ResearchParticipant
from research.forms import (
    ResearchStudyForm,
    ResearchStudyArmForm,
    ResearchParticipantIntakeForm,
    BulkParticipantGenerationForm,
)


# ============================================================================
# Utility Functions
# ============================================================================

def _user_can_manage_research_study(study: ResearchStudy, user) -> bool:
    """Check if user can manage a research study."""
    if not user.is_authenticated:
        return False
    if user.is_superuser:
        return True
    if user.groups.filter(name="atlas_editor").exists():
        return True
    return study.author.filter(pk=user.pk).exists()


def _get_next_cid():
    """Get the next available CID number."""
    from django.db.models import Max
    max_cid = CidUser.objects.aggregate(Max('cid'))['cid__max'] or 0
    return max(max_cid + 1, 100000)


def _choose_research_arm(study: ResearchStudy, candidate_group: str = "") -> ResearchStudyArm:
    """Select an arm for participant assignment based on study's randomisation strategy."""
    arms = list(study.arms.filter(active=True).select_related("packet").order_by("sort_order", "id"))
    if not arms:
        raise Http404("No active packets configured for this study")

    randomisation_mode = study.randomisation_mode

    # Completely random: weighted by allocation_weight
    if randomisation_mode == ResearchStudy.RandomisationMode.COMPLETELY_RANDOM:
        weights = [max(1, a.allocation_weight) for a in arms]
        return random.choices(arms, weights=weights, k=1)[0]

    # Target-based: allocate until arm reaches target, then move to next
    if randomisation_mode == ResearchStudy.RandomisationMode.TARGET_BASED:
        targets = study.allocation_targets or {}
        eligible = []
        
        for arm in arms:
            target = targets.get(str(arm.pk), 0)
            if target <= 0:  # No limit
                eligible.append(arm)
            else:
                current_count = ResearchParticipant.objects.filter(
                    study=study,
                    assigned_arm=arm,
                ).count()
                if current_count < target:
                    eligible.append(arm)
        
        if not eligible:
            # All targets met; fall back to balanced
            participant_qs = ResearchParticipant.objects.filter(study=study, assigned_arm__isnull=False)
            if study.balance_within_candidate_group and candidate_group:
                participant_qs = participant_qs.filter(candidate_group=candidate_group)
            counts = {arm.pk: participant_qs.filter(assigned_arm=arm).count() for arm in arms}
            min_count = min(counts.values())
            eligible = [arm for arm in arms if counts[arm.pk] == min_count]
        
        return random.choice(eligible)

    # Balanced within group: balance per candidate_group
    if randomisation_mode == ResearchStudy.RandomisationMode.BALANCED_WITHIN_GROUP and candidate_group:
        participant_qs = ResearchParticipant.objects.filter(
            study=study,
            assigned_arm__isnull=False,
            candidate_group=candidate_group,
        )
        counts = {arm.pk: participant_qs.filter(assigned_arm=arm).count() for arm in arms}
        min_count = min(counts.values())
        eligible = [arm for arm in arms if counts[arm.pk] == min_count]
        return random.choice(eligible)

    # Balanced (default): global balance across all participants
    participant_qs = ResearchParticipant.objects.filter(study=study, assigned_arm__isnull=False)
    counts = {arm.pk: participant_qs.filter(assigned_arm=arm).count() for arm in arms}
    min_count = min(counts.values())
    eligible = [arm for arm in arms if counts[arm.pk] == min_count]
    return random.choice(eligible)


def _ensure_research_participant_cid(participant: ResearchParticipant) -> CidUser:
    """Ensure participant has a CidUser (creates if needed); maintains 1-1 mapping."""
    if participant.cid_user is not None:
        return participant.cid_user

    # Try to allocate a unique CID
    created_cid = None
    for _ in range(10):
        try:
            try:
                next_cid = _get_next_cid()
            except Exception:
                next_cid = random.randint(100000, 999999)
            
            created_cid = CidUser.objects.create(
                cid=next_cid,
                passcode=get_random_string(8),
                name=f"study:{participant.study_id}:{participant.pseudo_id}",
                active=True,
            )
            break
        except Exception:
            continue

    if created_cid is None:
        raise RuntimeError("Unable to allocate CID user for participant")

    # Update participant with CidUser (1-1 relationship)
    try:
        participant.cid_user = created_cid
        participant.save(update_fields=["cid_user", "updated_at"])
    except Exception as e:
        # If 1-1 constraint violated, another participant may have been created; clean up
        created_cid.delete()
        raise RuntimeError(f"Unable to assign CID to participant: {e}")

    return created_cid


def _build_study_export_rows(study: ResearchStudy):
    """Build export rows for study results."""
    rows = []
    participants = (
        study.participants.select_related("assigned_arm", "assigned_arm__packet", "cid_user", "user_user")
        .order_by("pseudo_id")
    )

    for participant in participants:
        raw_score = None
        max_possible_score = None
        normalised_percent = None
        marked_answers = 0
        total_cases = 0
        attempt_start = None
        attempt_end = None
        attempt_completed = False

        if participant.assigned_arm and participant.cid_user:
            packet = participant.assigned_arm.packet
            casedetails = CaseDetail.objects.filter(collection=packet)
            total_cases = casedetails.count()
            answers = CidReportAnswer.objects.filter(
                question__in=casedetails,
                cid=participant.cid_user.cid,
            )
            marked_scores = [a.score for a in answers if a.score is not None]
            raw_score = sum(marked_scores) if marked_scores else 0
            marked_answers = len(marked_scores)

            # Get case max score from packet's marking config
            case_max = 10  # Default
            if packet.case_question_max_score:
                case_max = packet.case_question_max_score

            max_possible_score = total_cases * case_max
            if max_possible_score > 0:
                normalised_percent = round((raw_score / max_possible_score) * 100, 2)

            attempt = packet.get_cid_exams(cid_user=participant.cid_user).first()
            if attempt:
                attempt_start = attempt.start_time
                attempt_end = attempt.end_time
                attempt_completed = bool(attempt.completed)

        rows.append(
            {
                "study_id": study.pk,
                "study_name": study.name,
                "pseudo_id": participant.pseudo_id,
                "candidate_group": participant.candidate_group,
                "age_band": participant.age_band,
                "sex": participant.sex,
                "training_grade": participant.training_grade,
                "years_experience": participant.years_experience,
                "email": participant.email,
                "consented": participant.consented,
                "assigned_arm": participant.assigned_arm.name if participant.assigned_arm else "",
                "packet": participant.assigned_arm.packet.name if participant.assigned_arm else "",
                "assignment_method": participant.assignment_method,
                "assigned_at": participant.assigned_at,
                "cid": participant.cid_user.cid if participant.cid_user else "",
                "passcode": participant.cid_user.passcode if participant.cid_user else "",
                "attempt_start": attempt_start,
                "attempt_end": attempt_end,
                "attempt_completed": attempt_completed,
                "total_cases": total_cases,
                "marked_answers": marked_answers,
                "raw_score": raw_score,
                "max_possible_score": max_possible_score,
                "normalised_percent": normalised_percent,
            }
        )

    return rows


# ============================================================================
# Study Management Views
# ============================================================================

@login_required
def study_list(request):
    """List research studies accessible to the user."""
    studies = ResearchStudy.objects.all().prefetch_related("author")
    if not (request.user.is_superuser or request.user.groups.filter(name="atlas_editor").exists()):
        studies = studies.filter(author=request.user)
    return render(request, "research/study_list.html", {"studies": studies.distinct()})


@login_required
def study_create(request):
    """Create a new research study."""
    if request.method == "POST":
        form = ResearchStudyForm(request.POST)
        if form.is_valid():
            study = form.save()
            study.author.add(request.user)
            messages.success(request, f"Study '{study.name}' created successfully.")
            return redirect("research:study_detail", pk=study.pk)
    else:
        form = ResearchStudyForm()
    return render(request, "research/study_form.html", {"form": form, "is_create": True})


@login_required
def study_update(request, pk: int):
    """Update an existing research study."""
    study = get_object_or_404(ResearchStudy, pk=pk)
    if not _user_can_manage_research_study(study, request.user):
        raise PermissionDenied

    if request.method == "POST":
        form = ResearchStudyForm(request.POST, instance=study)
        if form.is_valid():
            form.save()
            messages.success(request, f"Study '{study.name}' updated successfully.")
            return redirect("research:study_detail", pk=study.pk)
    else:
        form = ResearchStudyForm(instance=study)

    return render(
        request,
        "research/study_form.html",
        {"form": form, "study": study, "is_create": False},
    )


@login_required
def study_detail(request, pk: int):
    """Display study detail with arms and participants."""
    study = get_object_or_404(ResearchStudy, pk=pk)
    if not _user_can_manage_research_study(study, request.user):
        raise PermissionDenied

    arm_form = ResearchStudyArmForm(study=study) if request.method != "POST" else None
    if request.method == "POST" and "add_arm" in request.POST:
        arm_form = ResearchStudyArmForm(request.POST, study=study)
        if arm_form.is_valid():
            arm = arm_form.save(commit=False)
            arm.study = study
            arm.save()
            messages.success(request, f"Arm '{arm.name}' added successfully.")
            return redirect("research:study_detail", pk=study.pk)

    arms = study.arms.select_related("packet").all()
    participants = study.participants.select_related("assigned_arm", "assigned_arm__packet", "cid_user").all()
    arm_rows = [
        (arm, participants.filter(assigned_arm=arm).count())
        for arm in arms
    ]
    return render(
        request,
        "research/study_detail.html",
        {
            "study": study,
            "arms": arms,
            "participants": participants,
            "arm_rows": arm_rows,
            "arm_form": arm_form,
        },
    )


@login_required
def study_arm_update(request, pk: int, arm_id: int):
    """Update a study arm."""
    study = get_object_or_404(ResearchStudy, pk=pk)
    if not _user_can_manage_research_study(study, request.user):
        raise PermissionDenied

    arm = get_object_or_404(ResearchStudyArm, pk=arm_id, study=study)
    if request.method == "POST":
        form = ResearchStudyArmForm(request.POST, instance=arm, study=study)
        if form.is_valid():
            form.save()
            messages.success(request, f"Arm '{arm.name}' updated successfully.")
            return redirect("research:study_detail", pk=study.pk)
    else:
        form = ResearchStudyArmForm(instance=arm, study=study)

    return render(
        request,
        "research/study_arm_form.html",
        {"study": study, "form": form, "arm": arm},
    )


@login_required
def study_export_csv(request, pk: int):
    """Export study results to CSV."""
    study = get_object_or_404(ResearchStudy, pk=pk)
    if not _user_can_manage_research_study(study, request.user):
        raise PermissionDenied

    rows = _build_study_export_rows(study)
    response = HttpResponse(content_type="text/csv")
    response["Content-Disposition"] = f'attachment; filename="study_{study.pk}_results.csv"'

    if rows:
        writer = csv.DictWriter(response, fieldnames=list(rows[0].keys()))
        writer.writeheader()
        for row in rows:
            writer.writerow(row)

    return response


@login_required
def study_export_json(request, pk: int):
    """Export study results to JSON."""
    study = get_object_or_404(ResearchStudy, pk=pk)
    if not _user_can_manage_research_study(study, request.user):
        raise PermissionDenied

    rows = _build_study_export_rows(study)
    return JsonResponse(rows, safe=False)


@login_required
def study_bulk_generate(request, pk: int):
    """Generate participants in bulk from CSV/JSON."""
    study = get_object_or_404(ResearchStudy, pk=pk)
    if not _user_can_manage_research_study(study, request.user):
        raise PermissionDenied

    generated_count = 0
    error_message = None

    if request.method == "POST":
        form = BulkParticipantGenerationForm(request.POST, request.FILES)
        if form.is_valid():
            try:
                upload_format = form.cleaned_data["format"]
                uploaded_file = request.FILES["file"]
                
                if upload_format == "csv":
                    # Parse CSV
                    decoded_file = uploaded_file.read().decode('utf-8').splitlines()
                    reader = csv.DictReader(decoded_file)
                    for row_num, row in enumerate(reader, start=2):  # +2 for header + 1-indexing
                        try:
                            pseudo_id = row.get("pseudo_id", "").strip()
                            candidate_group = row.get("candidate_group", "").strip()
                            email = row.get("email", "").strip()
                            
                            if not pseudo_id:
                                error_message = f"Row {row_num}: pseudo_id is required"
                                break
                            
                            participant, created = ResearchParticipant.objects.get_or_create(
                                study=study,
                                pseudo_id=pseudo_id,
                                defaults={
                                    "candidate_group": candidate_group,
                                    "email": email,
                                }
                            )
                            if created:
                                generated_count += 1
                        except Exception as e:
                            error_message = f"Row {row_num}: {str(e)}"
                            break

                elif upload_format == "json":
                    # Parse JSON
                    data = json.load(uploaded_file)
                    if not isinstance(data, list):
                        error_message = "JSON must be an array of participant objects"
                    else:
                        for row_num, row in enumerate(data, start=1):
                            try:
                                pseudo_id = row.get("pseudo_id", "").strip()
                                candidate_group = row.get("candidate_group", "").strip()
                                email = row.get("email", "").strip()
                                
                                if not pseudo_id:
                                    error_message = f"Record {row_num}: pseudo_id is required"
                                    break
                                
                                participant, created = ResearchParticipant.objects.get_or_create(
                                    study=study,
                                    pseudo_id=pseudo_id,
                                    defaults={
                                        "candidate_group": candidate_group,
                                        "email": email,
                                    }
                                )
                                if created:
                                    generated_count += 1
                            except Exception as e:
                                error_message = f"Record {row_num}: {str(e)}"
                                break

                if error_message:
                    messages.error(request, error_message)
                else:
                    messages.success(request, f"Generated {generated_count} new participants")
                    return redirect("research:study_detail", pk=study.pk)

            except Exception as e:
                error_message = f"File processing error: {str(e)}"
                messages.error(request, error_message)
    else:
        form = BulkParticipantGenerationForm()

    return render(
        request,
        "research/bulk_generate.html",
        {"study": study, "form": form},
    )


# ============================================================================
# Participant Entry & Access
# ============================================================================

def participant_entry(request, study_id: int, pseudo_id: str):
    """Participant landing page with on-site randomisation and intake form."""
    study = get_object_or_404(ResearchStudy, pk=study_id)
    if not study.active and not (request.user.is_authenticated and _user_can_manage_research_study(study, request.user)):
        raise Http404("Study not available")

    # Check generation mode
    if study.generation_mode == ResearchStudy.GenerationMode.BULK_ONLY:
        # Try to get existing participant
        try:
            participant = ResearchParticipant.objects.get(study=study, pseudo_id=pseudo_id)
        except ResearchParticipant.DoesNotExist:
            if request.user.is_authenticated and _user_can_manage_research_study(study, request.user):
                participant = ResearchParticipant.objects.create(study=study, pseudo_id=pseudo_id)
            else:
                raise Http404("Participant not found. Contact study administrator.")
    else:
        # AUTO_ON_ACCESS: create if doesn't exist
        participant, _ = ResearchParticipant.objects.get_or_create(
            study=study,
            pseudo_id=pseudo_id,
        )

    if request.user.is_authenticated and participant.user_user is None:
        participant.user_user = request.user

    if request.method == "POST":
        form = ResearchParticipantIntakeForm(
            request.POST,
            instance=participant,
            collect_demographics=study.collect_demographics,
        )
        if form.is_valid():
            participant = form.save(commit=False)
            participant.study = study

            with transaction.atomic():
                if participant.assigned_arm is None:
                    assigned_arm = _choose_research_arm(
                        study,
                        candidate_group=participant.candidate_group,
                    )
                    participant.assigned_arm = assigned_arm
                    participant.assignment_method = study.randomisation_mode
                    participant.assigned_at = timezone.now()

                participant.save()
                cid_user = _ensure_research_participant_cid(participant)
                
                # Store email on CidUser if provided
                if participant.email and not cid_user.email:
                    cid_user.email = participant.email
                    cid_user.save(update_fields=['email'])

                # Ensure packet can be accessed via existing CID flow
                if participant.assigned_arm is not None:
                    participant.assigned_arm.packet.valid_cid_users.add(cid_user)

            messages.success(request, "Intake complete. You can now start the study.")
            return redirect(
                "research:participant_entry",
                study_id=study.pk,
                pseudo_id=participant.pseudo_id,
            )
    else:
        form = ResearchParticipantIntakeForm(
            instance=participant,
            collect_demographics=study.collect_demographics,
        )

    # Build start URL if assigned
    assigned_packet = participant.assigned_arm.packet if participant.assigned_arm else None
    start_url = None
    if participant.assigned_arm and participant.cid_user:
        start_url = reverse(
            "atlas:collection_case_view_take",
            kwargs={
                "pk": assigned_packet.pk,
                "case_number": 0,
                "cid": participant.cid_user.cid,
                "passcode": participant.cid_user.passcode,
            },
        )

    return render(
        request,
        "research/participant_entry.html",
        {
            "study": study,
            "participant": participant,
            "form": form,
            "assigned_packet": assigned_packet,
            "start_url": start_url,
        },
    )
