import unittest
from datetime import datetime, timezone

from app.models.knowledge_base import KnowledgeBase
from app.models.knowledge_base_document import KnowledgeBaseDocument
from app.services.knowledge_base_service import KnowledgeBaseService


class StubKnowledgeBaseRepository:
    def __init__(self):
        self.items = {}

    def create(self, knowledge_base):
        self.items[knowledge_base.id] = knowledge_base
        return knowledge_base

    def get_by_id(self, knowledge_base_id):
        return self.items.get(knowledge_base_id)

    def list_paginated(self, tenant_id, page, page_size):
        return [item for item in self.items.values() if item.tenant_id == tenant_id]

    def count(self, tenant_id):
        return len([item for item in self.items.values() if item.tenant_id == tenant_id])


class StubKnowledgeBaseDocumentRepository:
    def __init__(self):
        self.items = {}

    def create(self, document):
        if document.created_at is None:
            document.created_at = datetime.now(timezone.utc)
        self.items[document.id] = document
        return document

    def update(self, document):
        if document.created_at is None:
            document.created_at = datetime.now(timezone.utc)
        self.items[document.id] = document
        return document

    def get_by_id(self, document_id, tenant_id=None, knowledge_base_id=None):
        item = self.items.get(document_id)
        if item is None:
            return None
        if tenant_id is not None and item.tenant_id != tenant_id:
            return None
        if knowledge_base_id is not None and item.knowledge_base_id != knowledge_base_id:
            return None
        return item

    def list_paginated(self, tenant_id, knowledge_base_id, page, page_size):
        return [
            item
            for item in self.items.values()
            if item.tenant_id == tenant_id and item.knowledge_base_id == knowledge_base_id
        ]

    def count(self, tenant_id, knowledge_base_id):
        return len(self.list_paginated(tenant_id, knowledge_base_id, 1, 100))

    def count_by_knowledge_base_ids(self, tenant_id, knowledge_base_ids):
        return {knowledge_base_id: self.count(tenant_id, knowledge_base_id) for knowledge_base_id in knowledge_base_ids}


class KnowledgeBaseServiceTestCase(unittest.TestCase):
    def setUp(self):
        self.service = KnowledgeBaseService(db=None)
        self.service.repository = StubKnowledgeBaseRepository()
        self.service.document_repository = StubKnowledgeBaseDocumentRepository()
        self.current_user = {"id": "user-admin", "tenant_id": "tenant-demo-001", "role": "admin"}
        self.knowledge_base = KnowledgeBase(
            id="kb-123",
            tenant_id=self.current_user["tenant_id"],
            name="合同知识库",
            description="测试知识库",
            status="active",
        )
        self.service.repository.create(self.knowledge_base)

    def test_failed_document_retry_without_source_explains_reupload(self):
        failed = self.service.upload_document(
            self.current_user,
            self.knowledge_base.id,
            "empty.txt",
            "text/plain",
            b"",
        )

        retried = self.service.retry_document(self.current_user, self.knowledge_base.id, failed["id"])

        self.assertEqual(retried["id"], failed["id"])
        self.assertEqual(retried["status"], "failed")
        self.assertEqual(retried["parse_status"], "failed")
        self.assertEqual(retried["index_status"], "blocked")
        self.assertTrue(retried["retryable"])
        self.assertIn("重新上传", retried["error_message"])

    def test_failed_document_retry_with_replacement_file_indexes_same_document(self):
        failed = self.service.upload_document(
            self.current_user,
            self.knowledge_base.id,
            "empty.txt",
            "text/plain",
            b"",
        )

        retried = self.service.retry_document(
            self.current_user,
            self.knowledge_base.id,
            failed["id"],
            file_name="rules-v2.txt",
            mime_type="text/plain",
            content="付款、保密和违约责任制度".encode("utf-8"),
        )

        self.assertEqual(retried["id"], failed["id"])
        self.assertEqual(retried["file_name"], "rules-v2.txt")
        self.assertEqual(retried["status"], "indexed")
        self.assertEqual(retried["parse_status"], "completed")
        self.assertEqual(retried["index_status"], "completed")
        self.assertFalse(retried["retryable"])
        self.assertIn("付款", retried["content_excerpt"])


if __name__ == "__main__":
    unittest.main()