from pathlib import Path

import pytest

from dotpromptz.rembg import run_rembg


def _install_fake_rembg(monkeypatch: pytest.MonkeyPatch, calls: dict[str, object]) -> None:
    def _fake_new_session(*, model_name: str) -> dict[str, str]:
        models = calls.setdefault('models', [])
        assert isinstance(models, list)
        models.append(model_name)
        return {'model': model_name}

    def _fake_remove(data: bytes, *, session: object) -> bytes:
        sessions = calls.setdefault('sessions', [])
        assert isinstance(sessions, list)
        sessions.append(session)
        return b'processed:' + data

    monkeypatch.setattr('dotpromptz.rembg._load_rembg_api', lambda: (_fake_new_session, _fake_remove))


def test_run_rembg_single_file_writes_output(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    calls: dict[str, object] = {}
    _install_fake_rembg(monkeypatch, calls)
    input_file = tmp_path / 'a.jpg'
    input_file.write_bytes(b'hello')

    output = tmp_path / 'out' / 'result.png'
    written = run_rembg(input_file, output_path=output, model='u2netp')

    assert written == [output]
    assert output.read_bytes() == b'processed:hello'
    assert calls['models'] == ['u2netp']


def test_run_rembg_directory_only_scans_top_level(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    calls: dict[str, object] = {}
    _install_fake_rembg(monkeypatch, calls)
    input_dir = tmp_path / 'images'
    input_dir.mkdir()
    (input_dir / 'a.jpg').write_bytes(b'a')
    nested = input_dir / 'nested'
    nested.mkdir()
    (nested / 'b.jpg').write_bytes(b'b')

    output_dir = tmp_path / 'output'
    written = run_rembg(input_dir, output_path=output_dir)

    assert written == [output_dir / 'a.png']
    assert (output_dir / 'a.png').read_bytes() == b'processed:a'
    assert not (output_dir / 'b.png').exists()


def test_run_rembg_file_requires_file_output(tmp_path: Path) -> None:
    input_file = tmp_path / 'a.jpg'
    input_file.write_bytes(b'a')
    output_dir = tmp_path / 'output'
    output_dir.mkdir()

    with pytest.raises(ValueError, match='must be a file'):
        run_rembg(input_file, output_path=output_dir)


def test_run_rembg_directory_requires_directory_output(tmp_path: Path) -> None:
    input_dir = tmp_path / 'images'
    input_dir.mkdir()
    (input_dir / 'a.jpg').write_bytes(b'a')
    output_file = tmp_path / 'output.png'
    output_file.write_bytes(b'x')

    with pytest.raises(ValueError, match='must be a directory'):
        run_rembg(input_dir, output_path=output_file)


def test_run_rembg_existing_output_requires_force(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    calls: dict[str, object] = {}
    _install_fake_rembg(monkeypatch, calls)
    input_file = tmp_path / 'a.jpg'
    input_file.write_bytes(b'a')
    output_file = tmp_path / 'result.png'
    output_file.write_bytes(b'old')

    with pytest.raises(ValueError, match='already exists'):
        run_rembg(input_file, output_path=output_file)

    written = run_rembg(input_file, output_path=output_file, force=True)

    assert written == [output_file]
    assert output_file.read_bytes() == b'processed:a'


def test_run_rembg_directory_reuses_single_session(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    calls: dict[str, object] = {}
    _install_fake_rembg(monkeypatch, calls)
    input_dir = tmp_path / 'images'
    input_dir.mkdir()
    (input_dir / 'a.jpg').write_bytes(b'a')
    (input_dir / 'b.jpeg').write_bytes(b'b')

    output_dir = tmp_path / 'output'
    written = run_rembg(input_dir, output_path=output_dir, model='birefnet-general')

    assert written == [output_dir / 'a.png', output_dir / 'b.png']
    assert calls['models'] == ['birefnet-general']
    sessions = calls['sessions']
    assert isinstance(sessions, list)
    assert len(sessions) == 2
    assert sessions[0] is sessions[1]


def test_run_rembg_directory_without_supported_files_raises(tmp_path: Path) -> None:
    input_dir = tmp_path / 'images'
    input_dir.mkdir()
    (input_dir / 'a.txt').write_text('hello', encoding='utf-8')

    with pytest.raises(ValueError, match='No supported image files found'):
        run_rembg(input_dir, output_path=tmp_path / 'output')
