import json
import unittest
from datetime import datetime, timezone

import httpx
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from app.clients.shared_ai_platform import RemoteSharedAIPlatformClient, build_shared_ai_platform_client
from app.core.config import settings
from app.db import Base
from app.models.conversation import Conversation
from app.models.message import Message
from shared_ai_platform_sdk import SharedAIPlatformError
from shared_ai_platform_sdk.client import SharedAIPlatformClient as SharedAiPlatformSdkClient
from shared_ai_platform_sdk.types import RequestContext, SSEEvent


class SharedAiPlatformSdkTestCase(unittest.TestCase):
    def test_stream_conversation_uses_default_v0_2_api_prefix(self):
        seen_paths = []

        def handler(request: httpx.Request) -> httpx.Response:
            seen_paths.append(request.url.path)
            return httpx.Response(
                200,
                headers={"content-type": "text/event-stream"},
                content=b'data: {"type":"end","conversation_id":"conv-1"}\n\n',
            )

        transport = httpx.MockTransport(handler)
        http_client = httpx.Client(base_url="http://testserver", transport=transport, timeout=30.0)
        sdk = SharedAiPlatformSdkClient(
            base_url="http://testserver",
            api_key="test-key",
            app_id="test-app",
            client=http_client,
        )

        try:
            events = list(
                sdk.stream_conversation(
                    RequestContext(tenant_id="tenant-1", user_id="user-1", user_role="admin"),
                    message="hello",
                )
            )
        finally:
            http_client.close()

        self.assertEqual(seen_paths, ["/api/v0.2/conversations/stream"])
        self.assertEqual(events[0].payload["conversation_id"], "conv-1")

    def test_stream_conversation_uses_configured_api_prefix(self):
        seen_paths = []

        def handler(request: httpx.Request) -> httpx.Response:
            seen_paths.append(request.url.path)
            return httpx.Response(
                200,
                headers={"content-type": "text/event-stream"},
                content=b'data: {"type":"end","conversation_id":"conv-1"}\n\n',
            )

        transport = httpx.MockTransport(handler)
        http_client = httpx.Client(base_url="http://testserver", transport=transport, timeout=30.0)
        sdk = SharedAiPlatformSdkClient(
            base_url="http://testserver",
            api_key="test-key",
            app_id="test-app",
            api_prefix="/api/v0.1",
            client=http_client,
        )

        try:
            events = list(
                sdk.stream_conversation(
                    RequestContext(tenant_id="tenant-1", user_id="user-1", user_role="admin"),
                    message="hello",
                )
            )
        finally:
            http_client.close()

        self.assertEqual(seen_paths, ["/api/v0.1/conversations/stream"])
        self.assertEqual(events[0].payload["conversation_id"], "conv-1")

    def test_stream_conversation_supports_split_data_and_json_lines(self):
        sse_payload = (
            'data:\n'
            '{"type":"delta","content":"hello"}\n\n'
            'data:\n'
            '{"type":"end","conversation_id":"conv-1"}\n\n'
        ).encode("utf-8")

        def handler(request: httpx.Request) -> httpx.Response:
            return httpx.Response(
                200,
                headers={"content-type": "text/event-stream"},
                content=sse_payload,
            )

        transport = httpx.MockTransport(handler)
        http_client = httpx.Client(base_url="http://testserver", transport=transport, timeout=30.0)
        sdk = SharedAiPlatformSdkClient(
            base_url="http://testserver",
            api_key="test-key",
            app_id="test-app",
            client=http_client,
        )

        try:
            events = list(
                sdk.stream_conversation(
                    RequestContext(tenant_id="tenant-1", user_id="user-1", user_role="admin"),
                    message="hello",
                )
            )
        finally:
            http_client.close()

        self.assertEqual([event.event_type for event in events], ["delta", "end"])
        self.assertEqual(events[0].payload["content"], "hello")
        self.assertEqual(events[1].payload["conversation_id"], "conv-1")


