import json
from unittest.mock import patch, MagicMock
from agent.api_client import APIClient


def test_parse_json_valid():
    text = '{"groups": {"g1": {"records": []}}}'
    result = APIClient._parse_json(text)
    assert result is not None
    assert 'groups' in result


def test_parse_json_with_markdown_wrapper():
    text = '```json\n{"key": "value"}\n```'
    result = APIClient._parse_json(text)
    assert result is not None
    assert result['key'] == 'value'


def test_parse_json_with_text_before():
    text = '说明\n{"key": "value"}\n以上'
    result = APIClient._parse_json(text)
    assert result is not None
    assert result['key'] == 'value'


def test_parse_json_invalid():
    text = '这不是有效的JSON。'
    result = APIClient._parse_json(text)
    assert result is None


@patch('agent.api_client.requests.post')
def test_generate_success(mock_post):
    mock_response = MagicMock()
    mock_response.json.return_value = {
        'choices': [{'message': {'content': '{"groups": {"g1": {"type": "json"}}}'}}]
    }
    mock_response.raise_for_status = MagicMock()
    mock_post.return_value = mock_response

    client = APIClient(api_key='test-key')
    result = client.generate("test prompt")

    assert result['groups']['g1']['type'] == 'json'


@patch('agent.api_client.requests.post')
def test_generate_retry_on_json_error(mock_post):
    bad_response = MagicMock()
    bad_response.json.return_value = {
        'choices': [{'message': {'content': 'invalid response'}}]
    }
    bad_response.raise_for_status = MagicMock()

    good_response = MagicMock()
    good_response.json.return_value = {
        'choices': [{'message': {'content': '{"result": "ok"}'}}]
    }
    good_response.raise_for_status = MagicMock()

    mock_post.side_effect = [bad_response, good_response]

    client = APIClient(api_key='test-key', max_retries=3)
    result = client.generate("test")

    assert result['result'] == 'ok'
    assert mock_post.call_count == 2