#
#  Copyright 2026 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 string

import pytest
from test_common import list_datasets, upload_documents
from configs import DOCUMENT_NAME_LIMIT, INVALID_API_TOKEN
from libs.auth import RAGFlowWebApiAuth
from utils.file_utils import create_txt_file

from concurrent.futures import ThreadPoolExecutor, as_completed


@pytest.mark.p1
@pytest.mark.usefixtures("clear_datasets")
class TestAuthorization:
    @pytest.mark.parametrize(
        "invalid_auth, expected_code, expected_message",
        [
            (None, 401, "<Unauthorized '401: Unauthorized'>"),
            (RAGFlowWebApiAuth(INVALID_API_TOKEN), 401, "<Unauthorized '401: Unauthorized'>"),
        ],
    )
    def test_invalid_auth(self, invalid_auth, expected_code, expected_message):
        res = upload_documents(invalid_auth)
        assert res["code"] == expected_code, res
        assert res["message"] == expected_message, res


class TestDocumentsUpload:
    @pytest.mark.p1
    def test_valid_single_upload(self, WebApiAuth, add_dataset_func, tmp_path):
        kb_id = add_dataset_func
        fp = create_txt_file(tmp_path / "ragflow_test.txt")
        res = upload_documents(WebApiAuth, {"kb_id": kb_id}, [fp])
        assert res["code"] == 0, res
        # New API returns "dataset_id" instead of "kb_id" due to key mapping
        assert res["data"][0]["dataset_id"] == kb_id, res
        assert res["data"][0]["name"] == fp.name, res

    @pytest.mark.p1
    @pytest.mark.parametrize(
        "generate_test_files",
        [
            "docx",
            "excel",
            "ppt",
            "image",
            "pdf",
            "txt",
            "md",
            "json",
            "eml",
            "html",
        ],
        indirect=True,
    )
    def test_file_type_validation(self, WebApiAuth, add_dataset_func, generate_test_files, request):
        kb_id = add_dataset_func
        fp = generate_test_files[request.node.callspec.params["generate_test_files"]]
        res = upload_documents(WebApiAuth, {"kb_id": kb_id}, [fp])
        assert res["code"] == 0, res
        # New API returns "dataset_id" instead of "kb_id" due to key mapping
        assert res["data"][0]["dataset_id"] == kb_id, res
        assert res["data"][0]["name"] == fp.name, res

    @pytest.mark.p3
    @pytest.mark.parametrize(
        "file_type",
        ["exe", "unknown"],
    )
    def test_unsupported_file_type(self, WebApiAuth, add_dataset_func, tmp_path, file_type):
        kb_id = add_dataset_func
        fp = tmp_path / f"ragflow_test.{file_type}"
        fp.touch()
        res = upload_documents(WebApiAuth, {"kb_id": kb_id}, [fp])
        assert res["code"] == 500, res
        assert res["message"] == f"ragflow_test.{file_type}: This type of file has not been supported yet!", res

    @pytest.mark.p2
    def test_missing_file(self, WebApiAuth, add_dataset_func):
        kb_id = add_dataset_func
        res = upload_documents(WebApiAuth, {"kb_id": kb_id})
        assert res["code"] == 101, res
        assert res["message"] == "No file part!", res

    @pytest.mark.p3
    def test_empty_file(self, WebApiAuth, add_dataset_func, tmp_path):
        kb_id = add_dataset_func
        fp = tmp_path / "empty.txt"
        fp.touch()

        res = upload_documents(WebApiAuth, {"kb_id": kb_id}, [fp])
        assert res["code"] == 0, res
        assert res["data"][0]["size"] == 0, res

    @pytest.mark.p3
    def test_filename_empty(self, WebApiAuth, add_dataset_func, tmp_path):
        kb_id = add_dataset_func

        fp = create_txt_file(tmp_path / "ragflow_test.txt")
        res = upload_documents(WebApiAuth, {"kb_id": kb_id}, [fp], filename_override="")
        assert res["code"] == 101, res
        assert res["message"] == "No file selected!", res

    @pytest.mark.p3
    def test_filename_exceeds_max_length(self, WebApiAuth, add_dataset_func, tmp_path):
        kb_id = add_dataset_func
        fp = create_txt_file(tmp_path / f"{'a' * (DOCUMENT_NAME_LIMIT - 4)}.txt")
        res = upload_documents(WebApiAuth, {"kb_id": kb_id}, [fp])
        assert res["code"] == 0, res
        assert res["data"][0]["name"] == fp.name, res

    @pytest.mark.p2
    def test_invalid_kb_id(self, WebApiAuth, tmp_path):
        fp = create_txt_file(tmp_path / "ragflow_test.txt")
        res = upload_documents(WebApiAuth, {"kb_id": "invalid_kb_id"}, [fp])
        assert res["code"] == 102, res
        assert res["message"] == "Can't find the dataset with ID invalid_kb_id!", res

    @pytest.mark.p2
    def test_duplicate_files(self, WebApiAuth, add_dataset_func, tmp_path):
        kb_id = add_dataset_func
        fp = create_txt_file(tmp_path / "ragflow_test.txt")
        res = upload_documents(WebApiAuth, {"kb_id": kb_id}, [fp, fp])
        assert res["code"] == 0, res
        assert len(res["data"]) == 2, res
        for i in range(len(res["data"])):
            # New API returns "dataset_id" instead of "kb_id" due to key mapping
            assert res["data"][i]["dataset_id"] == kb_id, res
            expected_name = fp.name
            if i != 0:
                expected_name = f"{fp.stem}({i}){fp.suffix}"
            assert res["data"][i]["name"] == expected_name, res

    @pytest.mark.p3
    def test_filename_special_characters(self, WebApiAuth, add_dataset_func, tmp_path):
        kb_id = add_dataset_func
        illegal_chars = '<>:"/\\|?*'
        translation_table = str.maketrans({char: "_" for char in illegal_chars})
        safe_filename = string.punctuation.translate(translation_table)
        fp = tmp_path / f"{safe_filename}.txt"
        fp.write_text("Sample text content")

        res = upload_documents(WebApiAuth, {"kb_id": kb_id}, [fp])
        assert res["code"] == 0, res
        assert len(res["data"]) == 1, res
        # New API returns "dataset_id" instead of "kb_id" due to key mapping
        assert res["data"][0]["dataset_id"] == kb_id, res
        assert res["data"][0]["name"] == fp.name, res

    @pytest.mark.p1
    def test_multiple_files(self, WebApiAuth, add_dataset_func, tmp_path):
        kb_id = add_dataset_func
        expected_document_count = 20
        fps = []
        for i in range(expected_document_count):
            fp = create_txt_file(tmp_path / f"ragflow_test_{i}.txt")
            fps.append(fp)
        res = upload_documents(WebApiAuth, {"kb_id": kb_id}, fps)
        assert res["code"] == 0, res

        res = list_datasets(WebApiAuth)
        assert res["data"][0]["document_count"] == expected_document_count, res

    @pytest.mark.p3
    def test_concurrent_upload(self, WebApiAuth, add_dataset_func, tmp_path):
        kb_id = add_dataset_func

        count = 20
        fps = []
        for i in range(count):
            fp = create_txt_file(tmp_path / f"ragflow_test_{i}.txt")
            fps.append(fp)

        with ThreadPoolExecutor(max_workers=5) as executor:
            futures = [executor.submit(upload_documents, WebApiAuth, {"kb_id": kb_id}, fps[i : i + 1]) for i in range(count)]
        responses = list(as_completed(futures))
        assert len(responses) == count, responses
        assert all(future.result()["code"] == 0 for future in futures), responses

        res = list_datasets(WebApiAuth)
        assert res["data"][0]["document_count"] == count, res


