<?php
namespace Isekai\AIToolbox\Api\AIToolboxBot;

use ApiBase;
use Isekai\AIToolbox\Api\ApiAIToolbox;
use Isekai\AIToolbox\Api\ApiAIToolboxBot;
use Isekai\AIToolbox\AIToolboxUtils;
use MediaWiki\MediaWikiServices;
use Wikimedia\ParamValidator\ParamValidator;
use Message;

class ApiReportUsage extends ApiBase {
    const TRANSACTION_TIMEOUT = 60 * 10;

    /** @var MediaWikiServices */
    private $services;

    /** @var \WANObjectCache */
    private $cache;

    /** @var ApiAIToolboxBot */
    private $mParent;

    public function __construct(ApiAIToolboxBot $main, $method) {
        parent::__construct($main->getMain(), $method);

        $this->mParent = $main;
        $this->services = MediaWikiServices::getInstance();
        $this->cache = $this->services->getMainWANObjectCache();
    }

    public function requireParameter(array $paramKeys) {
        $missing = [];
        foreach ($paramKeys as $paramKey) {
            if ($this->getParameter($paramKey) === null) {
                $missing[] = $paramKey;
            }
        }

        if (count($missing) > 0) {
            $this->dieWithError( [
                'apierror-missingparam',
                Message::listParam( array_map(
                    function ( $p ) {
                        return '<var>' . $this->encodeParamName( $p ) . '</var>';
                    },
                    $missing
                ) ),
                count( $missing ),
            ], 'missingparam' );
        }
    }

