#
#  Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
import random
import re

import pytest
from test_common import create_memory
from configs import INVALID_API_TOKEN
from libs.auth import RAGFlowWebApiAuth


class TestAuthorization:
    @pytest.mark.p2
    @pytest.mark.parametrize(
        "invalid_auth, expected_code, expected_message",
        [
            (None, 401, "<Unauthorized '401: Unauthorized'>"),
            (RAGFlowWebApiAuth(INVALID_API_TOKEN), 401, "<Unauthorized '401: Unauthorized'>"),
        ],
        ids=["empty_auth", "invalid_api_token"]
    )
    def test_auth_invalid(self, invalid_auth, expected_code, expected_message):
        res = create_memory(invalid_auth)
        assert res["code"] == expected_code, res
        assert res["message"] == expected_message, res


class TestMemoryCreate:
    @pytest.mark.p1
    @pytest.mark.parametrize("name", ["test_memory_name", "d" * 128])
    def test_name(self, WebApiAuth, name):
        payload = {
            "name": name,
            "memory_type": ["raw"] + random.choices(["semantic", "episodic", "procedural"], k=random.randint(0, 3)),
            "embd_id": "BAAI/bge-small-en-v1.5@Builtin",
            "llm_id": "glm-4-flash@ZHIPU-AI"
        }
        res = create_memory(WebApiAuth, payload)
        assert res["code"] == 0, res
        pattern = rf'^{name}|{name}(?:\((\d+)\))?$'
        escaped_name = re.escape(res["data"]["name"])
        assert re.match(pattern, escaped_name), res

    @pytest.mark.p2
    @pytest.mark.parametrize(
        "name, expected_message",
        [
            ("", "Memory name cannot be empty or whitespace."),
            (" ", "Memory name cannot be empty or whitespace."),
            ("a" * 129, f"Memory name '{'a'*129}' exceeds limit of 128."),
        ],
        ids=["empty_name", "space_name", "too_long_name"],
    )
    def test_name_invalid(self, WebApiAuth, name, expected_message):
        payload = {
            "name": name,
            "memory_type": ["raw"] + random.choices(["semantic", "episodic", "procedural"], k=random.randint(0, 3)),
            "embd_id": "BAAI/bge-small-en-v1.5@Builtin",
            "llm_id": "glm-4-flash@ZHIPU-AI"
        }
        res = create_memory(WebApiAuth, payload)
        assert res["message"] == expected_message, res

    @pytest.mark.p2
    @pytest.mark.parametrize("name", ["invalid_type_name", "memory_alpha"])
    def test_type_invalid(self, WebApiAuth, name):
        payload = {
            "name": name,
            "memory_type": ["something"],
            "embd_id": "BAAI/bge-small-en-v1.5@Builtin",
            "llm_id": "glm-4-flash@ZHIPU-AI"
        }
        res = create_memory(WebApiAuth, payload)
        assert res["message"] == f"Memory type '{ {'something'} }' is not supported.", res

    @pytest.mark.p3
    def test_name_duplicated(self, WebApiAuth):
        name = "duplicated_name_test"
        payload = {
            "name": name,
            "memory_type": ["raw"] + random.choices(["semantic", "episodic", "procedural"], k=random.randint(0, 3)),
            "embd_id": "BAAI/bge-small-en-v1.5@Builtin",
            "llm_id": "glm-4-flash@ZHIPU-AI"
        }
        res1 = create_memory(WebApiAuth, payload)
        assert res1["code"] == 0, res1

        res2 = create_memory(WebApiAuth, payload)
        assert res2["code"] == 0, res2