import unittest
from unittest.mock import patch, mock_open
import hashlib
import csv
from collections import OrderedDict
import src.utils.llm_utils as llm_utils
from src import config

class TestLLMUtils(unittest.TestCase):
    def setUp(self):
        self.filename = "test_file.txt"
        self.prompt = "This is a test prompt."
        self.response = "This is a test response."
        self.model_type = "test_model"
        self.input_cost_per_1k = 0.01
        self.output_cost_per_1k = 0.02
        self.mock_stats = {self.filename: {"requests": 0, "tokens_sent": 0, "tokens_received": 0, "total_time": 0}}
        config.MODEL_STATS = self.mock_stats
        config.GLOBAL_STATS = {}

    def test_update_stats(self):
        llm_utils.update_stats(self.filename, 10, 20, 5, self.model_type)
        self.assertEqual(config.MODEL_STATS[self.filename]["requests"], 1)
        self.assertEqual(config.MODEL_STATS[self.filename]["tokens_sent"], 10)
        self.assertEqual(config.MODEL_STATS[self.filename]["tokens_received"], 20)
        self.assertEqual(config.MODEL_STATS[self.filename]["total_time"], 5)
        self.assertEqual(config.GLOBAL_STATS["total_requests"], 1)
        self.assertEqual(config.GLOBAL_STATS["total_tokens_sent"], 10)
        self.assertEqual(config.GLOBAL_STATS["total_tokens_received"], 20)
        self.assertEqual(config.GLOBAL_STATS["total_time"], 5)

    def test_count_tokens(self):
        self.assertEqual(llm_utils.count_tokens("123456"), 1)
        self.assertEqual(llm_utils.count_tokens("1234567"), 2)
        self.assertEqual(llm_utils.count_tokens(""), 0)

    def test_create_cost_log_entry(self):
        entry = llm_utils.create_cost_log_entry(
            self.filename, "caller_func", self.prompt, self.response, self.input_cost_per_1k, self.output_cost_per_1k
        )
        self.assertEqual(entry["Filename"], self.filename)
        self.assertEqual(entry["Caller Function"], "caller_func")
        self.assertEqual(entry["Input Prompt Length"], len(self.prompt))
        self.assertEqual(entry["Output Response Length"], len(self.response))
        self.assertIn("Total Cost", entry)

    def test_get_cache_key(self):
        model_id = "test_model"
        expected_hash = hashlib.sha256(f"{model_id}:{self.prompt}".encode()).hexdigest()
        self.assertEqual(llm_utils.get_cache_key(self.prompt, model_id), expected_hash)

    @patch("builtins.open", new_callable=mock_open)
    @patch("csv.DictWriter.writerow")
    def test_invoke_claude_local_mode(self, mock_csv_writer, mock_open_file):
        config.RUN_MODE = "local"
        model_id = config.MODEL_ID_CLAUDE2
        
        with patch("src.utils.llm_utils.local_claude_2", return_value=self.response):
            response = llm_utils.invoke_claude(self.prompt, model_id, self.filename)
            self.assertEqual(response, self.response)
            mock_csv_writer.assert_called()

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