    public function execute() {
        // Check permission
        $this->checkUserRightsAny('aitoolboxbot');

        $step = $this->getParameter('step');

        $result = $this->getResult();
        if ($step) {
            switch ($step) {
                case 'start':
                    $this->requireParameter(['userid', 'useraction']);

                    $userId = $this->getParameter('userid');
                    $userAction = $this->getParameter('useraction');
                    $botId = $this->getParameter('botid');
                    $tokens = $this->getParameter('tokens');
                    $pointUsage = $this->getParameter('pointusage');

                    $user = $this->services->getUserFactory()->newFromId($userId);

                    if (!$user->isRegistered()) {
                        $this->addError('apierror-isekai-ai-toolbox-user-not-exists', 'user-not-exists');
                        return false;
                    }

                    $permissionManager = $this->services->getPermissionManager();
                    $permissionMap = ApiAIToolbox::$permissionMap;

                    $permissionKey = $permissionMap[$userAction] ?? null;
                    if ($permissionKey && !$user->isAllowedAny($permissionKey)) {
                        $this->addError('apierror-isekai-ai-toolbox-nopermission', 'nopermission', [
                            'permission' => $permissionKey,
                            'user' => $user->getName(),
                            'user_permissions' => $permissionManager->getUserPermissions($this->getUser()),
                        ]);
                        return false;
                    }
                    
                    $noCost = false;
                    if ($user->isAllowedAny('aitoolbox-unlimited')) {
                        $noCost = true;
                    }

                    $transactionid = bin2hex(random_bytes(8));

                    $pointType = MediaWikiServices::getInstance()->getMainConfig()->get('IsekaiAIToolboxPointType');

                    $data = [
                        'userid' => $userId,
                        'useraction' => $userAction,
                        'pointtype' => $pointType,
                        'pointusage' => $pointUsage,
                        'botid' => $botId,
                        'tokens' => $tokens,
                        'step' => $step,
                    ];

                    if (!$noCost) {
                        /** @var \Isekai\UserPoints\Service\IsekaiUserPointsFactory */
                        $pointFactory = $this->services->getService('IsekaiUserPoints');
                        $pointService = $pointFactory->newFromUserId($userId, $pointType);

                        if (!$pointService->hasEnoughPoints($pointUsage)) { // User doesn't have enough points
                            $this->addError('apierror-isekai-ai-toolbox-noenoughpoints', 'noenoughpoints');
                            return false;
                        }

                        $pointTransactionId = $pointService->startConsumePointsTransaction($pointUsage, 600, 'isekai-ai-toolbox', $userAction);
                        $data['pointtransaction'] = $pointTransactionId;
                    }

                    $this->cache->set(
                        $this->cache->makeKey('aitoolboxbot', 'reportusage', 'transaction', $transactionid),
                        $data,
                        self::TRANSACTION_TIMEOUT
                    );

                    $result->addValue(['aitoolboxbot', $this->getModuleName()], 'success', 1);
                    $result->addValue(['aitoolboxbot', $this->getModuleName()], 'transactionid', $transactionid);
                    break;
                case 'end':
                    $transactionId = $this->getParameter('transactionid');
                    $transactionData = $this->cache->get(
                        $this->cache->makeKey('aitoolboxbot', 'reportusage', 'transaction', $transactionId)
                    );

                    $isSuccess = true;

                    if ($transactionData) {
                        $pointTransactionId = $transactionData['pointtransaction'] ?? null;
                        $pointType = $transactionData['pointtype'] ?? null;
                        if ($pointTransactionId && $pointType) {
                            /** @var \Isekai\UserPoints\Service\IsekaiUserPointsFactory */
                            $pointFactory = $this->services->getService('IsekaiUserPoints');
                            $pointService = $pointFactory->newFromUserId($transactionData['userid'], $pointType);
    
                            $ret = $pointService->commitConsumePointsTransaction($pointTransactionId);
                            if (!$ret) {
                                $this->addWarning('apiwarn-isekai-ai-toolbox-pointtransactionfailed', 'pointtransactionfailed');
                                $isSuccess = false;
                            }
                        }

                        $realTokens = $this->getParameter('tokens') ?? $transactionData['tokens'];

                        $dbw = $this->services->getDBLoadBalancer()->getConnection(DB_PRIMARY);
                        $dbw->insert(
                            'aitoolbox_usage',
                            [
                                'user_id' => $transactionData['userid'],
                                'action' => $transactionData['useraction'],
                                'bot_id' => $transactionData['botid'],
                                'tokens' => $transactionData['tokens'],
                                'real_tokens' => $realTokens,
                                'timestamp' => $dbw->timestamp(),
                                'success' => 1,
                            ],
                            __METHOD__
                        );

                        $this->cache->delete(
                            $this->cache->makeKey('aitoolboxbot', 'reportusage', 'transaction', $transactionId)
                        );
                    } else {
                        $this->addWarning('apierror-isekai-ai-toolbox-invalidtransaction', 'invalidtransaction');
                        $isSuccess = false;
                    }

                    $result->addValue(['aitoolboxbot', $this->getModuleName()], 'success', $isSuccess ? 1 : 0);
                    break;
                case 'cancel':
                    $transactionId = $this->getParameter('transactionid');
                    $transactionData = $this->cache->get(
                        $this->cache->makeKey('aitoolboxbot', 'reportusage', 'transaction', $transactionId)
                    );

                    $isSuccess = true;

                    if ($transactionData) {
                        $pointTransactionId = $transactionData['pointtransaction'] ?? null;
                        $pointType = $transactionData['pointtype'] ?? null;
                        if ($pointTransactionId && $pointType) {
                            /** @var \Isekai\UserPoints\Service\IsekaiUserPointsFactory */
                            $pointFactory = $this->services->getService('IsekaiUserPoints');
                            $pointService = $pointFactory->newFromUserId($transactionData['userid'], $pointType);
    
                            $ret = $pointService->rollbackConsumePointsTransaction($pointTransactionId);
                            if (!$ret) {
                                $this->addWarning('apiwarn-isekai-ai-toolbox-pointtransactionfailed', 'pointtransactionfailed');
                                $isSuccess = false;
                            }
                        }

                        $reportedError = $this->getParameter('error');

                        $dbw = $this->services->getDBLoadBalancer()->getConnection(DB_PRIMARY);
                        $dbw->insert(
                            'aitoolbox_usage',
                            [
                                'user_id' => $transactionData['userid'],
                                'action' => $transactionData['useraction'],
                                'bot_id' => $transactionData['botid'],
                                'tokens' => $transactionData['tokens'],
                                'timestamp' => $dbw->timestamp(),
                                'success' => 0,
                                'error' => $reportedError,
                            ],
                            __METHOD__
                        );

                        $this->cache->delete(
                            $this->cache->makeKey('aitoolboxbot', 'reportusage', 'transaction', $transactionId)
                        );
                    } else {
                        $this->addWarning('apierror-isekai-ai-toolbox-invalidtransaction', 'invalidtransaction');
                        $isSuccess = false;
                    }

                    $result->addValue(['aitoolboxbot', $this->getModuleName()], 'success', $isSuccess ? 1 : 0);
                    break;
            }
        } else {
            $this->requireParameter(['userid', 'useraction']);

            $userId = $this->getParameter('userid');
            $userAction = strtolower($this->getParameter('useraction'));
            $botId = $this->getParameter('botid');
            $realTokens = $tokens = $this->getParameter('tokens');
            $pointUsage =  $this->getParameter('pointusage');

            $user = $this->services->getUserFactory()->newFromId($userId);

            if (!$user->isRegistered()) {
                $this->addError('apierror-isekai-ai-toolbox-nopermission', 'nopermission');
                return false;
            }

            $permissionMap = ApiAIToolbox::$permissionMap;

            $permissionKey = $permissionMap[$userAction] ?? null;
            if ($permissionKey && !$user->isAllowed($permissionKey)) {
                $this->addError('apierror-isekai-ai-toolbox-nopermission', 'nopermission');
                return false;
            }
            
            $noCost = false;
            if ($userAction === 'chatcomplete' && $user->isAllowed('aitoolbox-unlimited')) {
                $noCost = true;
            }

            $isSuccess = true;
            
            $pointType = MediaWikiServices::getInstance()->getMainConfig()->get('IsekaiAIToolboxPointType');

            if (!$noCost) {
                /** @var \Isekai\UserPoints\Service\IsekaiUserPointsFactory */
                $pointFactory = $this->services->getService('IsekaiUserPoints');
                $pointService = $pointFactory->newFromUserId($userId, $pointType);

                if (!$pointService->hasEnoughPoints($pointUsage)) { // User doesn't have enough points
                    $this->addError('apierror-isekai-ai-toolbox-notenoughpoints', 'notenoughpoints');
                    return false;
                }

                $ret = $pointService->consumePoints($pointUsage, true, false, 'isekai-ai-toolbox', $userAction);
                if (!$ret) {
                    $this->addWarning('apiwarn-isekai-ai-toolbox-pointtransactionfailed', 'pointtransactionfailed');
                    $isSuccess = false;
                }
            }

            $reportedError = $this->getParameter('error');

            $dbw = $this->services->getDBLoadBalancer()->getConnection(DB_PRIMARY);
            $dbw->insert(
                'aitoolbox_usage',
                [
                    'user_id' => $userId,
                    'action' => $userAction,
                    'bot_id' => $botId,
                    'tokens' => $tokens,
                    'real_tokens' => $realTokens,
                    'timestamp' => $dbw->timestamp(),
                    'success' => $reportedError ? 0 : 1,
                    'error' => $reportedError,
                ],
                __METHOD__
            );

            $result->addValue(['aitoolboxbot', $this->getModuleName()], 'success', $isSuccess ? 1 : 0);
        }

        return true;
    }

