"""Approved PM exclusion contracts and additive Finance upgrade; all inputs synthetic."""
from __future__ import annotations

import copy
from dataclasses import replace
import json
import os
from pathlib import Path
import tempfile
import unittest

from openpyxl import load_workbook

from arr_ingestion.contracts import DeliveryEnvelope, IngestionError
from arr_ingestion.postgres import DatabaseConfig, PostgresIngestionRepository, _outcome_counts
from arr_ingestion.validation import DeliveryValidator, _validate_structured_payload_v4
from arr_processing.policy import load_processor_policy
from tests.local_postgres import TemporaryPostgres
from tests import test_arr_direct_data as direct_helpers
from tests.test_arr_ingestion_validation import build_delivery, policy
from tests.test_arr_opera_daily_ingest import reservation, xml_document


PROJECT = Path(__file__).resolve().parents[1]


def pm_row(*, status=None):
    value = reservation(1, rate_code="", rate_amount="0", company="").replace(
        "<ROOM_CATEGORY_LABEL>SYNTHETIC ROOM TYPE</ROOM_CATEGORY_LABEL>",
        "<ROOM_CATEGORY_LABEL> pm </ROOM_CATEGORY_LABEL>",
    ).replace("<DISP_ROOM_NO>SYNTHETIC-ROOM-1</DISP_ROOM_NO>", "<DISP_ROOM_NO/>")
    if status:
        value = value.replace("</G_RESERVATION>", f"<RESV_STATUS>{status}</RESV_STATUS></G_RESERVATION>")
    return value


