Lotty/src/backend/api/v1/experiments/endpoints.py

from http import HTTPStatus
from uuid import UUID

from django.http import Http404, HttpRequest
from django.shortcuts import get_object_or_404
from ninja import Router
from ninja.errors import AuthenticationError

from api.v1.experiments.schemas import (
    ApprovalOut,
    ApproveIn,
    CompleteIn,
    ExperimentCreateIn,
    ExperimentListOut,
    ExperimentOut,
    ExperimentUpdateIn,
    LogOut,
    OutcomeOut,
    PauseIn,
    RejectIn,
    RequestChangesIn,
    VariantCreateIn,
    VariantOut,
    VariantUpdateIn,
)
from apps.experiments.models import (
    Experiment,
    ExperimentOutcome,
    ExperimentStatus,
    Variant,
)
from apps.experiments.selectors import (
    experiment_approvals,
    experiment_get,
    experiment_list,
    experiment_logs,
    variant_list,
)
from apps.experiments.services import (
    experiment_approve,
    experiment_archive,
    experiment_complete,
    experiment_create,
    experiment_pause,
    experiment_reject,
    experiment_reopen,
    experiment_request_changes,
    experiment_resume,
    experiment_start,
    experiment_submit_for_review,
    experiment_update,
    variant_create,
    variant_delete,
    variant_update,
)
from apps.flags.models import FeatureFlag
from apps.users.auth.bearer import (
    jwt_bearer,
    require_admin_or_approver,
    require_admin_or_experimenter,
)
from apps.users.models import User

router = Router(tags=["experiments"], auth=jwt_bearer)


def _get_user(request: HttpRequest) -> User:
    user = getattr(request, "auth", None)
    if not isinstance(user, User):
        raise AuthenticationError
    return user


@router.get(
    "",
    response={HTTPStatus.OK: ExperimentListOut},
    summary="List experiments",
)
def list_experiments(
    request: HttpRequest,
    status: ExperimentStatus | None = None,
    flag_id: UUID | None = None,
    search: str | None = None,
    limit: int = 50,
    offset: int = 0,
) -> tuple[HTTPStatus, ExperimentListOut]:
    qs = experiment_list(
        status=status,
        flag_id=flag_id,
        search=search,
    )
    total = qs.count()
    items = [
        ExperimentOut.from_experiment(exp)
        for exp in qs[offset : offset + limit]
    ]
    return HTTPStatus.OK, ExperimentListOut(count=total, items=items)


@router.post(
    "",
    response={HTTPStatus.CREATED: ExperimentOut},
    summary="Create experiment",
)
@require_admin_or_experimenter
def create_experiment(
    request: HttpRequest,
    payload: ExperimentCreateIn,
) -> tuple[HTTPStatus, ExperimentOut]:
    user = _get_user(request)
    flag = get_object_or_404(FeatureFlag, pk=payload.flag_id)
    exp = experiment_create(
        flag=flag,
        name=payload.name,
        owner=user,
        description=payload.description,
        hypothesis=payload.hypothesis,
        traffic_allocation=payload.traffic_allocation,
        targeting_rules=payload.targeting_rules,
    )
    exp = experiment_get(exp.pk)
    return HTTPStatus.CREATED, ExperimentOut.from_experiment(exp)


@router.get(
    "/{experiment_id}",
    response={HTTPStatus.OK: ExperimentOut},
    summary="Get experiment",
)
def get_experiment(
    request: HttpRequest,
    experiment_id: UUID,
) -> tuple[HTTPStatus, ExperimentOut]:
    exp = experiment_get(experiment_id)
    if not exp:
        raise Http404
    return HTTPStatus.OK, ExperimentOut.from_experiment(exp)


