#!/usr/bin/env python3
"""
Unit tests for JSON utilities in jsonUtils.py
Tests parseJsonWithModel, extractJsonString, tryParseJson, repairBrokenJson.
"""

import pytest
import json

from modules.shared.jsonUtils import (
    parseJsonWithModel,
    extractJsonString,
    tryParseJson,
    repairBrokenJson
)
from modules.datamodels.datamodelWorkflow import ActionDefinition, AiResponse


class TestExtractJsonString:
    """Test extractJsonString function"""
    
    def test_extractJsonString_plain_json(self):
        """Test extracting plain JSON"""
        text = '{"key": "value"}'
        result = extractJsonString(text)
        assert result == '{"key": "value"}'
    
    def test_extractJsonString_with_code_fences(self):
        """Test extracting JSON from code fences"""
        text = '```json\n{"key": "value"}\n```'
        result = extractJsonString(text)
        assert result == '{"key": "value"}'
    
    def test_extractJsonString_with_extra_text(self):
        """Test extracting JSON with extra text"""
        text = 'Some text before {"key": "value"} some text after'
        result = extractJsonString(text)
        assert result == '{"key": "value"}'


class TestTryParseJson:
    """Test tryParseJson function"""
    
    def test_tryParseJson_valid_json(self):
        """Test parsing valid JSON"""
        obj, error, cleaned = tryParseJson('{"key": "value"}')
        assert error is None
        assert isinstance(obj, dict)
        assert obj["key"] == "value"
    
    def test_tryParseJson_invalid_json(self):
        """Test parsing invalid JSON"""
        obj, error, cleaned = tryParseJson('{"key": "value"')
        assert error is not None
        assert obj is None
    
    def test_tryParseJson_with_code_fences(self):
        """Test parsing JSON with code fences"""
        obj, error, cleaned = tryParseJson('```json\n{"key": "value"}\n```')
        assert error is None
        assert isinstance(obj, dict)
        assert obj["key"] == "value"


class TestParseJsonWithModel:
    """Test parseJsonWithModel function"""
    
    def test_parseJsonWithModel_valid_json(self):
        """Test parsing valid JSON into Pydantic model"""
        jsonStr = '{"action": "ai.process", "actionObjective": "Process documents"}'
        result = parseJsonWithModel(jsonStr, ActionDefinition)
        assert isinstance(result, ActionDefinition)
        assert result.action == "ai.process"
        assert result.actionObjective == "Process documents"
    
    def test_parseJsonWithModel_with_code_fences(self):
        """Test parsing JSON with code fences"""
        jsonStr = '```json\n{"action": "ai.process", "actionObjective": "Process"}\n```'
        result = parseJsonWithModel(jsonStr, ActionDefinition)
        assert isinstance(result, ActionDefinition)
        assert result.action == "ai.process"
    
    def test_parseJsonWithModel_invalid_json_raises(self):
        """Test that invalid JSON raises ValueError"""
        jsonStr = '{"action": "ai.process"'
        with pytest.raises(ValueError):
            parseJsonWithModel(jsonStr, ActionDefinition)
    
    def test_parseJsonWithModel_empty_string_raises(self):
        """Test that empty string raises ValueError"""
        with pytest.raises(ValueError):
            parseJsonWithModel("", ActionDefinition)
    
    def test_parseJsonWithModel_list_wraps_first_item(self):
        """Test that list JSON wraps first item"""
        jsonStr = '[{"action": "ai.process", "actionObjective": "Process"}]'
        result = parseJsonWithModel(jsonStr, ActionDefinition)
        assert isinstance(result, ActionDefinition)
        assert result.action == "ai.process"
    
    def test_parseJsonWithModel_aiResponse(self):
        """Test parsing AiResponse model"""
        jsonStr = '{"content": "Test content", "metadata": {"title": "Test"}}'
        result = parseJsonWithModel(jsonStr, AiResponse)
        assert isinstance(result, AiResponse)
        assert result.content == "Test content"
        assert result.metadata is not None
        assert result.metadata.title == "Test"


class TestRepairBrokenJson:
    """Test repairBrokenJson function"""
    
    def test_repairBrokenJson_incomplete_json(self):
        """Test repairing incomplete JSON"""
        brokenJson = '{"key": "value"'
        repaired = repairBrokenJson(brokenJson)
        # Should attempt to repair or return None
        assert repaired is None or isinstance(repaired, dict)
    
    def test_repairBrokenJson_missing_closing_brace(self):
        """Test repairing JSON with missing closing brace"""
        brokenJson = '{"documents": [{"sections": [{"id": "section_1"}]}'
        repaired = repairBrokenJson(brokenJson)
        # Should attempt to repair
        assert repaired is None or isinstance(repaired, dict)


if __name__ == "__main__":
    pytest.main([__file__, "-v"])