class PMIngestionContractTests(unittest.TestCase):
    def generated(self, root, *rows):
        raw, store, _ = build_delivery(xml_document(*rows), root)
        envelope = DeliveryEnvelope.from_dict(json.loads(raw))
        payload = json.loads(store.objects[envelope.artifacts["structured_result_json"].object_key])
        return raw, store, envelope, payload

    def test_pm_missing_required_fields_is_audited_before_validation_and_pricing(self):
        with tempfile.TemporaryDirectory() as temporary:
            raw, store, _, payload = self.generated(Path(temporary), pm_row(), reservation(2))
            verified = DeliveryValidator(store, policy()).validate(raw)
            self.assertEqual(verified.envelope.status, "success")
            self.assertEqual(payload["outcome_counts"]["excluded_pm"], 1)
            self.assertEqual(payload["output_rows"], 1)
            row = payload["records"][0]
            self.assertEqual(row["outcome"], "excluded_pm")
            self.assertEqual(row["decision_codes"], ["ROOM_TYPE_PM_EXCLUDED"])
            for field in ("real_price", "total_price", "kb_amount", "channel_key", "pricing_method"):
                self.assertIsNone(row[field])

    def test_pm_count_is_required_only_for_44_and_older_versions_stay_readable(self):
        with tempfile.TemporaryDirectory() as temporary:
            _, _, envelope, payload = self.generated(Path(temporary), reservation(1))
            self.assertEqual(payload["processor_version"], "4.4.0")
            self.assertEqual(payload["outcome_counts"]["excluded_pm"], 0)
            for field in ("excluded_pm", "excluded_cancelled"):
                missing = copy.deepcopy(payload)
                del missing["outcome_counts"][field]
                with self.assertRaisesRegex(IngestionError, "outcome counts contract"):
                    _validate_structured_payload_v4(missing, envelope)
            for version in ("4.3.0", "4.2.0"):
                older = copy.deepcopy(payload)
                older["processor_version"] = version
                with self.assertRaisesRegex(IngestionError, "outcome counts contract"):
                    _validate_structured_payload_v4(older, replace(envelope, processor_version=version))
                del older["outcome_counts"]["excluded_pm"]
                if version == "4.2.0":
                    del older["outcome_counts"]["excluded_cancelled"]
                _validate_structured_payload_v4(older, replace(envelope, processor_version=version))

    def test_pm_record_cannot_be_relabelled_or_given_price_facts(self):
        with tempfile.TemporaryDirectory() as temporary:
            _, _, envelope, payload = self.generated(Path(temporary), pm_row(), reservation(2))
            changes = {
                "room_category_label": "PM-SUITE",
                "decision_codes": ["RATE_CODE_EXCLUDED"],
                "real_price": 0,
                "total_price": 0,
                "kb_amount": 0,
                "channel_key": "Group",
                "pricing_method": "zero_price_exception",
            }
            for field, value in changes.items():
                changed = copy.deepcopy(payload)
                changed["records"][0][field] = value
                with self.subTest(field=field), self.assertRaisesRegex(IngestionError, "PM exclusion"):
                    _validate_structured_payload_v4(changed, envelope)
            older = copy.deepcopy(payload)
            older["processor_version"] = "4.3.0"
            del older["outcome_counts"]["excluded_pm"]
            with self.assertRaisesRegex(IngestionError, "structured outcome is invalid"):
                _validate_structured_payload_v4(older, replace(envelope, processor_version="4.3.0"))
            changed = copy.deepcopy(payload)
            changed["outcome_counts"]["excluded_pm"] = 0
            with self.assertRaisesRegex(IngestionError, "outcomes do not reconcile"):
                _validate_structured_payload_v4(changed, envelope)

    def test_pm_balances_price_review_and_failure_and_cancelled_takes_precedence(self):
        cases = (
            (reservation(2, rate_amount="1800"), "review_required"),
            (reservation(2, departure="2026-07-26"), "failed"),
        )
        for row, status in cases:
            with self.subTest(status=status), tempfile.TemporaryDirectory() as temporary:
                raw, store, _, payload = self.generated(Path(temporary), pm_row(), row)
                self.assertEqual(DeliveryValidator(store, policy()).validate(raw).envelope.status, status)
                self.assertEqual(payload["outcome_counts"]["excluded_pm"], 1)
        with tempfile.TemporaryDirectory() as temporary:
            raw, store, _, payload = self.generated(Path(temporary), pm_row(status="CXL"), reservation(2))
            DeliveryValidator(store, policy()).validate(raw)
            self.assertEqual(payload["outcome_counts"]["excluded_cancelled"], 1)
            self.assertEqual(payload["outcome_counts"]["excluded_pm"], 0)

    def test_finance_defaults_only_old_pm_counts_and_rejects_44_missing_keys(self):
        counts = dict(retained=2, excluded_rate_code=1, duplicate=3,
                      validation_failed=4, price_unmatched=5)
        self.assertEqual(_outcome_counts({"outcome_counts": counts}), (2, 1, 3, 4, 5, 0, 0))
        counts["excluded_cancelled"] = 6
        self.assertEqual(_outcome_counts({"processor_version": "4.3.0", "outcome_counts": counts}),
                         (2, 1, 3, 4, 5, 6, 0))
        with self.assertRaises(IngestionError):
            _outcome_counts({"processor_version": "4.4.0", "outcome_counts": counts})
        counts["excluded_pm"] = 7
        self.assertEqual(_outcome_counts({"processor_version": "4.4.0", "outcome_counts": counts}),
                         (2, 1, 3, 4, 5, 6, 7))
        del counts["excluded_cancelled"]
        with self.assertRaises(IngestionError):
            _outcome_counts({"processor_version": "4.4.0", "outcome_counts": counts})