@router.patch(
    "/{experiment_id}",
    response={HTTPStatus.OK: ExperimentOut},
    summary="Update experiment",
)
@require_admin_or_experimenter
def update_experiment(
    request: HttpRequest,
    experiment_id: UUID,
    payload: ExperimentUpdateIn,
) -> tuple[HTTPStatus, ExperimentOut]:
    user = _get_user(request)
    exp = get_object_or_404(
        Experiment.objects.select_related("flag", "owner"),
        pk=experiment_id,
    )
    fields = payload.model_dump(exclude_none=True)
    exp = experiment_update(experiment=exp, user=user, **fields)
    exp = experiment_get(exp.pk)
    return HTTPStatus.OK, ExperimentOut.from_experiment(exp)


@router.get(
    "/{experiment_id}/variants",
    response={HTTPStatus.OK: list[VariantOut]},
    summary="List variants",
)
def list_variants(
    request: HttpRequest,
    experiment_id: UUID,
) -> tuple[HTTPStatus, list[VariantOut]]:
    get_object_or_404(Experiment, pk=experiment_id)
    variants = variant_list(experiment_id)
    return HTTPStatus.OK, [VariantOut.model_validate(v) for v in variants]


@router.post(
    "/{experiment_id}/variants",
    response={HTTPStatus.CREATED: VariantOut},
    summary="Create variant",
)
@require_admin_or_experimenter
def create_variant(
    request: HttpRequest,
    experiment_id: UUID,
    payload: VariantCreateIn,
) -> tuple[HTTPStatus, VariantOut]:
    user = _get_user(request)
    exp = get_object_or_404(
        Experiment.objects.select_related("flag", "owner"),
        pk=experiment_id,
    )
    v = variant_create(
        experiment=exp,
        user=user,
        name=payload.name,
        value=payload.value,
        weight=payload.weight,
        is_control=payload.is_control,
    )
    return HTTPStatus.CREATED, VariantOut.model_validate(v)


@router.patch(
    "/{experiment_id}/variants/{variant_id}",
    response={HTTPStatus.OK: VariantOut},
    summary="Update variant",
)
@require_admin_or_experimenter
def update_variant(
    request: HttpRequest,
    experiment_id: UUID,
    variant_id: UUID,
    payload: VariantUpdateIn,
) -> tuple[HTTPStatus, VariantOut]:
    user = _get_user(request)
    v = get_object_or_404(
        Variant.objects.select_related(
            "experiment__flag", "experiment__owner"
        ),
        pk=variant_id,
        experiment_id=experiment_id,
    )
    v = variant_update(
        variant=v,
        user=user,
        **payload.model_dump(exclude_none=True),
    )
    return HTTPStatus.OK, VariantOut.model_validate(v)


@router.delete(
    "/{experiment_id}/variants/{variant_id}",
    response={HTTPStatus.NO_CONTENT: None},
    summary="Delete variant",
)
@require_admin_or_experimenter
def delete_variant(
    request: HttpRequest,
    experiment_id: UUID,
    variant_id: UUID,
) -> tuple[HTTPStatus, None]:
    user = _get_user(request)
    v = get_object_or_404(
        Variant.objects.select_related("experiment__owner"),
        pk=variant_id,
        experiment_id=experiment_id,
    )
    variant_delete(variant=v, user=user)
    return HTTPStatus.NO_CONTENT, None


@router.post(
    "/{experiment_id}/submit-for-review",
    response={HTTPStatus.OK: ExperimentOut},
    summary="Submit experiment for review",
)
@require_admin_or_experimenter
def submit_for_review(
    request: HttpRequest,
    experiment_id: UUID,
) -> tuple[HTTPStatus, ExperimentOut]:
    user = _get_user(request)
    exp = get_object_or_404(
        Experiment.objects.select_related("flag", "owner"),
        pk=experiment_id,
    )
    exp = experiment_submit_for_review(experiment=exp, user=user)
    exp = experiment_get(exp.pk)
    return HTTPStatus.OK, ExperimentOut.from_experiment(exp)


