import json
import os
import tempfile
from agent.output_writer import OutputWriter


def test_write_json_files():
    with tempfile.TemporaryDirectory() as tmpdir:
        writer = OutputWriter(tmpdir)

        ai_result = {
            "groups": {
                "g1": {
                    "program": "TEST",
                    "records": [
                        {"input": {"R01": {"R01-FIELD": "A001"}}}
                    ]
                }
            }
        }

        written = writer.write("TESTPGM", ai_result, input_type="file")

        assert len(written) == 1
        json_path = written["g1/json"]
        assert os.path.exists(json_path)

        with open(json_path, 'r', encoding='utf-8') as f:
            content = json.load(f)
        assert content['program'] == 'TEST'
        assert len(content['records']) == 1


def test_write_mixed_type():
    with tempfile.TemporaryDirectory() as tmpdir:
        writer = OutputWriter(tmpdir)

        ai_result = {
            "groups": {
                "g1": {
                    "program": "MIXED",
                    "records": [{"input": {"R01": {"F1": "X"}}}],
                    "sql": "INSERT INTO T VALUES ('x');"
                }
            }
        }

        written = writer.write("MIXEDPGM", ai_result, input_type="mixed")

        assert len(written) == 2
        assert os.path.exists(written["g1/json"])
        assert os.path.exists(written["g1/sql"])


def test_write_multiple_groups():
    with tempfile.TemporaryDirectory() as tmpdir:
        writer = OutputWriter(tmpdir)

        ai_result = {
            "groups": {
                "g1": {"records": [{}]},
                "g2": {"records": [{}]},
                "g3": {"records": [{}]},
            }
        }

        written = writer.write("MULTI", ai_result, input_type="file")

        assert len(written) == 3
        for g in ['g1', 'g2', 'g3']:
            assert f'{g}/json' in written
            assert os.path.exists(written[f'{g}/json'])


def test_directory_structure():
    with tempfile.TemporaryDirectory() as tmpdir:
        writer = OutputWriter(tmpdir)
        ai_result = {"groups": {"g1": {"records": [{}]}}}
        writer.write("MYPROG", ai_result, input_type="file")

        expected_dir = os.path.join(tmpdir, "MYPROG", "g1")
        assert os.path.isdir(expected_dir)
        assert os.path.isfile(os.path.join(expected_dir, "MYPROG_g1.json"))