import asyncio
from types import SimpleNamespace


class _AwaitableValue:
    def __init__(self, value):
        self._value = value

    def __await__(self):
        async def _coro():
            return self._value

        return _coro().__await__()


class _DummyFiles(dict):
    def getlist(self, key):
        value = self.get(key, [])
        if isinstance(value, list):
            return value
        return [value]


class _DummyFile:
    def __init__(self, filename):
        self.filename = filename
        self.closed = False
        self.stream = self

    def close(self):
        self.closed = True


class _DummyRequest:
    def __init__(self, form=None, files=None, args=None):
        self._form = form or {}
        self._files = files or _DummyFiles()
        self.args = args or {}

    @property
    def form(self):
        return _AwaitableValue(self._form)

    @property
    def files(self):
        return _AwaitableValue(self._files)


def _run(coro):
    return asyncio.run(coro)


@pytest.mark.p2
class TestDocumentsUploadUnit:
    """Unit tests for document upload using upload_documents helper function"""

    def test_missing_kb_id(self, WebApiAuth, tmp_path):
        """Test that missing KB ID returns error"""
        # When kb_id is empty, the API should return an error
        fp = create_txt_file(tmp_path / "ragflow_test.txt")
        res = upload_documents(WebApiAuth, {"kb_id": ""}, [fp])
        assert res["code"] == 100
        assert res["message"] == "<MethodNotAllowed '405: Method Not Allowed'>"

    def test_missing_file_part(self, WebApiAuth, add_dataset_func):
        """Test that missing file part returns error"""
        kb_id = add_dataset_func
        # Call without files - should return error for missing file
        res = upload_documents(WebApiAuth, {"kb_id": kb_id})
        assert res["code"] == 101
        assert "file" in res["message"].lower()

    def test_empty_filename_closes_files(self, WebApiAuth, add_dataset_func, tmp_path):
        """Test that empty filename returns error"""
        kb_id = add_dataset_func
        # Create a file with empty name by using filename_override
        fp = create_txt_file(tmp_path / "ragflow_test.txt")
        res = upload_documents(WebApiAuth, {"kb_id": kb_id}, [fp], filename_override="")
        assert res["code"] == 101
        assert "file" in res["message"].lower() or "selected" in res["message"].lower()

    def test_invalid_kb_id_raises(self, WebApiAuth, tmp_path):
        """Test that invalid KB ID returns error"""
        fp = create_txt_file(tmp_path / "ragflow_test.txt")
        res = upload_documents(WebApiAuth, {"kb_id": "invalid_kb_id"}, [fp])
        # The API should return an error for invalid KB ID
        assert res["code"] == 102
        assert "Can't find the dataset" in res["message"] or "not found" in res["message"].lower()

    def test_no_permission(self, WebApiAuth, tmp_path):
        """Test that no permission returns error"""
        # Create a file and try to upload to a dataset we don't have access to
        # This test would require setting up a dataset without permission
        # For now, we skip this test as it requires specific setup
        pytest.skip("Requires dataset without permission setup")

    def test_thread_pool_errors(self, WebApiAuth, add_dataset_func, tmp_path):
        """Test that thread pool errors are handled"""
        kb_id = add_dataset_func
        # Upload a file with unsupported type
        fp = tmp_path / "test.exe"
        fp.write_text("test")
        res = upload_documents(WebApiAuth, {"kb_id": kb_id}, [fp])
        # Should return error for unsupported file type
        assert res["code"] == 500
        assert "supported" in res["message"].lower() or "type" in res["message"].lower()

    def test_empty_upload_result(self, WebApiAuth, add_dataset_func, tmp_path):
        """Test that empty upload result returns error"""
        kb_id = add_dataset_func
        # Create an empty file
        fp = tmp_path / "empty.txt"
        fp.write_text("")
        res = upload_documents(WebApiAuth, {"kb_id": kb_id}, [fp])
        # Empty file might cause issues
        # The exact behavior depends on the implementation
        # Just verify we get a response
        assert "code" in res


