import json
import unittest
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.presentation.preprocessing.previous_messages_stack import PreviousMessagesStack
from analogai.deepthink.presentation.preprocessing.preprocessor import preprocess_text


class TestPreprocessorIntegration(unittest.TestCase):
    def setUp(self):
        class _TestGateway(LlmGateway):
            def __init__(self):
                self._prompt = ""

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

            def set_model(self, model_id: str) -> None:
                return

            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_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())
    
    def tearDown(self):
        LlmGatewayFactory.set_default_gateway_type(self._prev_gateway_type)
        LlmGatewayFactory.set_default_gateway_instance(self._prev_gateway_instance)


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()

    def test_preprocessor_chain(self):
        input_text = "I am your creator"
        preprocessed = preprocess_text(input_text, PreviousMessagesStack())
        self.assertIsNotNone(preprocessed)
        self.assertIsInstance(preprocessed, str)

    def test_preprocessor_with_pronouns(self):
        input_text = "You are my friend"
        preprocessed = preprocess_text(input_text, PreviousMessagesStack())
        self.assertIsNotNone(preprocessed)
        self.assertIsInstance(preprocessed, str)

    def test_preprocessor_with_causal(self):
        input_text = "If it rains, I will stay home"
        preprocessed = preprocess_text(input_text, PreviousMessagesStack())
        self.assertIsNotNone(preprocessed)
        self.assertIsInstance(preprocessed, str)

    def test_pronoun_resolution_in_conversation(self):
        store = PreviousMessagesStack()
        first_message = "john is mortal"
        second_message = "he is philosopher"
        
        preprocess_text(first_message, store)
        result = preprocess_text(second_message, store)
        
        self.assertIn("john", result.lower())
        result_words = result.lower().split()
        self.assertNotIn("he", result_words)

    def test_previous_messages_stack_persists_across_calls(self):
        stack = PreviousMessagesStack()
        first_message = "socrates is human"
        second_message = "socrates sells iphone 12"
        
        preprocess_text(first_message, stack)
        history_after_first = list(stack.get_history())
        self.assertEqual(len(history_after_first), 1)
        
        preprocess_text(second_message, stack)
        history_after_second = list(stack.get_history())
        self.assertEqual(len(history_after_second), 2)


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

