tests/test_ui_input_choice.py

import unittest
from tempfile import NamedTemporaryFile

from shop_ledger.processor import extract_document_text, prepare_document_input
from shop_ledger.ui import (
    add_to_ledger,
    apply_row_correction,
    ask_ledger,
    ask_ledger_chat,
    ask_ledger_voice_chat,
    choose_input,
    compose_chart,
    generate_daily_brief,
    initial_ask_chat,
    run_command_palette,
)


class InputChoiceTests(unittest.TestCase):
    def test_auto_asks_when_text_and_audio_exist(self):
        choice = choose_input("paid Ravi 1200", "/tmp/audio.wav", None, "Auto")

        self.assertEqual(choice["status"], "conflict")
        self.assertIn("Multiple inputs", choice["notice"])

    def test_text_choice_uses_text_when_audio_exists(self):
        choice = choose_input("paid Ravi 1200", "/tmp/audio.wav", None, "Text note")

        self.assertEqual(choice["status"], "ready")
        self.assertEqual(choice["source"], "text")

    def test_auto_uses_audio_when_audio_is_only_input(self):
        choice = choose_input("", "/tmp/audio.wav", None, "Auto")

        self.assertEqual(choice["status"], "ready")
        self.assertEqual(choice["source"], "audio")

    def test_auto_uses_document_when_document_is_only_input(self):
        choice = choose_input("", None, "/tmp/receipt.pdf", "Auto")

        self.assertEqual(choice["status"], "ready")
        self.assertEqual(choice["source"], "document")

    def test_document_text_extraction_reads_plain_text_files(self):
        with NamedTemporaryFile("w", suffix=".txt") as handle:
            handle.write("paid Ravi 1200 for rice bags")
            handle.flush()

            text = extract_document_text(handle.name)

        self.assertIn("Ravi", text)

    def test_document_image_preparation_creates_data_url(self):
        from PIL import Image

        with NamedTemporaryFile(suffix=".png") as handle:
            Image.new("RGB", (8, 8), color="white").save(handle.name)

            document = prepare_document_input(handle.name)

        self.assertEqual(document["kind"], "image")
        self.assertTrue(document["image_urls"][0].startswith("data:image/jpeg;base64,"))

    def test_successful_text_add_clears_written_note(self):
        def fake_process(note, currency, image_urls=None):
            return {
                "entries": [
                    {
                        "date": "2026-06-11",
                        "direction": "expense",
                        "counterparty": "Ravi",
                        "item": "rice bags",
                        "quantity": "",
                        "amount": 1200,
                        "currency": currency,
                        "category": "inventory",
                        "payment_status": "paid",
                        "due_date": "",
                        "confidence": 0.9,
                        "reminder": "",
                    }
                ],
                "reminders": [],
                "questions": [],
                "model_used": "fake",
            }

        output = add_to_ledger("paid Ravi 1200", None, None, "Auto", "LKR", [], fake_process)

        self.assertEqual(len(output[6]), 1)
        self.assertEqual(output[7]["value"], "")
        self.assertEqual(output[10]["value"], "Auto")
        self.assertIn("Added 1 row", output[11])

    def test_successful_document_add_sends_image_urls_and_clears_file(self):
        captured = {}

        def fake_process(note, currency, image_urls=None):
            captured["note"] = note
            captured["image_urls"] = image_urls
            return {
                "entries": [
                    {
                        "date": "2026-06-11",
                        "direction": "expense",
                        "counterparty": "Ravi",
                        "item": "rice bags",
                        "quantity": "",
                        "amount": 1200,
                        "currency": currency,
                        "category": "inventory",
                        "payment_status": "paid",
                        "due_date": "",
                        "confidence": 0.9,
                        "reminder": "",
                    }
                ],
                "reminders": [],
                "questions": [],
                "model_used": "fake",
            }

        with NamedTemporaryFile("w", suffix=".txt") as handle:
            handle.write("paid Ravi 1200 for rice bags")
            handle.flush()

            output = add_to_ledger("", None, handle.name, "Document", "LKR", [], fake_process)

        self.assertIn("paid Ravi", captured["note"])
        self.assertIsNone(captured["image_urls"])
        self.assertEqual(output[9]["value"], None)
        self.assertIn("Added 1 row", output[11])

    def test_generate_daily_brief_uses_supplied_function(self):
        rows = [{"amount": 1200, "currency": "LKR", "direction": "expense", "payment_status": "paid"}]

        markdown = generate_daily_brief(
            rows,
            "LKR",
            lambda supplied_rows, currency: {"brief": f"{len(supplied_rows)} rows in {currency}", "model_used": "fake"},
        )

        self.assertIn("1 rows in LKR", markdown)
        self.assertIn("fake", markdown)

    def test_ask_ledger_uses_supplied_function(self):
        rows = [{"amount": 7500, "currency": "LKR", "payment_status": "due"}]

        markdown = ask_ledger(
            rows,
            "Who owes me most?",
            "LKR",
            lambda supplied_rows, question, currency: {"answer": f"{question} / {len(supplied_rows)}", "model_used": "fake"},
        )

        self.assertIn("Who owes me most?", markdown)
        self.assertIn("fake", markdown)

    def test_ask_ledger_chat_appends_messages_and_clears_input(self):
        rows = [{"amount": 7500, "currency": "LKR", "payment_status": "due"}]

        history, next_question = ask_ledger_chat(
            rows,
            "Who owes me most?",
            initial_ask_chat(),
            "LKR",
            lambda supplied_rows, question, currency: {"answer": "Nimal owes LKR 7,500.", "model_used": "fake"},
        )

        self.assertEqual(next_question, "")
        self.assertEqual(history[-2]["role"], "user")
        self.assertEqual(history[-1]["role"], "assistant")
        self.assertIn("Nimal", history[-1]["content"])

    def test_ask_ledger_voice_chat_transcribes_and_answers(self):
        history, next_question, next_audio = ask_ledger_voice_chat(
            [{"counterparty": "Nimal", "amount": 7500, "payment_status": "due", "currency": "LKR"}],
            "/tmp/question.wav",
            initial_ask_chat(),
            "LKR",
            lambda rows, question, currency: {"answer": f"Answered: {question}", "model_used": "fake"},
            transcribe_fn=lambda path: "Who owes me most?",
        )

        self.assertEqual(next_question, "")
        self.assertIsNone(next_audio)
        self.assertIn("Who owes me most?", history[-2]["content"])
        self.assertIn("Answered", history[-1]["content"])

    def test_ask_ledger_voice_chat_handles_empty_transcript(self):
        history, _, next_audio = ask_ledger_voice_chat(
            [],
            "/tmp/question.wav",
            initial_ask_chat(),
            "LKR",
            lambda rows, question, currency: {"answer": "unused", "model_used": "fake"},
            transcribe_fn=lambda path: "",
        )

        self.assertIsNone(next_audio)
        self.assertIn("could not hear", history[-1]["content"])

    def test_run_command_palette_uses_current_rows(self):
        rows = [{"payment_status": "due", "counterparty": "Nimal", "amount": 7500, "currency": "LKR", "item": "tea"}]

        output = run_command_palette(rows, "Show unpaid")

        self.assertIn("Nimal", output)

    def test_compose_chart_returns_markdown_figure_and_clears_input(self):
        rows = [{"payment_status": "due", "counterparty": "Nimal", "amount": 7500, "currency": "LKR"}]

        markdown, figure, next_question = compose_chart(
            rows,
            "Who owes me?",
            lambda supplied_rows, question: {"chart": "due_by_party", "reason": "Dues", "model_used": "fake"},
        )

        self.assertIn("Due radar", markdown)
        self.assertTrue(hasattr(figure, "to_plotly_json"))
        self.assertEqual(next_question, "")

    def test_apply_row_correction_updates_state_and_confidence(self):
        rows = [
            {
                "date": "2026-06-11",
                "direction": "income",
                "counterparty": "",
                "item": "tea packets",
                "quantity": "",
                "amount": 750,
                "currency": "LKR",
                "category": "sales",
                "payment_status": "due",
                "due_date": "",
                "confidence": 0.42,
                "reminder": "",
            }
        ]

        output = apply_row_correction(rows, 1, "counterparty", "Nimal", "LKR")

        updated_rows = output[6]
        self.assertEqual(updated_rows[0]["counterparty"], "Nimal")
        self.assertEqual(updated_rows[0]["confidence"], 0.9)
        self.assertIn("Updated row 1", output[-1])

    def test_apply_row_correction_rejects_bad_amount(self):
        rows = [
            {
                "date": "2026-06-11",
                "direction": "income",
                "counterparty": "Nimal",
                "item": "tea packets",
                "quantity": "",
                "amount": 750,
                "currency": "LKR",
                "category": "sales",
                "payment_status": "due",
                "due_date": "",
                "confidence": 0.42,
                "reminder": "",
            }
        ]

        output = apply_row_correction(rows, 1, "amount", "many rupees", "LKR")

        self.assertEqual(output[6][0]["amount"], 750)
        self.assertIn("need a number", output[-1])


if __name__ == "__main__":
    unittest.main()