@router.post(
    "/{experiment_id}/approve",
    response={HTTPStatus.OK: ExperimentOut},
    summary="Approve experiment",
)
@require_admin_or_approver
def approve(
    request: HttpRequest,
    experiment_id: UUID,
    payload: ApproveIn,
) -> tuple[HTTPStatus, ExperimentOut]:
    user = _get_user(request)
    exp = get_object_or_404(
        Experiment.objects.select_related("flag", "owner"),
        pk=experiment_id,
    )
    exp = experiment_approve(
        experiment=exp,
        approver=user,
        comment=payload.comment,
    )
    exp = experiment_get(exp.pk)
    return HTTPStatus.OK, ExperimentOut.from_experiment(exp)


@router.post(
    "/{experiment_id}/reject",
    response={HTTPStatus.OK: ExperimentOut},
    summary="Reject experiment",
)
@require_admin_or_approver
def reject(
    request: HttpRequest,
    experiment_id: UUID,
    payload: RejectIn,
) -> tuple[HTTPStatus, ExperimentOut]:
    user = _get_user(request)
    exp = get_object_or_404(
        Experiment.objects.select_related("flag", "owner"),
        pk=experiment_id,
    )
    exp = experiment_reject(
        experiment=exp,
        user=user,
        comment=payload.comment,
    )
    exp = experiment_get(exp.pk)
    return HTTPStatus.OK, ExperimentOut.from_experiment(exp)


@router.post(
    "/{experiment_id}/request-changes",
    response={HTTPStatus.OK: ExperimentOut},
    summary="Request changes (return to draft)",
)
@require_admin_or_approver
def request_changes(
    request: HttpRequest,
    experiment_id: UUID,
    payload: RequestChangesIn,
) -> tuple[HTTPStatus, ExperimentOut]:
    user = _get_user(request)
    exp = get_object_or_404(
        Experiment.objects.select_related("flag", "owner"),
        pk=experiment_id,
    )
    exp = experiment_request_changes(
        experiment=exp,
        user=user,
        comment=payload.comment,
    )
    exp = experiment_get(exp.pk)
    return HTTPStatus.OK, ExperimentOut.from_experiment(exp)


@router.post(
    "/{experiment_id}/start",
    response={HTTPStatus.OK: ExperimentOut},
    summary="Start experiment",
)
@require_admin_or_experimenter
def start(
    request: HttpRequest,
    experiment_id: UUID,
) -> tuple[HTTPStatus, ExperimentOut]:
    user = _get_user(request)
    exp = get_object_or_404(
        Experiment.objects.select_related("flag", "owner"),
        pk=experiment_id,
    )
    exp = experiment_start(experiment=exp, user=user)
    exp = experiment_get(exp.pk)
    return HTTPStatus.OK, ExperimentOut.from_experiment(exp)


@router.post(
    "/{experiment_id}/pause",
    response={HTTPStatus.OK: ExperimentOut},
    summary="Pause experiment",
)
@require_admin_or_experimenter
def pause(
    request: HttpRequest,
    experiment_id: UUID,
    payload: PauseIn,
) -> tuple[HTTPStatus, ExperimentOut]:
    user = _get_user(request)
    exp = get_object_or_404(
        Experiment.objects.select_related("flag", "owner"),
        pk=experiment_id,
    )
    exp = experiment_pause(
        experiment=exp,
        user=user,
        comment=payload.comment,
    )
    exp = experiment_get(exp.pk)
    return HTTPStatus.OK, ExperimentOut.from_experiment(exp)


@router.post(
    "/{experiment_id}/resume",
    response={HTTPStatus.OK: ExperimentOut},
    summary="Resume experiment",
)
@require_admin_or_experimenter
def resume(
    request: HttpRequest,
    experiment_id: UUID,
) -> tuple[HTTPStatus, ExperimentOut]:
    user = _get_user(request)
    exp = get_object_or_404(
        Experiment.objects.select_related("flag", "owner"),
        pk=experiment_id,
    )
    exp = experiment_resume(experiment=exp, user=user)
    exp = experiment_get(exp.pk)
    return HTTPStatus.OK, ExperimentOut.from_experiment(exp)