@pytest.mark.p2
class TestWebCrawlUnit:
    def test_invalid_url(self, document_rest_api_module, monkeypatch):
        module = document_rest_api_module
        monkeypatch.setattr(
            module,
            "request",
            _DummyRequest(form={"name": "doc", "url": "not-a-url"}, args={"type": "web"}),
        )
        res = _run(module.upload_document(dataset_id="kb1"))
        assert res["code"] == 101
        assert res["message"] == "The URL format is invalid"

    def test_invalid_kb_id(self, document_rest_api_module, monkeypatch):
        module = document_rest_api_module
        monkeypatch.setattr(module, "is_valid_url", lambda _url: True)
        monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (False, None))
        monkeypatch.setattr(
            module,
            "request",
            _DummyRequest(form={"name": "doc", "url": "http://example.com"}, args={"type": "web"}),
        )
        res = _run(module.upload_document(dataset_id="missing"))
        assert res["code"] == 102
        assert "Can't find the dataset" in res["message"]

    def test_no_permission(self, document_rest_api_module, monkeypatch):
        module = document_rest_api_module
        kb = SimpleNamespace(id="kb1", tenant_id="tenant1", name="kb", parser_id="parser", pipeline_id="pipe", parser_config={})
        monkeypatch.setattr(module, "is_valid_url", lambda _url: True)
        monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, kb))
        monkeypatch.setattr(module, "check_kb_team_permission", lambda *_args, **_kwargs: False)
        monkeypatch.setattr(
            module,
            "request",
            _DummyRequest(form={"name": "doc", "url": "http://example.com"}, args={"type": "web"}),
        )
        res = _run(module.upload_document(dataset_id="kb1"))
        assert res["code"] == 109
        assert res["message"] == "No authorization."

    def test_download_failure(self, document_rest_api_module, monkeypatch):
        module = document_rest_api_module
        kb = SimpleNamespace(id="kb1", tenant_id="tenant1", name="kb", parser_id="parser", pipeline_id="pipe", parser_config={})
        monkeypatch.setattr(module, "is_valid_url", lambda _url: True)
        monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, kb))
        monkeypatch.setattr(module, "check_kb_team_permission", lambda *_args, **_kwargs: True)
        monkeypatch.setattr(module, "html2pdf", lambda _url: None)
        monkeypatch.setattr(
            module,
            "request",
            _DummyRequest(form={"name": "doc", "url": "http://example.com"}, args={"type": "web"}),
        )
        res = _run(module.upload_document(dataset_id="kb1"))
        assert res["code"] == 100
        assert "Download failure" in res["message"]

    def test_unsupported_type(self, document_rest_api_module, monkeypatch):
        module = document_rest_api_module
        kb = SimpleNamespace(id="kb1", tenant_id="tenant1", name="kb", parser_id="parser", pipeline_id="pipe", parser_config={})
        monkeypatch.setattr(module, "is_valid_url", lambda _url: True)
        monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, kb))
        monkeypatch.setattr(module, "check_kb_team_permission", lambda *_args, **_kwargs: True)
        monkeypatch.setattr(module, "html2pdf", lambda _url: b"%PDF-1.4")
        monkeypatch.setattr(module.FileService, "get_root_folder", lambda _uid: {"id": "root"})
        monkeypatch.setattr(module.FileService, "init_knowledgebase_docs", lambda *_args, **_kwargs: None)
        monkeypatch.setattr(module.FileService, "get_kb_folder", lambda *_args, **_kwargs: {"id": "kb_root"})
        monkeypatch.setattr(module.FileService, "new_a_file_from_kb", lambda *_args, **_kwargs: {"id": "kb_folder"})
        monkeypatch.setattr(module, "duplicate_name", lambda *_args, **_kwargs: "bad.exe")
        monkeypatch.setattr(
            module,
            "request",
            _DummyRequest(form={"name": "doc", "url": "http://example.com"}, args={"type": "web"}),
        )
        res = _run(module.upload_document(dataset_id="kb1"))
        assert res["code"] == 100
        assert "supported yet" in res["message"]

    @pytest.mark.parametrize(
        "filename,filetype,expected_parser",
        [
            ("image.png", "visual", "picture"),
            ("sound.mp3", "aural", "audio"),
            ("deck.pptx", "doc", "presentation"),
            ("mail.eml", "doc", "email"),
        ],
    )
    def test_success_parser_overrides(self, document_rest_api_module, monkeypatch, filename, filetype, expected_parser):
        module = document_rest_api_module
        kb = SimpleNamespace(id="kb1", tenant_id="tenant1", name="kb", parser_id="parser", pipeline_id="pipe", parser_config={})
        captured = {}

        class _Storage:
            def obj_exist(self, *_args, **_kwargs):
                return False

            def put(self, *_args, **_kwargs):
                captured["put"] = True

        def insert_doc(doc):
            captured["doc"] = doc

        monkeypatch.setattr(module, "is_valid_url", lambda _url: True)
        monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, kb))
        monkeypatch.setattr(module, "check_kb_team_permission", lambda *_args, **_kwargs: True)
        monkeypatch.setattr(module, "html2pdf", lambda _url: b"%PDF-1.4")
        monkeypatch.setattr(module.FileService, "get_root_folder", lambda _uid: {"id": "root"})
        monkeypatch.setattr(module.FileService, "init_knowledgebase_docs", lambda *_args, **_kwargs: None)
        monkeypatch.setattr(module.FileService, "get_kb_folder", lambda *_args, **_kwargs: {"id": "kb_root"})
        monkeypatch.setattr(module.FileService, "new_a_file_from_kb", lambda *_args, **_kwargs: {"id": "kb_folder"})
        monkeypatch.setattr(module, "duplicate_name", lambda *_args, **_kwargs: filename)
        monkeypatch.setattr(module, "filename_type", lambda _name: filetype)
        monkeypatch.setattr(module, "thumbnail", lambda *_args, **_kwargs: "")
        monkeypatch.setattr(module, "get_uuid", lambda: "doc-1")
        monkeypatch.setattr(module.settings, "STORAGE_IMPL", _Storage())
        monkeypatch.setattr(module.DocumentService, "insert", insert_doc)
        monkeypatch.setattr(module.FileService, "add_file_from_kb", lambda *_args, **_kwargs: None)
        monkeypatch.setattr(
            module,
            "request",
            _DummyRequest(form={"name": "doc", "url": "http://example.com"}, args={"type": "web"}),
        )

        res = _run(module.upload_document(dataset_id="kb1"))
        assert res["code"] == 0
        assert captured["doc"]["parser_id"] == expected_parser
        assert captured["put"] is True

    def test_exception_path(self, document_rest_api_module, monkeypatch):
        module = document_rest_api_module
        kb = SimpleNamespace(id="kb1", tenant_id="tenant1", name="kb", parser_id="parser", pipeline_id="pipe", parser_config={})

        class _Storage:
            def obj_exist(self, *_args, **_kwargs):
                return False

            def put(self, *_args, **_kwargs):
                return None

        def insert_doc(_doc):
            raise RuntimeError("boom")

        monkeypatch.setattr(module, "is_valid_url", lambda _url: True)
        monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, kb))
        monkeypatch.setattr(module, "check_kb_team_permission", lambda *_args, **_kwargs: True)
        monkeypatch.setattr(module, "html2pdf", lambda _url: b"%PDF-1.4")
        monkeypatch.setattr(module.FileService, "get_root_folder", lambda _uid: {"id": "root"})
        monkeypatch.setattr(module.FileService, "init_knowledgebase_docs", lambda *_args, **_kwargs: None)
        monkeypatch.setattr(module.FileService, "get_kb_folder", lambda *_args, **_kwargs: {"id": "kb_root"})
        monkeypatch.setattr(module.FileService, "new_a_file_from_kb", lambda *_args, **_kwargs: {"id": "kb_folder"})
        monkeypatch.setattr(module, "duplicate_name", lambda *_args, **_kwargs: "doc.pdf")
        monkeypatch.setattr(module, "filename_type", lambda _name: "pdf")
        monkeypatch.setattr(module, "thumbnail", lambda *_args, **_kwargs: "")
        monkeypatch.setattr(module, "get_uuid", lambda: "doc-1")
        monkeypatch.setattr(module.settings, "STORAGE_IMPL", _Storage())
        monkeypatch.setattr(module.DocumentService, "insert", insert_doc)
        monkeypatch.setattr(module.FileService, "add_file_from_kb", lambda *_args, **_kwargs: None)
        monkeypatch.setattr(
            module,
            "request",
            _DummyRequest(form={"name": "doc", "url": "http://example.com"}, args={"type": "web"}),
        )

        res = _run(module.upload_document(dataset_id="kb1"))
        assert res["code"] == 100