@unittest.skipUnless(os.environ.get("ARR_TEST_LOCAL_POSTGRES") == "1", "owned PostgreSQL opt-in required")
class PMFinanceMigrationTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.database = TemporaryPostgres().__enter__()
        cls.addClassCleanup(cls.database.__exit__, None, None, None)

    def test_additive_upgrade_preserves_cancelled_history_and_role_privileges(self):
        import psycopg
        self.database.reset_database(schema_version=20)
        repository = PostgresIngestionRepository(DatabaseConfig("owned-fixture"), connect=self.database.connect)
        repository.assert_cancelled_scope_schema()
        with self.assertRaises(IngestionError) as missing:
            repository.assert_pm_scope_schema()
        self.assertEqual(missing.exception.code, "DATABASE_MIGRATION_MISSING")
        with self.database.connect(autocommit=True) as connection:
            artifact = connection.execute("""INSERT INTO ingestion.artifacts
                (artifact_kind,storage_provider,bucket_alias,object_key,original_filename,sha256,byte_size)
                VALUES ('opera_xml','local_fixture','fixture','old/source.xml','source.xml',%s,0) RETURNING id""",
                ("a" * 64,)).fetchone()[0]
            output_ids = {}
            for kind in ("daily_xlsx", "result_json", "structured_result_json"):
                output_ids[kind] = connection.execute("""INSERT INTO ingestion.artifacts
                    (artifact_kind,storage_provider,bucket_alias,object_key,original_filename,sha256,byte_size)
                    VALUES (%s,'local_fixture','fixture',%s,%s,%s,0) RETURNING id""",
                    (kind, "old/" + kind, kind, "d" * 64)).fetchone()[0]
            run = connection.execute("""INSERT INTO ingestion.processing_runs
                (run_key,pipeline_type,source_artifact_id,result_artifact_id,run_status,result_delivery_mode,business_date,
                 delivered_processor_version,delivered_rule_set_sha256,result_schema_version,
                 delivery_sha256,validated_at,finished_at)
                VALUES ('old-cancelled','opera_daily',%s,%s,'accepted','artifact_callback','2026-07-27',
                        '4.3.0',%s,'4.0',%s,now(),now()) RETURNING id""",
                (artifact, output_ids["result_json"], "b" * 64, "c" * 64)).fetchone()[0]
            version = connection.execute("""INSERT INTO finance.daily_versions
                (business_date,version_no,processing_run_id,source_artifact_id,daily_report_artifact_id,
                 result_json_artifact_id,structured_result_artifact_id,version_status,
                 processor_version,rule_set_sha256,result_schema_version,result_sha256,source_rows,
                 retained_rows,excluded_rate_code_rows,duplicate_rows,validation_failed_rows,
                 price_unmatched_rows,excluded_cancelled_rows,validated_at,result_delivery_mode)
                VALUES ('2026-07-27',1,%s,%s,%s,%s,%s,'validated','4.3.0',%s,'4.0',%s,
                        2,1,0,0,0,0,1,now(),'artifact_callback') RETURNING id""",
                (run, artifact, output_ids["daily_xlsx"], output_ids["result_json"],
                 output_ids["structured_result_json"], "b" * 64, "c" * 64)).fetchone()[0]
            connection.execute("""INSERT INTO finance.daily_records
                (daily_version_id,source_sequence,source_location,outcome,decision_codes,
                 booking_source_match_status)
                VALUES (%s,1,'reservation[1]','excluded_cancelled',
                        ARRAY['RESERVATION_CANCELLED_EXCLUDED'],'missing_group_code')""", (version,))
            # Processor 4.3 included this PM row. A schema upgrade must preserve it.
            connection.execute("""INSERT INTO finance.daily_records
                (daily_version_id,source_sequence,source_location,outcome,decision_codes,block_code,
                 adults,children,company_name,company_key,confirmation_no,disp_room_no,
                 effective_rate_amount,full_name,no_of_rooms,products,rate_code,room_category_label,
                 arrival,departure,nights,real_price,total_price,channel_key,pricing_method,
                 booking_source_match_status)
                VALUES (%s,2,'reservation[2]','retained',ARRAY['PRICE_REFERENCE_EXACT'],'',2,0,
                        'SYNTHETIC COMPANY','SYNTHETIC COMPANY','SYNTHETIC CONF','SYNTHETIC ROOM',
                        900,'SYNTHETIC GUEST',1,'SYNTHETIC','GRPA1','PM','2026-07-27','2026-07-28',
                        1,900,900,'Group','price_reference_exact','missing_group_code')""", (version,))
            before = connection.execute("SELECT to_jsonb(v) FROM finance.daily_versions v WHERE id=%s", (version,)).fetchone()[0]
            records_before = connection.execute("SELECT to_jsonb(r) FROM finance.daily_records r ORDER BY source_sequence").fetchall()
            guards = connection.execute("SELECT tgname,pg_get_triggerdef(oid) FROM pg_trigger WHERE NOT tgisinternal ORDER BY tgname").fetchall()
            connection.execute("CREATE ROLE arr_pm_fixture")
            connection.execute("CREATE ROLE arr_pm_column_fixture")
            connection.execute("GRANT USAGE ON SCHEMA finance TO arr_pm_fixture,arr_pm_column_fixture")
            connection.execute("GRANT SELECT,INSERT ON finance.daily_versions TO arr_pm_fixture")
            connection.execute("GRANT INSERT (excluded_cancelled_rows) ON finance.daily_versions TO arr_pm_column_fixture")
            connection.execute((PROJECT / "database/021_daily_pm_exclusion.sql").read_text(), prepare=False)
            repository.assert_data_source_schema()
            after = connection.execute("SELECT to_jsonb(v) FROM finance.daily_versions v WHERE id=%s", (version,)).fetchone()[0]
            self.assertEqual(after.pop("excluded_pm_rows"), 0)
            self.assertEqual(after, before)
            self.assertEqual(connection.execute("SELECT to_jsonb(r) FROM finance.daily_records r ORDER BY source_sequence").fetchall(), records_before)
            self.assertEqual(connection.execute("SELECT tgname,pg_get_triggerdef(oid) FROM pg_trigger WHERE NOT tgisinternal ORDER BY tgname").fetchall(), guards)
            self.assertEqual(connection.execute("SELECT has_column_privilege('arr_pm_fixture','finance.daily_versions','excluded_pm_rows','INSERT'),has_table_privilege('arr_pm_fixture','finance.daily_versions','UPDATE'),has_column_privilege('arr_pm_column_fixture','finance.daily_versions','excluded_pm_rows','INSERT')").fetchone(), (True, False, False))
            for statement in (
                "UPDATE finance.daily_versions SET excluded_pm_rows=-1 WHERE id=%s",
                "UPDATE finance.daily_versions SET excluded_pm_rows=1 WHERE id=%s",
            ):
                with self.assertRaises(psycopg.errors.CheckViolation):
                    connection.execute(statement, (version,))
            connection.execute("UPDATE finance.daily_versions SET excluded_cancelled_rows=0,excluded_pm_rows=1 WHERE id=%s", (version,))
            with self.assertRaisesRegex(psycopg.errors.RaiseException, "PM-exclusion facts exist"):
                connection.execute((PROJECT / "database/021_daily_pm_exclusion.down.sql").read_text(), prepare=False)
            connection.execute("ROLLBACK")
            connection.execute("SET ROLE arr_pm_fixture")
            with self.assertRaises(psycopg.errors.InsufficientPrivilege):
                connection.execute("UPDATE finance.daily_versions SET excluded_pm_rows=0 WHERE id=%s", (version,))
            connection.execute("RESET ROLE")

    def test_empty_rollback_preserves_020_and_allows_reapply(self):
        self.database.reset_database()
        repository = PostgresIngestionRepository(DatabaseConfig("owned-fixture"), connect=self.database.connect)
        repository.assert_data_source_schema()
        with self.database.connect(autocommit=True) as connection:
            connection.execute((PROJECT / "database/021_daily_pm_exclusion.down.sql").read_text(), prepare=False)
        repository.assert_cancelled_scope_schema()
        with self.assertRaises(IngestionError) as missing:
            repository.assert_pm_scope_schema()
        self.assertEqual(missing.exception.code, "DATABASE_MIGRATION_MISSING")
        with self.database.connect(autocommit=True) as connection:
            connection.execute((PROJECT / "database/021_daily_pm_exclusion.sql").read_text(), prepare=False)
        repository.assert_data_source_schema()