@router.post(
    "/{experiment_id}/complete",
    response={HTTPStatus.OK: ExperimentOut},
    summary="Complete experiment with outcome",
)
@require_admin_or_experimenter
def complete(
    request: HttpRequest,
    experiment_id: UUID,
    payload: CompleteIn,
) -> tuple[HTTPStatus, ExperimentOut]:
    user = _get_user(request)
    exp = get_object_or_404(
        Experiment.objects.select_related("flag", "owner"),
        pk=experiment_id,
    )
    exp = experiment_complete(
        experiment=exp,
        user=user,
        outcome=payload.outcome,
        rationale=payload.rationale,
        winning_variant_id=(
            str(payload.winning_variant_id)
            if payload.winning_variant_id
            else None
        ),
    )
    exp = experiment_get(exp.pk)
    return HTTPStatus.OK, ExperimentOut.from_experiment(exp)


@router.post(
    "/{experiment_id}/archive",
    response={HTTPStatus.OK: ExperimentOut},
    summary="Archive experiment",
)
@require_admin_or_experimenter
def archive(
    request: HttpRequest,
    experiment_id: UUID,
) -> tuple[HTTPStatus, ExperimentOut]:
    user = _get_user(request)
    exp = get_object_or_404(
        Experiment.objects.select_related("flag", "owner"),
        pk=experiment_id,
    )
    exp = experiment_archive(experiment=exp, user=user)
    exp = experiment_get(exp.pk)
    return HTTPStatus.OK, ExperimentOut.from_experiment(exp)


@router.post(
    "/{experiment_id}/reopen",
    response={HTTPStatus.OK: ExperimentOut},
    summary="Reopen rejected experiment",
)
@require_admin_or_experimenter
def reopen(
    request: HttpRequest,
    experiment_id: UUID,
) -> tuple[HTTPStatus, ExperimentOut]:
    user = _get_user(request)
    exp = get_object_or_404(
        Experiment.objects.select_related("flag", "owner"),
        pk=experiment_id,
    )
    exp = experiment_reopen(experiment=exp, user=user)
    exp = experiment_get(exp.pk)
    return HTTPStatus.OK, ExperimentOut.from_experiment(exp)


@router.get(
    "/{experiment_id}/logs",
    response={HTTPStatus.OK: list[LogOut]},
    summary="Get experiment audit logs",
)
def get_logs(
    request: HttpRequest,
    experiment_id: UUID,
) -> tuple[HTTPStatus, list[LogOut]]:
    get_object_or_404(Experiment, pk=experiment_id)
    logs = experiment_logs(experiment_id)
    return HTTPStatus.OK, [LogOut.from_log(log) for log in logs]


@router.get(
    "/{experiment_id}/approvals",
    response={HTTPStatus.OK: list[ApprovalOut]},
    summary="Get experiment approvals",
)
def get_approvals(
    request: HttpRequest,
    experiment_id: UUID,
) -> tuple[HTTPStatus, list[ApprovalOut]]:
    get_object_or_404(Experiment, pk=experiment_id)
    approvals = experiment_approvals(experiment_id)
    return HTTPStatus.OK, [ApprovalOut.from_approval(a) for a in approvals]


@router.get(
    "/{experiment_id}/outcome",
    response={HTTPStatus.OK: OutcomeOut},
    summary="Get experiment outcome",
)
def get_outcome(
    request: HttpRequest,
    experiment_id: UUID,
) -> tuple[HTTPStatus, OutcomeOut]:
    exp = get_object_or_404(Experiment, pk=experiment_id)
    outcome = (
        ExperimentOutcome.objects.filter(
            experiment=exp,
        )
        .select_related("winning_variant", "decided_by")
        .first()
    )
    if not outcome:
        raise Http404
    return HTTPStatus.OK, OutcomeOut.from_outcome(outcome)