from __future__ import annotations

import tempfile
import unittest
from pathlib import Path

from us_cpa.document_extractors import extract_document_facts


class DocumentExtractorTests(unittest.TestCase):
    def test_extracts_common_w2_fields(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            path = Path(temp_dir) / "w2.txt"
            path.write_text(
                "Form W-2 Wage and Tax Statement\n"
                "Employee: Jane Doe\n"
                "Box 1 Wages, tips, other compensation 50000\n"
                "Box 2 Federal income tax withheld 6000\n"
                "Box 16 State wages, tips, etc. 50000\n"
                "Box 17 State income tax 1200\n"
                "Box 3 Social security wages 50000\n"
                "Box 5 Medicare wages and tips 50000\n"
            )

            extracted = extract_document_facts(path)

            self.assertEqual(extracted["taxpayer.fullName"], "Jane Doe")
            self.assertEqual(extracted["wages"], 50000.0)
            self.assertEqual(extracted["federalWithholding"], 6000.0)
            self.assertEqual(extracted["stateWages"], 50000.0)
            self.assertEqual(extracted["stateWithholding"], 1200.0)
            self.assertEqual(extracted["socialSecurityWages"], 50000.0)
            self.assertEqual(extracted["medicareWages"], 50000.0)

    def test_extracts_common_1099_patterns(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            div_path = Path(temp_dir) / "1099-div.txt"
            div_path.write_text("Form 1099-DIV\nRecipient: Jane Doe\nBox 1a Total ordinary dividends 250\n")
            ret_path = Path(temp_dir) / "1099-r.txt"
            ret_path.write_text("Form 1099-R\nRecipient: Jane Doe\nBox 1 Gross distribution 10000\n")
            misc_path = Path(temp_dir) / "1099-misc.txt"
            misc_path.write_text("Form 1099-MISC\nRecipient: Jane Doe\nBox 3 Other income 900\n")

            self.assertEqual(extract_document_facts(div_path)["ordinaryDividends"], 250.0)
            self.assertEqual(extract_document_facts(ret_path)["retirementDistribution"], 10000.0)
            self.assertEqual(extract_document_facts(misc_path)["otherIncome"], 900.0)

    def test_extracts_prior_year_return_summary_values(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            path = Path(temp_dir) / "prior-return.txt"
            path.write_text(
                "2024 Form 1040 Summary\n"
                "Adjusted gross income 72100\n"
                "Taxable income 49800\n"
                "Refund 2100\n"
            )

            extracted = extract_document_facts(path)

            self.assertEqual(extracted["priorYear.adjustedGrossIncome"], 72100.0)
            self.assertEqual(extracted["priorYear.taxableIncome"], 49800.0)
            self.assertEqual(extracted["priorYear.refund"], 2100.0)


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