    public function getAllowedParams($flags = 0) {
        return [
            'userid' => [
                ParamValidator::PARAM_TYPE => 'integer',
                ParamValidator::PARAM_DEFAULT => null,
                ParamValidator::PARAM_REQUIRED => false,
            ],
            'useraction' => [
                ParamValidator::PARAM_TYPE => [
                    'embeddingpage',
                    'chatcomplete',
                ],
                ParamValidator::PARAM_DEFAULT => null,
                ParamValidator::PARAM_REQUIRED => false,
            ],
            'botid' => [
                ParamValidator::PARAM_TYPE => 'string',
                ParamValidator::PARAM_DEFAULT => null,
                ParamValidator::PARAM_REQUIRED => false,
            ],
            'tokens' => [
                ParamValidator::PARAM_TYPE => 'integer',
                ParamValidator::PARAM_DEFAULT => 100,
                ParamValidator::PARAM_REQUIRED => false,
            ],
            'pointusage' => [
                ParamValidator::PARAM_TYPE => 'integer',
                ParamValidator::PARAM_DEFAULT => 0,
                ParamValidator::PARAM_REQUIRED => false,
            ],
            'error' => [
                ParamValidator::PARAM_TYPE => 'string',
                ParamValidator::PARAM_DEFAULT => null,
                ParamValidator::PARAM_REQUIRED => false,
            ],
            'step' => [
                ParamValidator::PARAM_TYPE => [
                    'start',
                    'end',
                    'cancel',
                ],
                ParamValidator::PARAM_DEFAULT => null,
                ParamValidator::PARAM_REQUIRED => false,
            ],
            'transactionid' => [
                ParamValidator::PARAM_TYPE => 'string',
                ParamValidator::PARAM_DEFAULT => null,
                ParamValidator::PARAM_REQUIRED => false,
            ],
        ];
    }

    public function getCacheMode($params) {
		return 'private';
	}

    public function getParent() {
        return $this->mParent;
    }
}