App/Payments/PaymentsRoutes.py

from fastapi import APIRouter, HTTPException, status
from typing import List
from .Model import Payment
from App.Users.Model import User
from App.Plans.Model import Plan
from .Schema import (
    CreatePaymentRequest,
    UpdatePaymentStatusRequest,
    PaymentResponse,
    UpdatePaymentUserRequest,
    BaseResponse,
    PaymentListResponse,
)
from .Schema import PaymentMethod

payment_router = APIRouter(tags=["Payments"])


@payment_router.post("/payment/create", response_model=BaseResponse)
async def create_payment(request: CreatePaymentRequest, internal=False):
    # If payment method is "Lipa Number", transaction_id is required
    if (
        request.payment_method == PaymentMethod.LIPA_NUMBER
        and not request.transaction_id
    ):
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Transaction ID is required for Lipa Number payments",
        )

    # Check if plan exists
    plan = await Plan.get_or_none(id=request.plan_id)
    if request.plan_id and not plan:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND, detail="Plan not found"
        )

    # Create new payment without linking a user initially
    payment = await Payment.create(
        user_id=request.user_id,
        plan=plan,
        amount=request.amount,
        payment_method=request.payment_method,
        transaction_id=request.transaction_id,
        status="pending",  # Default status
    )

    if internal:
        return payment

    # If payment method is "cash", attempt to create a subscription
    if payment.payment_method == PaymentMethod.CASH:
        await payment.create_subscription_if_cash()

    await payment.save()

    return BaseResponse(
        code=200,
        message="Payment created successfully",
        payload={"payment_id": str(payment.id)},
    )


@payment_router.get("/payment/{payment_id}", response_model=PaymentResponse)
async def get_payment(payment_id: str):
    payment = await Payment.get_or_none(id=payment_id)
    if not payment:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND, detail="Payment not found"
        )

    return PaymentResponse(
        id=str(payment.id),
        user_id=payment.user_id,
        plan_id=payment.plan_id,
        amount=payment.amount,
        payment_method=payment.payment_method,
        status=payment.status,
        transaction_id=payment.transaction_id,
        created_time=payment.created_time,
        updated_time=payment.updated_time,
    )


@payment_router.get("/payments/unlinked", response_model=PaymentListResponse)
async def get_unlinked_payments():
    unlinked_payments = await Payment.filter(user_id=None)
    result = [
        PaymentResponse(
            id=str(payment.id),
            user_id=payment.user_id,
            plan_id=payment.plan_id,
            amount=payment.amount,
            payment_method=payment.payment_method,
            status=payment.status,
            transaction_id=payment.transaction_id,
            created_time=payment.created_time,
            updated_time=payment.updated_time,
        )
        for payment in unlinked_payments
    ]

    return PaymentListResponse(payments=result, total_count=len(result))


@payment_router.put("/payment/{payment_id}/link-user", response_model=BaseResponse)
async def link_user_to_payment(payment_id: str, request: UpdatePaymentUserRequest):
    payment = await Payment.get_or_none(id=payment_id)
    if not payment:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND, detail="Payment not found"
        )

    user = await User.get_or_none(id=request.user_id)
    if not user:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND, detail="User not found"
        )

    # Link payment to the user
    payment.user_id = request.user_id
    await payment.save()

    # If the payment method is "Lipa Number" or "M-Pesa", add to user's balance directly
    if (
        payment.payment_method in [PaymentMethod.LIPA_NUMBER, PaymentMethod.MPESA]
        and not payment.plan
    ):
        user.balance += payment.amount
        await user.save()

        # Update payment status to indicate it was assigned to balance
        payment.status = "balance-assigned"
        await payment.save()

        return BaseResponse(
            code=200,
            message="Payment linked to user and amount added to balance successfully",
            payload={
                "payment_id": str(payment.id),
                "user_id": request.user_id,
                "new_balance": user.balance,
            },
        )

    return BaseResponse(
        code=200,
        message="Payment linked to user successfully",
        payload={"payment_id": str(payment.id), "user_id": request.user_id},
    )