from agent.prompt_builder import PromptBuilder
from agent.models import ProgramMeta, FileInfo, CopyField


def test_build_basic_prompt():
    meta = ProgramMeta(
        program_id='ZAN04MAT', program_name='取消マッチング処理',
        system_name='残業統計管理システム',
        pgm_type='メイン', pgm_pattern='マッチング(1:1)',
        summary_lines=['取消申請のマッチング'], prerequisites=[],
        files=[], keys=[], modules=[],
        process_detail='1.初期処理...\n2.主処理...',
        output_records='### 出力ファイル1...',
        input_type='file', copy_fields={}, db_tables={}
    )

    builder = PromptBuilder()
    prompt = builder.build(
        meta=meta,
        rules_text='# マッチング(1:1) データ生成規則\n...',
        group_descriptions=['两端不一致', '逆方向两端不一致', '中间不一致'],
        group_count=3
    )

    assert 'ZAN04MAT' in prompt
    assert 'マッチング(1:1)' in prompt
    assert '1.初期処理' in prompt
    assert '生成するグループ数: 3' in prompt
    assert 'g1: 两端不一致' in prompt
    assert 'g2: 逆方向两端不一致' in prompt
    assert 'g3: 中间不一致' in prompt
    assert 'JSON形式' in prompt


def test_build_with_copy_fields():
    fields = {
        'R01': [
            CopyField(level=3, name='R01-APPL-ID', raw_name='(A)APPL-ID',
                      pic_type='X(008)', pic_bytes=8),
            CopyField(level=3, name='R01-EMP-ID', raw_name='(A)EMP-ID',
                      pic_type='9(008)', pic_bytes=8),
        ]
    }

    meta = ProgramMeta(
        program_id='TEST', program_name='', system_name='',
        pgm_type='メイン', pgm_pattern='マッチング(1:1)',
        summary_lines=[], prerequisites=[],
        files=[
            FileInfo(no=1, file_db_name='INPUT-FILE', identifier='R01',
                     dd_name='TESTR01', io='I', copy_group='TESTREC',
                     format='FB', record_len=80, medium='PS', remarks='')
        ],
        keys=[], modules=[],
        process_detail='', output_records='',
        input_type='file', copy_fields=fields, db_tables={}
    )

    builder = PromptBuilder()
    prompt = builder.build(
        meta=meta,
        rules_text='# 規則',
        group_descriptions=['テスト'],
        group_count=1
    )

    assert 'R01-APPL-ID' in prompt
    assert 'X(008)' in prompt
    assert 'R01-EMP-ID' in prompt


def test_build_db_prompt():
    meta = ProgramMeta(
        program_id='TESTDB', program_name='DB更新', system_name='',
        pgm_type='メイン', pgm_pattern='DB更新',
        summary_lines=[], prerequisites=[],
        files=[
            FileInfo(no=1, file_db_name='DB-TABLE', identifier='DB',
                     dd_name='', io='I', copy_group='',
                     format='', record_len=0, medium='DB', remarks='')
        ],
        keys=[], modules=[],
        process_detail='', output_records='',
        input_type='db', copy_fields={}, db_tables={}
    )

    builder = PromptBuilder()
    prompt = builder.build(
        meta=meta,
        rules_text='# DB更新規則',
        group_descriptions=['テスト'],
        group_count=1
    )

    assert 'DB' in prompt or 'SQL' in prompt
    assert 'INSERT' in prompt