class RemoteSharedAiPlatformClientStreamTestCase(unittest.TestCase):
    def test_missing_shared_ai_config_uses_local_history_fallback(self):
        original_base_url = settings.shared_ai_platform_base_url
        original_api_key = settings.shared_ai_platform_api_key
        settings.shared_ai_platform_base_url = ""
        settings.shared_ai_platform_api_key = ""
        engine = create_engine("sqlite:///:memory:", future=True)
        Base.metadata.create_all(engine)
        session = sessionmaker(bind=engine, future=True)()
        try:
            now = datetime.now(timezone.utc)
            session.add(
                Conversation(
                    id="conv-local-config-missing",
                    tenant_id="tenant-1",
                    user_id="user-1",
                    title="Local fallback",
                    message_count=1,
                    last_message_at=now,
                    created_at=now,
                    updated_at=now,
                )
            )
            session.commit()

            client = build_shared_ai_platform_client(session)
            payload = client.list_conversations({"id": "user-1", "tenant_id": "tenant-1", "role": "admin"}, 1, 20)
            chunks = list(client.stream_answer({"id": "user-1", "tenant_id": "tenant-1", "role": "admin"}, None, "hello"))

            self.assertEqual(payload["source"], "local_fallback")
            self.assertEqual(payload["items"][0]["id"], "conv-local-config-missing")
            self.assertIn("SHARED_AI_PLATFORM_BASE_URL", payload["warning"])
            self.assertIn("SHARED_AI_PLATFORM_API_KEY", payload["warning"])
            self.assertIn("模型服务未配置", chunks[0])
        finally:
            session.close()
            engine.dispose()
            settings.shared_ai_platform_base_url = original_base_url
            settings.shared_ai_platform_api_key = original_api_key

    def test_remote_stream_ignores_start_and_usage_and_maps_delta(self):
        calls = []

        class StubSdk:
            def stream_conversation(self, context, *, message, chat_context=None, provider=None, template_key=None, conversation_id=None):
                calls.append(
                    {
                        "message": message,
                        "chat_context": chat_context,
                        "template_key": template_key,
                        "conversation_id": conversation_id,
                    }
                )
                yield SSEEvent(event_type="start", payload={"type": "start", "conversation_id": "conv-1"})
                yield SSEEvent(event_type="delta", payload={"type": "delta", "content": "hello"})
                yield SSEEvent(event_type="usage", payload={"type": "usage", "usage": {"provider_mode": "stream"}})
                yield SSEEvent(event_type="sources", payload={"type": "sources", "sources": [{"title": "source", "content": "body"}]})
                yield SSEEvent(event_type="end", payload={"type": "end", "conversation_id": "conv-1"})

            def run_legal_review(self, context, *, title, text, contract_type=None, review_as="party_a", provider=None):
                return {
                    "basic_info": {
                        "contract_type": contract_type or "采购合同",
                        "contract_title": title,
                        "party_a": "甲方",
                        "party_b": "乙方",
                        "amount": "100万元",
                        "term": "自2026-01-01至2026-12-31",
                    },
                    "risk_summary": {"high": 1, "medium": 0, "low": 0, "overall": "high"},
                    "clauses": [{"title": "付款条款", "content": "签署后5日内支付50%", "clause_type": "payment", "clause_label": "付款条款", "summary": "签署后5日内支付50%"}],
                    "risk_items": [{"title": "存在自动续约安排", "level": "high", "description": "需核查自动续约", "suggestion": "补充退出机制"}],
                    "review_opinion": "建议人工复核",
                    "recommendation": "manual_review",
                    "confidence_summary": {"provider_mode": "provider-enhanced"},
                    "contract_structure": {"title": title, "payment_milestones": [], "clause_relationships": [], "missing_sections": [], "clause_type_counts": {"payment": 1}},
                }

        client = RemoteSharedAIPlatformClient(
            base_url="http://testserver",
            app_id="test-app",
            api_key="test-key",
            timeout_seconds=30.0,
        )
        client.sdk = StubSdk()
        current_user = {"id": "user-1", "tenant_id": "tenant-1", "role": "admin"}

        chunks = list(client.stream_answer(current_user, None, "hello", {"metadata": {"reviewTitle": "demo"}}, template_key="legal_chat_default"))

        self.assertEqual(
            chunks,
            [
                'data: {"type": "text", "content": "hello"}\n\n',
                'data: {"type": "sources", "sources": [{"title": "source", "content": "body"}]}\n\n',
                'data: {"type": "end", "conversation_id": "conv-1"}\n\n',
            ],
        )
        self.assertEqual(calls[0]["template_key"], None)
        self.assertEqual(calls[0]["chat_context"]["metadata"]["reviewTitle"], "demo")

    def test_remote_stream_retries_without_template_key_when_remote_chat_template_fails_before_payload(self):
        calls = []

        class StubSdk:
            def stream_conversation(self, context, *, message, chat_context=None, provider=None, template_key=None, conversation_id=None):
                calls.append(
                    {
                        "message": message,
                        "chat_context": chat_context,
                        "template_key": template_key,
                        "conversation_id": conversation_id,
                    }
                )
                if template_key == "general-assistant.chat":
                    raise RuntimeError("peer closed connection without sending complete message body (incomplete chunked read)")
                yield SSEEvent(event_type="delta", payload={"type": "delta", "content": "fallback hello"})
                yield SSEEvent(event_type="end", payload={"type": "end", "conversation_id": "conv-fallback"})

        client = RemoteSharedAIPlatformClient(
            base_url="http://testserver",
            app_id="test-app",
            api_key="test-key",
            timeout_seconds=30.0,
        )
        client.sdk = StubSdk()
        current_user = {"id": "user-1", "tenant_id": "tenant-1", "role": "admin"}

        chunks = list(client.stream_answer(current_user, None, "hello", {"metadata": {"reviewTitle": "demo"}}, template_key="general-assistant.chat"))

        self.assertEqual(
            chunks,
            [
                'data: {"type": "text", "content": "fallback hello"}\n\n',
                'data: {"type": "end", "conversation_id": "conv-fallback"}\n\n',
            ],
        )
        self.assertEqual([call["template_key"] for call in calls], ["general-assistant.chat", None])

    def test_remote_stream_skips_non_remote_template_keys(self):
        calls = []

        class StubSdk:
            def stream_conversation(self, context, *, message, chat_context=None, provider=None, template_key=None, conversation_id=None):
                calls.append(template_key)
                yield SSEEvent(event_type="delta", payload={"type": "delta", "content": "hello"})
                yield SSEEvent(event_type="end", payload={"type": "end", "conversation_id": "conv-1"})

        client = RemoteSharedAIPlatformClient(
            base_url="http://testserver",
            app_id="test-app",
            api_key="test-key",
            timeout_seconds=30.0,
        )
        client.sdk = StubSdk()
        current_user = {"id": "user-1", "tenant_id": "tenant-1", "role": "admin"}

        chunks = list(client.stream_answer(current_user, None, "hello", None, template_key="negotiation_strategy_default"))

        self.assertEqual(
            chunks,
            [
                'data: {"type": "text", "content": "hello"}\n\n',
                'data: {"type": "end", "conversation_id": "conv-1"}\n\n',
            ],
        )
        self.assertEqual(calls, [None])

    def test_remote_document_payloads_are_normalized_for_source_metadata(self):
        class StubSdk:
            def list_documents(self, context, knowledge_base_id, *, page=1, page_size=20):
                return {
                    "items": [
                        {
                            "id": "doc-1",
                            "knowledge_base_id": knowledge_base_id,
                            "file_name": "civil-code.md",
                            "mime_type": "text/markdown",
                            "file_size": 128,
                            "status": "indexed",
                            "metadata": {
                                "source_type": "statute",
                                "issuer": "全国人民代表大会",
                                "effective_date": "2021-01-01",
                                "jurisdiction": "CN",
                                "tags": ["民法典", "合同"],
                            },
                        }
                    ],
                    "page": page,
                    "page_size": page_size,
                    "total": 1,
                }

            def upload_document(self, context, knowledge_base_id, *, file_name, mime_type, content):
                return {
                    "id": "doc-2",
                    "knowledge_base_id": knowledge_base_id,
                    "file_name": file_name,
                    "mime_type": mime_type,
                    "file_size": len(content),
                    "status": "indexed",
                    "tags_json": ["制度"],
                    "metadata_json": {
                        "source_type": "policy",
                        "authority": "企业法务部",
                        "effective_date": "2026-05-01",
                        "region": "CN-SH",
                    },
                }

            def retry_document(self, context, knowledge_base_id, document_id, *, file_name=None, mime_type=None, content=None):
                return {
                    "id": document_id,
                    "knowledge_base_id": knowledge_base_id,
                    "file_name": file_name or "retry.txt",
                    "mime_type": mime_type or "text/plain",
                    "file_size": len(content or b""),
                    "status": "indexed",
                    "metadata": {
                        "source_type": "judicial-interpretation",
                        "authority": "最高人民法院",
                        "effective_date": "2024-02-01",
                        "region": "CN",
                    },
                }

        client = RemoteSharedAIPlatformClient(
            base_url="http://testserver",
            app_id="test-app",
            api_key="test-key",
            timeout_seconds=30.0,
        )
        client.sdk = StubSdk()
        current_user = {"id": "user-1", "tenant_id": "tenant-1", "role": "admin"}

        listed = client.list_documents(current_user, "kb-1", 1, 20)
        uploaded = client.upload_document(current_user, "kb-1", "policy.md", "text/markdown", b"hello")
        retried = client.retry_document(current_user, "kb-1", "doc-3", file_name="retry.md", mime_type="text/markdown", content=b"world")

        self.assertEqual(listed["items"][0]["source_type"], "statute")
        self.assertEqual(listed["items"][0]["authority"], "全国人民代表大会")
        self.assertEqual(listed["items"][0]["effective_date"], "2021-01-01")
        self.assertEqual(listed["items"][0]["region"], "CN")
        self.assertEqual(listed["items"][0]["tags"], ["民法典", "合同"])
        self.assertEqual(uploaded["authority"], "企业法务部")
        self.assertEqual(uploaded["region"], "CN-SH")
        self.assertEqual(uploaded["tags"], ["制度"])
        self.assertEqual(retried["source_type"], "judicial-interpretation")
        self.assertEqual(retried["authority"], "最高人民法院")

    def test_remote_knowledge_bases_are_sorted_latest_first(self):
        class StubSdk:
            def list_knowledge_bases(self, context, *, page=1, page_size=20):
                return {
                    "items": [
                        {
                            "id": "kb-old",
                            "name": "old",
                            "document_count": 1,
                            "created_at": "2026-05-12T00:00:00Z",
                            "updated_at": "2026-05-12T00:00:00Z",
                        },
                        {
                            "id": "kb-new",
                            "name": "new",
                            "document_count": 2,
                            "created_at": "2026-05-13T00:00:00Z",
                            "updated_at": "2026-05-13T00:00:00Z",
                        },
                    ],
                    "page": page,
                    "page_size": page_size,
                    "total": 2,
                }

        client = RemoteSharedAIPlatformClient(
            base_url="http://testserver",
            app_id="test-app",
            api_key="test-key",
            timeout_seconds=30.0,
        )
        client.sdk = StubSdk()

        payload = client.list_knowledge_bases({"id": "user-1", "tenant_id": "tenant-1", "role": "admin"}, 1, 100)

        self.assertEqual([item["id"] for item in payload["items"]], ["kb-new", "kb-old"])
        self.assertEqual(payload["items"][0]["retrieval_mode"], "hybrid")

    def test_remote_knowledge_bases_support_sdk_without_pagination_arguments(self):
        class StubSdk:
            def list_knowledge_bases(self, context):
                return {
                    "items": [
                        {"id": "kb-1", "name": "one", "created_at": "2026-05-10T00:00:00Z"},
                        {"id": "kb-2", "name": "two", "created_at": "2026-05-11T00:00:00Z"},
                        {"id": "kb-3", "name": "three", "created_at": "2026-05-12T00:00:00Z"},
                    ]
                }

        client = RemoteSharedAIPlatformClient(
            base_url="http://testserver",
            app_id="test-app",
            api_key="test-key",
            timeout_seconds=30.0,
        )
        client.sdk = StubSdk()

        payload = client.list_knowledge_bases({"id": "user-1", "tenant_id": "tenant-1", "role": "admin"}, 2, 2)

        self.assertEqual(payload["page"], 2)
        self.assertEqual(payload["page_size"], 2)
        self.assertEqual(payload["total"], 3)
        self.assertEqual([item["id"] for item in payload["items"]], ["kb-1"])
        self.assertEqual(payload["items"][0]["retrieval_mode"], "hybrid")

    def test_remote_knowledge_base_publish_chain_calls_sdk(self):
        class StubSdk:
            def list_knowledge_base_versions(self, context, knowledge_base_id):
                return [{"id": "kbver-1", "knowledge_base_id": knowledge_base_id, "version_label": "v1"}]

            def publish_knowledge_base(self, context, knowledge_base_id):
                return {"id": knowledge_base_id, "status": "active", "current_version_label": "v1"}

            def unpublish_knowledge_base(self, context, knowledge_base_id):
                return {"id": knowledge_base_id, "status": "archived", "current_version_label": "v1"}

            def rollback_knowledge_base(self, context, knowledge_base_id, *, target_version_label, change_summary=None):
                return {
                    "id": knowledge_base_id,
                    "status": "active",
                    "current_version_label": target_version_label,
                    "change_summary": change_summary,
                }

        client = RemoteSharedAIPlatformClient(
            base_url="http://testserver",
            app_id="test-app",
            api_key="test-key",
            timeout_seconds=30.0,
        )
        client.sdk = StubSdk()
        current_user = {"id": "user-1", "tenant_id": "tenant-1", "role": "admin"}

        versions = client.list_knowledge_base_versions(current_user, "kb-1")
        published = client.publish_knowledge_base(current_user, "kb-1")
        unpublished = client.unpublish_knowledge_base(current_user, "kb-1")
        rolled_back = client.rollback_knowledge_base(current_user, "kb-1", "v1", "rollback")

        self.assertEqual(versions[0]["version_label"], "v1")
        self.assertEqual(published["current_version_label"], "v1")
        self.assertEqual(unpublished["status"], "archived")
        self.assertEqual(rolled_back["current_version_label"], "v1")

    def test_remote_knowledge_base_versions_falls_back_to_current_snapshot_when_sdk_fails(self):
        class StubSdk:
            def list_knowledge_base_versions(self, context, knowledge_base_id):
                raise SharedAIPlatformError(500, "version endpoint unavailable")

            def list_knowledge_bases(self, context, page=1, page_size=100):
                return {
                    "items": [
                        {
                            "id": "kb-1",
                            "name": "Pilot KB",
                            "description": "demo",
                            "status": "active",
                            "retrieval_mode": "hybrid",
                            "current_version_label": "v2",
                            "published_at": "2026-06-20T00:00:00+00:00",
                            "published_by": "user-admin",
                            "created_at": "2026-06-19T00:00:00+00:00",
                            "updated_at": "2026-06-20T00:00:00+00:00",
                        }
                    ],
                    "page": page,
                    "page_size": page_size,
                    "total": 1,
                }

        client = RemoteSharedAIPlatformClient(
            base_url="http://testserver",
            app_id="test-app",
            api_key="test-key",
            timeout_seconds=30.0,
        )
        client.sdk = StubSdk()

        versions = client.list_knowledge_base_versions({"id": "user-1", "tenant_id": "tenant-1", "role": "admin"}, "kb-1")

        self.assertEqual(len(versions), 1)
        self.assertEqual(versions[0]["version_label"], "v2")
        self.assertIn("共享 AI 版本历史暂时不可用", versions[0]["change_summary"])

    def test_remote_conversation_history_falls_back_to_local_records_when_sdk_fails(self):
        engine = create_engine("sqlite:///:memory:", future=True)
        Base.metadata.create_all(engine)
        session = sessionmaker(bind=engine, future=True)()
        try:
            now = datetime.now(timezone.utc)
            session.add(
                Conversation(
                    id="conv-local-1",
                    tenant_id="tenant-1",
                    user_id="user-1",
                    title="Local history",
                    message_count=2,
                    last_message_at=now,
                    created_at=now,
                    updated_at=now,
                )
            )
            session.add_all(
                [
                    Message(
                        id="msg-local-1",
                        tenant_id="tenant-1",
                        conversation_id="conv-local-1",
                        user_id="user-1",
                        role="user",
                        content="hello",
                        created_at=now,
                    ),
                    Message(
                        id="msg-local-2",
                        tenant_id="tenant-1",
                        conversation_id="conv-local-1",
                        user_id="user-1",
                        role="assistant",
                        content="world",
                        sources_json=[{"title": "source", "content": "body"}],
                        created_at=now,
                    ),
                ]
            )
            session.commit()

            class StubSdk:
                def list_conversations(self, context, *, page=1, page_size=20):
                    raise RuntimeError("connect timeout")

                def get_conversation_detail(self, context, conversation_id):
                    raise RuntimeError("connect timeout")

            client = RemoteSharedAIPlatformClient(
                base_url="http://testserver",
                app_id="test-app",
                api_key="test-key",
                timeout_seconds=30.0,
                db=session,
            )
            client.sdk = StubSdk()
            current_user = {"id": "user-1", "tenant_id": "tenant-1", "role": "admin"}

            payload = client.list_conversations(current_user, 1, 20)
            detail = client.get_conversation_detail(current_user, "conv-local-1")

            self.assertEqual(payload["source"], "local_fallback")
            self.assertEqual(payload["total"], 1)
            self.assertEqual(payload["items"][0]["id"], "conv-local-1")
            self.assertIn("connect timeout", payload["warning"])
            self.assertEqual(detail["source"], "local_fallback")
            self.assertEqual(len(detail["messages"]), 2)
            self.assertEqual(detail["messages"][1]["sources"][0]["title"], "source")
        finally:
            session.close()
            engine.dispose()

    def test_enhance_contract_generation_extracts_html_from_structured_json(self):
        class StubClient(RemoteSharedAIPlatformClient):
            def stream_answer(self, current_user, conversation_id, message, review_context=None, template_key=None):
                payload = '{"title":"担保合同","html":"第一条 合同正文","text":"备用文本"}'
                yield f'data: {{"type":"text","content":{json.dumps(payload, ensure_ascii=False)} }}\n\n'
                yield 'data: {"type":"end","conversation_id":"conv-1"}\n\n'

        client = StubClient(
            base_url="http://testserver",
            app_id="test-app",
            api_key="test-key",
            timeout_seconds=30.0,
        )
        result = client.enhance_contract_generation(
            {"id": "user-1", "tenant_id": "tenant-1", "role": "admin"},
            prompt_runtime={"prompt_key": "contract_generation_default", "prompt_version": "contract_generation_default.v1.0"},
            prompt="请补强合同正文",
            draft_content="原始草稿",
        )

        self.assertEqual(result["content"], "第一条 合同正文")
        self.assertEqual(result["raw_content"], '{"title":"担保合同","html":"第一条 合同正文","text":"备用文本"}')
        self.assertEqual(result["structured_output"]["title"], "担保合同")


    def test_enhance_contract_generation_strips_html_from_structured_json(self):
        class StubClient(RemoteSharedAIPlatformClient):
            def stream_answer(self, current_user, conversation_id, message, review_context=None, template_key=None):
                payload = '{"title":"Guarantee Contract","html":"<h1>Guarantee Contract</h1><p>Article 1 guarantee scope.</p>"}'
                yield f'data: {{"type":"text","content":{json.dumps(payload, ensure_ascii=False)} }}\n\n'
                yield 'data: {"type":"end","conversation_id":"conv-1"}\n\n'

        client = StubClient(
            base_url="http://testserver",
            app_id="test-app",
            api_key="test-key",
            timeout_seconds=30.0,
        )
        result = client.enhance_contract_generation(
            {"id": "user-1", "tenant_id": "tenant-1", "role": "admin"},
            prompt_runtime={"prompt_key": "contract_generation_default", "prompt_version": "contract_generation_default.v1.0"},
            prompt="enhance contract",
            draft_content="draft",
        )

        self.assertEqual(result["content"], "Guarantee Contract\nArticle 1 guarantee scope.")
        self.assertNotIn("<h1>", result["content"])


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