@unittest.skipUnless(os.environ.get("ARR_TEST_LOCAL_POSTGRES") == "1", "owned PostgreSQL opt-in required")
class PMFinanceProcessingTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.database = TemporaryPostgres().__enter__()
        cls.addClassCleanup(cls.database.__exit__, None, None, None)
        cls.policy = load_processor_policy(PROJECT)

    setUp = direct_helpers.DirectDataPostgresTests.setUp
    configure = direct_helpers.DirectDataTests.configure
    execute = direct_helpers.DirectDataTests.execute
    sql = direct_helpers.DirectDataPostgresTests.sql
    count = direct_helpers.DirectDataPostgresTests.count
    monthly_worker = direct_helpers.DirectDataPostgresTests.monthly_worker

    def make_first_pm(self):
        row = self.transport.rows[0]
        row["roomStay"]["roomType"] = "PM"
        row["roomStay"]["currentRoomInfo"]["roomType"] = "PM"
        for rate in row["roomStay"]["roomRates"]:
            rate["roomType"] = "PM"

    def test_direct_pm_is_audited_and_daily_monthly_exclude_it_without_refetch(self):
        self.make_first_pm()
        self.assertEqual(self.execute().status, "succeeded")
        self.assertEqual(self.sql("SELECT result_schema_version,source_rows,retained_rows,excluded_pm_rows FROM finance.daily_versions"), [("5.0", 6, 5, 1)])
        self.assertEqual(self.sql("SELECT outcome,decision_codes,real_price,channel_key FROM finance.daily_records WHERE source_sequence=1"), [("excluded_pm", ["ROOM_TYPE_PM_EXCLUDED"], None, None)])
        self.assertEqual(self.count("finance.v_active_daily_facts"), 5)
        self.assertEqual(self.monthly_worker().process_next().status, "published")
        key, size = self.sql("SELECT a.object_key,a.byte_size FROM reporting.monthly_runs r JOIN ingestion.artifacts a ON a.id=r.workbook_artifact_id")[0]
        output = self.files / "monthly.xlsx"
        self.store.materialize(key, output, size)
        workbook = load_workbook(output, data_only=False)
        try:
            self.assertEqual(sum(cell.data_type == "f" for sheet in workbook for row in sheet for cell in row), 5)
        finally:
            workbook.close()
        calls = len(self.transport.calls)
        self.assertEqual(self.execute().status, "succeeded")
        self.assertEqual(len(self.transport.calls), calls)
        self.assertEqual(self.count("finance.daily_versions"), 1)

    def test_xml_pm_is_audited_without_required_field_or_price_review(self):
        outcome = self.coordinator.submit("source.xml", xml_document(pm_row(), reservation(2)).encode())
        self.assertEqual(outcome["status"], "succeeded")
        self.assertEqual(self.sql("SELECT result_schema_version,source_rows,retained_rows,excluded_pm_rows FROM finance.daily_versions"), [("4.0", 2, 1, 1)])
        self.assertEqual(self.sql("SELECT outcome FROM finance.daily_records ORDER BY source_sequence"), [("excluded_pm",), ("retained",)])
        self.assertEqual(self.count("finance.v_active_daily_facts"), 1)
        self.assertEqual(self.monthly_worker().process_next().status, "published")

    def test_xml_failed_batch_keeps_pm_audit_without_monthly_commit(self):
        outcome = self.coordinator.submit("source.xml", xml_document(pm_row(), reservation(2, departure="2026-07-26")).encode())
        self.assertEqual(outcome["status"], "failed")
        self.assertEqual(self.sql("SELECT version_status,source_rows,excluded_pm_rows,validation_failed_rows FROM finance.daily_versions"), [("rejected", 2, 1, 1)])
        self.assertEqual(self.count("finance.current_daily_versions"), 0)
        self.assertEqual(self.monthly_worker().process_next().status, "idle")

    def test_database_pm_record_classification_guard_rejects_wrong_room_or_price(self):
        import psycopg
        self.make_first_pm()
        self.assertEqual(self.execute().status, "succeeded")
        with self.database.connect(autocommit=True) as connection:
            for statement in (
                "UPDATE finance.daily_records SET room_category_label='PM-SUITE' WHERE source_sequence=1",
                "UPDATE finance.daily_records SET real_price=0 WHERE source_sequence=1",
                "UPDATE finance.daily_records SET channel_key='Group' WHERE source_sequence=1",
                "UPDATE finance.daily_records SET decision_codes=ARRAY['RATE_CODE_EXCLUDED'] WHERE source_sequence=1",
            ):
                with self.subTest(statement=statement), self.assertRaises(psycopg.errors.CheckViolation):
                    connection.execute(statement)


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