import os

import numpy as np
import pytest
from PIL import Image

WEIGHTS = os.environ.get("CELLQUANT_WEIGHTS", os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "weights", "cellquant.pt"))


@pytest.mark.skipif(not os.path.exists(WEIGHTS), reason="trained weights not present")
def test_analyze_files_end_to_end(tmp_path, monkeypatch):
    monkeypatch.delenv("GEMINI_API_KEY", raising=False)   # never call Gemini from the tests
    import app
    rng = np.random.default_rng(0)
    img = (rng.random((200, 240, 3)) * 50).astype(np.uint8)
    p1 = tmp_path / "a.png"
    Image.fromarray(img).save(p1)
    import tifffile
    p2 = tmp_path / "b.tif"
    tifffile.imwrite(str(p2), img)
    rows, gallery, zip_path, summary = app.analyze_files([str(p1), str(p2)], 0.5, 0.5, 50, "head", "contains")
    from cellquant.labels import COUNT_COLUMNS
    assert len(rows) == 2 and rows[0][0] == "a.png" and len(rows[0]) == 1 + len(COUNT_COLUMNS) == 9
    for r in rows:   # "All nuclei" after the four nucleus labels, "All puncta" after the two punctum labels
        assert r[5] == sum(r[1:5]) and r[8] == sum(r[6:8])
    assert len(gallery) == 2 and os.path.exists(gallery[0][0])
    assert gallery[0][1] == f"a.png: {rows[0][5]} nuclei, {rows[0][8]} puncta"
    assert os.path.exists(zip_path) and summary.startswith("Analysed 2 images in ") and "on the CPU" in summary
    assert "Gemini summary is off" in summary          # no key: the write-up is replaced by a one-line notice
    assert app._gpu_seconds(3) == 15 + 3 * {1: 4, 4: 6, 8: 10}[app.TTA] and app._gpu_seconds(0) == 15 + {1: 4, 4: 6, 8: 10}[app.TTA]
    with pytest.raises(Exception):
        app.analyze_files([], 0.5, 0.5, 50, "head", "contains")
    import csv
    import io
    import json
    import zipfile
    with zipfile.ZipFile(zip_path) as zf:
        names = zf.namelist()
        counts_csv = list(csv.reader(io.StringIO(zf.read("counts.csv").decode())))
        results = json.loads(zf.read("results.json"))
    assert "counts.csv" in names and "objects.csv" in names and "results.json" in names
    assert counts_csv[0] == ["image"] + list(COUNT_COLUMNS)
    assert counts_csv[0][5] == "All nuclei" and counts_csv[0][8] == "All puncta"
    assert [[r[0]] + [int(v) for v in r[1:]] for r in counts_csv[1:]] == rows
    assert results["cellquant_version"] == app.__version__
    assert results["images"]["a.png"]["totals"] == {"All nuclei": rows[0][5], "All puncta": rows[0][8]}
    assert results["images"]["a.png"]["totals"]["All nuclei"] == sum(results["images"]["a.png"]["counts"][lab] for lab in COUNT_COLUMNS[:4])


def test_build_demo_constructs():
    pytest.importorskip("gradio")
    import app
    demo = app.build_demo()
    assert demo is not None