import json
import unittest

from analogai.deepthink.application.value_objects.user_agent import UserAgent
from analogai.deepthink.infrastructure.gateways.llm.gateway import LlmGateway
from analogai.deepthink.infrastructure.gateways.llm.enums.llm_provider import LlmProvider
from analogai.deepthink.infrastructure.gateways.llm.llm_factory import LlmGatewayFactory
from analogai.deepthink.infrastructure.gateways.llm.llm_gateway_provider_factory import (
    LlmGatewayProviderFactory,
)
from analogai.deepthink.presentation.preprocessing.previous_messages_stack import PreviousMessagesStack
from analogai.deepthink.tests.functional.storage_cleaner import cleanup_storages
from analogai.deepthink.startup.bootstrap import create_deepthink_engine


class TestThinkerPreviousMessageStackIntegration(unittest.TestCase):
    def setUp(self):
        cleanup_storages()
        class _TestGateway(LlmGateway):
            def __init__(self):
                self._prompt = ""
                self._model = None

            def set_system_prompt(self, prompt: str) -> None:
                self._prompt = prompt

            def set_model(self, model_id: str) -> None:
                self._model = model_id

            def generate_completion(self, user_message: str) -> str:
                sentence_text = user_message.strip()
                last_line = sentence_text.splitlines()[-1] if sentence_text else ""
                if last_line.startswith("Sentence:"):
                    sentence_text = last_line.replace("Sentence:", "", 1).strip()
                history_subject = _extract_last_subject(user_message)
                if history_subject:
                    sentence_text = _resolve_pronouns(sentence_text, history_subject)
                parsed = {
                    "intent_type": "making_statement",
                    "subject": sentence_text,
                    "predicate": {"relation": "is", "complement": "unknown"},
                }
                return json.dumps({"sentences": [parsed]})

            async def generate_streaming_completion(self, user_message: str):
                yield self.generate_completion(user_message)

        self._prev_provider_create_gateway = LlmGatewayProviderFactory.create_gateway

        def _create_gateway_for_test(
            cls,
            gateway_type=None,
            system_prompt: str | None = None,
            max_tokens=None,
            model: str | None = None,
        ):
            gateway = _TestGateway()
            if model:
                gateway.set_model(model)
            if system_prompt is not None:
                gateway.set_system_prompt(system_prompt)
            return gateway

        LlmGatewayProviderFactory.create_gateway = classmethod(_create_gateway_for_test)

        self._prev_gateway_type = LlmGatewayFactory._default_gateway_type
        self._prev_gateway_instance = LlmGatewayFactory._default_gateway_instance
        LlmGatewayFactory.set_default_gateway_type(LlmProvider.OPENAI)
        LlmGatewayFactory.set_default_gateway_instance(_TestGateway())

        self._deepthink = create_deepthink_engine(
            llm_gateway_instance=_TestGateway(),
            llm_provider=LlmProvider.OPENAI,
            llm_model="test-model",
        )

    def tearDown(self):
        LlmGatewayProviderFactory.create_gateway = self._prev_provider_create_gateway
        LlmGatewayFactory.set_default_gateway_type(self._prev_gateway_type)
        LlmGatewayFactory.set_default_gateway_instance(self._prev_gateway_instance)
        cleanup_storages()

    def test_conversation_saved_per_user_after_think(self):
        user_id = "user-a"
        agent_id = "agent-1"
        stack = PreviousMessagesStack()
        user_agent = UserAgent(user_id=user_id, agent_id=agent_id, user_authority=1.0)
        self._deepthink.deepthink("socrates is human", user_agent=user_agent, previous_messages_stack=stack)
        history = list(stack.get_history())
        self.assertEqual(len(history), 2)
        self.assertIn("socrates", history[0].lower())
        self.assertIn("human", history[0].lower())

    def test_conversation_adapted_to_stack_and_used_for_pronoun_resolution(self):
        user_id = "user-b"
        agent_id = "agent-1"
        stack = PreviousMessagesStack()
        user_agent = UserAgent(user_id=user_id, agent_id=agent_id, user_authority=1.0)
        self._deepthink.deepthink("john is mortal", user_agent=user_agent, previous_messages_stack=stack)
        result = self._deepthink.deepthink("he is philosopher", user_agent=user_agent, previous_messages_stack=stack)
        self.assertIsNotNone(result)
        self.assertIn("john", result.lower())
        self.assertNotIn(" he ", f" {result.lower()} ")
        history = list(stack.get_history())
        self.assertEqual(len(history), 4)

    def test_conversations_isolated_per_user(self):
        stack_user1 = PreviousMessagesStack()
        stack_user2 = PreviousMessagesStack()
        user_agent_1 = UserAgent(user_id="user-1", agent_id="agent-1", user_authority=1.0)
        user_agent_2 = UserAgent(user_id="user-2", agent_id="agent-1", user_authority=1.0)
        self._deepthink.deepthink("alice likes apples", user_agent=user_agent_1, previous_messages_stack=stack_user1)
        self._deepthink.deepthink("bob likes oranges", user_agent=user_agent_2, previous_messages_stack=stack_user2)
        history_user1 = list(stack_user1.get_history())
        history_user2 = list(stack_user2.get_history())
        self.assertEqual(history_user1[0], "alice likes apples")
        self.assertEqual(history_user2[0], "bob likes oranges")

    def test_conversation_adapted_to_stack_returns_persisted_history(self):
        user_id = "user-c"
        agent_id = "agent-1"
        stack = PreviousMessagesStack()
        user_agent = UserAgent(user_id=user_id, agent_id=agent_id, user_authority=1.0)
        self._deepthink.deepthink("maria reads books", user_agent=user_agent, previous_messages_stack=stack)
        history = list(stack.get_history())
        self.assertGreater(len(history), 0)
        self.assertIn("maria reads books", history)


def _extract_last_subject(message: str) -> str | None:
    for line in reversed(message.splitlines()):
        line = line.strip()
        if line.lower().startswith("q:"):
            content = line[2:].strip()
            for sep in (" is ", " are "):
                if sep in content:
                    return content.split(sep)[0].strip()
            return content.split()[0] if content else None
    return None


def _resolve_pronouns(sentence: str, subject: str) -> str:
    replacements = {
        " he ": f" {subject} ",
        " she ": f" {subject} ",
        " they ": f" {subject} ",
        " him ": f" {subject} ",
        " her ": f" {subject} ",
    }
    padded = f" {sentence.strip()} "
    lowered = padded.lower()
    for token, replacement in replacements.items():
        if token in lowered:
            lowered = lowered.replace(token, replacement)
    return lowered.strip()


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