"""Request-scoped masking manifest and response validation invariants."""
from __future__ import annotations

import json
import os
import re
import sys
import tempfile
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from unittest.mock import patch

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

os.environ["KANSHO_DB_PATH"] = str(Path(tempfile.gettempdir()) / "kansho-privacy-manifest-test.sqlite")
os.environ["KANSHO_FAKE_PROVIDER"] = "1"
os.environ["KANSHO_FAKE_DETECT"] = "1"
Path(os.environ["KANSHO_DB_PATH"]).unlink(missing_ok=True)

from fastapi.testclient import TestClient
from entity_detect import install_test_spans, reset_detect_test_hooks
from identity_store import remember_mapping
from main import app
from privacy_gateway import (
    ERROR_EGRESS_VALIDATION,
    ActiveReplacement,
    GatewayRequest,
    MaskingManifest,
    PrivacyGatewayError,
    _demask,
    _process_response,
    _validate_response,
    _visible_response,
    complete,
    last_trace,
    mask_prompt,
    reset_debug,
    validate_pre_egress,
)
from providers import ChatResult


def expect(ok: bool, message: str) -> None:
    if not ok:
        raise SystemExit(f"FAIL: {message}")
    print(f"OK  {message}")


def _no_label_keys(payload) -> bool:
    if isinstance(payload, dict):
        if "local_label" in payload:
            return False
        return all(_no_label_keys(value) for value in payload.values())
    if isinstance(payload, list):
        return all(_no_label_keys(item) for item in payload)
    return True


def _run_gateway(profile_id: str, rendered: str, *, purpose: str = "dialogue_turn"):
    return complete(
        GatewayRequest(
            prompt_id="manifest-test",
            purpose=purpose,
            data_class="B",
            profile_id=profile_id,
            payload={"rendered": rendered, "source_text": rendered},
        )
    )


def main() -> None:
    reset_debug()
    anna = [{"local_label": "Anna", "token": "PERSON:01", "entity_type": "PERSON"}]
    lars = [{"local_label": "Lars", "token": "PERSON:02", "entity_type": "PERSON"}]
    sushi = [{"local_label": "Sushi", "token": "PERSON:01", "entity_type": "PERSON"}]
    hamburg = [{"local_label": "Hamburg", "token": "PLACE:01", "entity_type": "PLACE"}]

    inactive = mask_prompt("Heute nur Markt und Kirschen.", anna, "journal_generate")
    expect(inactive.replacements == (), "historical mapping is inactive when absent from the prompt")
    kept = _validate_response("Anna kam später zum Markt.", inactive)
    expect("Anna kam später" in kept, "inactive mapping does not treat the same string as a leak")

    mixed = mask_prompt("Anna war in Hamburg.", anna + lars + hamburg, "journal_generate")
    active_tokens = {item.token for item in mixed.replacements}
    expect(active_tokens == {"PERSON:01", "PLACE:01"}, "only actually masked mappings are active")
    expect("PERSON:02" not in active_tokens, "unused profile mapping stays inactive")

    food_only = mask_prompt("Ich esse gerne Sushi.", sushi, "journal_generate")
    expect(food_only.replacements == (), "intentional homonym is not an active identity")
    food_reply = _validate_response("Heute nur von Sushi essen erzählt.", food_only)
    expect("Sushi essen" in food_reply, "unmasked homonym does not block a normal reply")

    kinship = mask_prompt("Meine Frau Sushi kam später.", sushi, "journal_generate")
    expect(len(kinship.replacements) == 1, "identity mention of the same word is active")
    normalized, meta = _process_response("Sushi kam vorbei.", kinship)
    expect(normalized.startswith("[[PERSON:01]]"), "active cleartext is remapped to the request token")
    expect(meta.get("response_normalization") == "active_cleartext_normalized", "normalization is named")
    expect(_demask(normalized, kinship) == "Sushi kam vorbei.", "normalized reply demasks locally")
    homonym_ok = _validate_response("Danach Sushi essen.", kinship)
    expect("Sushi essen" in homonym_ok, "homonym in the reply uses the same classification rule")

    placeholder = _validate_response("[[PERSON:01]] war ruhig.", kinship)
    expect(placeholder.startswith("[[PERSON:01]]"), "active placeholder passes validation")
    expect(_demask(placeholder, kinship) == "Sushi war ruhig.", "demask happens only after validation")
    leftover = _demask("[[PERSON:01]] und [[PERSON:99]]", kinship)
    expect("Sushi" in leftover and "[[PERSON:99]]" in leftover, "inactive placeholder is not rematerialized")
    expect("Clarissa" not in leftover, "inactive mapping label is not introduced by demask")
    import privacy_gateway as gw

    expect(not hasattr(gw, "IDENTITY_LEAK_RETRY"), "automatic identity-leak retry instruction is gone")
    expect(
        _demask("[[ person:01 ]] war ruhig.", kinship) == "Sushi war ruhig.",
        "demask ignores case and inner spacing",
    )
    expect(
        _demask("[[PERSON:1]] war ruhig.", kinship) == "Sushi war ruhig.",
        "demask accepts an unpadded numeric token",
    )
    expect(
        _demask("[[…]] und [[...]] blieben.", kinship) == "[[…]] und [[...]] blieben.",
        "generic ellipsis placeholders are not rematerialized",
    )

    alias_rows = [
        {
            "local_label": "Sushi",
            "token": "PERSON:07",
            "canonical_label": "Clarissa",
            "demask_label": "Clarissa",
            "aliases": ["Sushi"],
            "labels": ["Clarissa", "Sushi"],
            "entity_type": "PERSON",
        }
    ]
    aliased = mask_prompt("Sushi kam ins Wohnzimmer.", alias_rows, "journal_generate")
    expect(len(aliased.replacements) == 1, "confirmed alias activates its token")
    spellings = {label.casefold() for label in aliased.replacements[0].all_labels()}
    expect(spellings == {"sushi", "clarissa"}, "request manifest keeps every spelling of the active token")
    expect(_demask("[[PERSON:07]] kam.", aliased) == "Clarissa kam.", "confirmed alias demasks to the canonical spelling")
    canonical_clear, _meta = _process_response("Clarissa kam vorbei.", aliased)
    expect(canonical_clear.startswith("[[PERSON:07]]"), "canonical spelling of an active alias token is remapped")
    expect(
        _visible_response("Sushi kam vorbei.", aliased) == "Sushi kam vorbei.",
        "alias plaintext keeps the model surface form",
    )
    expect(
        _visible_response("Clarissa kam vorbei.", aliased) == "Clarissa kam vorbei.",
        "canonical plaintext keeps the model surface form",
    )
    expect(
        _visible_response("[[PERSON:07]] kam vorbei.", aliased) == "Clarissa kam vorbei.",
        "emitted placeholder still rehydrates to the canonical spelling",
    )

    group_prompt = "Heute die Gruppe am Markt."
    group_start = group_prompt.index("die Gruppe")
    group_rows = [
        {
            "local_label": "die Gruppe",
            "token": "PERSON:04",
            "canonical_label": "die Gruppe",
            "demask_label": "die Gruppe",
            "entity_type": "PERSON",
            "start": group_start,
            "end": group_start + len("die Gruppe"),
            "source": "request_local",
        }
    ]
    group = mask_prompt(group_prompt, group_rows, "journal_generate")
    expect(len(group.replacements) == 1, "observed multi-word label is active")
    mixed_raw = (
        "Die Gruppe verstand kein Wort. "
        "Später fragte ich, ob die Gruppe noch da sei. "
        "[[PERSON:04]] blieb. "
        "die Gruppen standen abseits."
    )
    tokenized, meta = _process_response(mixed_raw, group)
    expect(meta.get("response_normalization") == "active_cleartext_normalized", "check copy still remaps active cleartext")
    expect(tokenized.startswith("[[PERSON:04]]"), "check copy tokenizes sentence-initial cleartext")
    expect("die Gruppe verstand" in _demask(tokenized, group), "check roundtrip uses the stored label")
    visible = _visible_response(mixed_raw, group)
    expect(visible.startswith("Die Gruppe verstand"), "stored text keeps sentence-initial capitalization")
    expect("ob die Gruppe noch da sei" in visible, "stored text keeps mid-sentence form")
    expect("die Gruppe blieb" in visible, "emitted placeholder is rehydrated locally")
    expect("die Gruppen standen" in visible, "unmatched plural surface form is left unchanged")
    expect(visible != _demask(tokenized, group), "visible output is not the rematerialized check copy")
    leftover_unknown = _visible_response("[[PERSON:04]] und [[PERSON:99]] und [[…]]", group)
    expect(leftover_unknown.startswith("die Gruppe"), "known placeholder is rehydrated")
    expect("[[PERSON:99]]" in leftover_unknown, "unknown placeholder is not materialized")
    expect("[[…]]" in leftover_unknown, "generic placeholder is not materialized")

    overlap = mask_prompt(
        "Anna-Lena kam vorbei.",
        [
            {"local_label": "Anna", "token": "PERSON:01"},
            {"local_label": "Anna-Lena", "token": "PERSON:08"},
        ],
        "journal_generate",
    )
    expect([item.token for item in overlap.replacements] == ["PERSON:08"], "only the actually replaced overlapping label is active")
    expect("[[PERSON:08]]" in overlap.masked_text, "longer label is the one replaced")
    expect("Anna-Lena" not in overlap.masked_text, "replaced longer label is gone")

    mixed_word = "Ich aß Sushi. Sushi kam später."
    second = mixed_word.find("Sushi", mixed_word.find("Sushi") + 1)
    span_only = mask_prompt(
        mixed_word,
        [
            {
                "local_label": "Sushi",
                "token": "PERSON:01",
                "entity_type": "PERSON",
                "start": second,
                "end": second + 5,
                "text": "Sushi",
                "source": "request_local",
            }
        ],
        "journal_generate",
    )
    expect(span_only.masked_text.startswith("Ich aß Sushi."), "undetected same wording stays unmasked")
    expect("[[PERSON:01]] kam später." in span_only.masked_text, "only the detected span is masked")
    validate_pre_egress(span_only.masked_text, span_only)

    overlapping_spans = mask_prompt(
        "Anna-Lena kam vorbei.",
        [
            {
                "local_label": "Anna",
                "token": "PERSON:01",
                "entity_type": "PERSON",
                "start": 0,
                "end": 4,
                "text": "Anna",
                "source": "request_local",
            },
            {
                "local_label": "Anna-Lena",
                "token": "PERSON:08",
                "entity_type": "PERSON",
                "start": 0,
                "end": 9,
                "text": "Anna-Lena",
                "source": "request_local",
            },
        ],
        "journal_generate",
    )
    expect("[[PERSON:08]]" in overlapping_spans.masked_text, "longest overlapping span wins")
    expect("[[PERSON:01]]" not in overlapping_spans.masked_text, "shorter overlapping span is dropped")

    subject = mask_prompt(
        "Clarissa kaufte Kirschen.",
        [{"local_label": "Clarissa", "token": "PERSON:99"}],
        "journal_generate",
    )
    expect([item.token for item in subject.replacements] == ["PERSON:99"], "person as subject of kaufte remains identity")

    leaky = MaskingManifest(
        masked_text="Anna war am Markt.",
        available_mapping_count=1,
        replacements=(ActiveReplacement("PERSON:01", "PERSON", 1, "Anna"),),
    )
    pre = False
    try:
        validate_pre_egress("Anna war am Markt.", leaky)
    except PrivacyGatewayError as exc:
        pre = exc.code == ERROR_EGRESS_VALIDATION
        expect(exc.diagnostics.get("pre_egress_validation") == "failed", "pre-egress failure is named")
    expect(pre, "remaining active identity occurrence fails closed before the provider")
    validate_pre_egress("[[PERSON:01]] war am Markt.", leaky)

    def check(prompt: str, mappings: list[dict], reply: str):
        manifest = mask_prompt(prompt, mappings, "journal_generate")
        try:
            return "ok", _validate_response(reply, manifest), [item.token for item in manifest.replacements]
        except PrivacyGatewayError as exc:
            return exc.code, None, [item.token for item in manifest.replacements]

    with ThreadPoolExecutor(max_workers=2) as pool:
        first = pool.submit(check, "Anna kam vorbei.", anna, "Anna kam.")
        second = pool.submit(check, "Heute nur Markt.", anna + lars, "Anna kam.")
        code_a, _text_a, tokens_a = first.result()
        code_b, text_b, tokens_b = second.result()
    expect(code_a == "ok" and "[[PERSON:01]]" in (_text_a or ""), "parallel request A normalizes its own active cleartext")
    expect(tokens_a == ["PERSON:01"], "parallel request A keeps its own active set")
    expect(code_b == "ok" and "Anna kam" in (text_b or ""), "parallel request B is not blocked by A's mapping")
    expect(tokens_b == [], "parallel request B does not inherit A's active set")

    with TestClient(app) as client:
        setup = client.post(
            "/api/auth/setup",
            json={"email": "lars@example.test", "name": "Lars", "password": "test-pass"},
        )
        profile_id = setup.json()["profile_id"]
        remember_mapping(profile_id, "Anna", "PERSON:01")
        remember_mapping(profile_id, "Clarissa", "PERSON:99")

        result = _run_gateway(profile_id, "Ich war mit Anna am Markt.")
        expect(result.allowed, "masked identity request is allowed")
        expect(result.diagnostics.get("available_mapping_count", 0) >= 1, "diagnostics count request mappings")
        expect(result.diagnostics.get("active_mapping_count") >= 1, "diagnostics count actually masked mappings")
        expect("PERSON:01" in (result.diagnostics.get("active_tokens") or []), "diagnostics name active tokens")
        expect(result.trace.get("pre_egress_validation") == "ok", "pre-egress is recorded")
        expect(result.trace.get("response_validation") == "ok", "response validation is recorded")
        expect(_no_label_keys(result.diagnostics), "compact diagnostics have no local_label keys")
        expect(_no_label_keys(result.trace.get("log")), "run log has no local_label keys")
        expect(_no_label_keys(last_trace()), "last compact trace has no local_label keys")
        compact_blob = json.dumps(result.diagnostics)
        expect("local_label" not in compact_blob, "compact diagnostics omit the mapping field name")
        expect("Anna" not in compact_blob, "compact diagnostics omit the clear name")

        provider_calls = {"n": 0}

        def boom(_messages, _policy):
            provider_calls["n"] += 1
            raise AssertionError("provider must not be called after pre-egress failure")

        def leaky_mask(rendered, mappings, purpose):
            return MaskingManifest(
                masked_text=rendered,
                available_mapping_count=len(mappings or []),
                replacements=(ActiveReplacement("PERSON:01", "PERSON", 1, "Anna"),),
            )

        blocked_egress = False
        with patch("privacy_gateway.mask_prompt", leaky_mask), patch("privacy_gateway.complete_model", boom):
            try:
                _run_gateway(profile_id, "Anna war am Markt.")
            except PrivacyGatewayError as exc:
                blocked_egress = exc.code == ERROR_EGRESS_VALIDATION
        expect(blocked_egress, "pre-egress failure aborts the request")
        expect(provider_calls["n"] == 0, "provider is not called when an active identity remains")

        seen = []

        def leak_once(messages, _policy):
            seen.append(messages[0].get("content") or "")
            return ChatResult(
                content="Anna stand den ganzen Nachmittag am Markt.",
                model="fake",
                usage={"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18, "cost": 0.01},
                context_compression="not_applicable",
            )

        with patch("privacy_gateway.complete_model", leak_once):
            repaired = _run_gateway(profile_id, "Ich war mit Anna am Markt.")
        expect(len(seen) == 1, "active cleartext does not start a second generate call")
        expect("[[PERSON:01]]" in seen[0] and "Anna" not in seen[0].replace("[[PERSON:01]]", ""), "egress stays masked")
        expect("Anna" in (repaired.content or ""), "model plaintext identity spelling is kept")
        expect(repaired.diagnostics.get("generate_calls") == 1, "exactly one generate call")
        expect(
            repaired.diagnostics.get("response_normalization") == "active_cleartext_normalized",
            "active cleartext is recorded as local normalization",
        )
        expect("Anna" not in json.dumps(repaired.diagnostics), "compact diagnostics omit the clear name")

        def spaced_token(_messages, _policy):
            return ChatResult(
                content="[[ person:01 ]] stand den ganzen Nachmittag am Markt.",
                model="fake",
                usage={},
                context_compression="not_applicable",
            )

        with patch("privacy_gateway.complete_model", spaced_token):
            restored = _run_gateway(profile_id, "Ich war mit Anna am Markt.")
        expect(
            "Anna stand den ganzen Nachmittag" in (restored.content or ""),
            "gateway demasks case and spacing variants",
        )
        expect("[[" not in (restored.content or ""), "no leftover placeholder after a successful demask")
        expect(repaired.diagnostics.get("prompt_tokens") == 10, "single call keeps prompt tokens")
        expect(repaired.diagnostics.get("completion_tokens") == 8, "single call keeps completion tokens")
        expect(repaired.diagnostics.get("total_tokens") == 18, "single call keeps total tokens")
        expect(abs(float(repaired.diagnostics.get("cost") or 0) - 0.01) < 1e-9, "single call keeps cost")

        prompt_group = "Ich traf die Gruppe."
        group_at = prompt_group.index("die Gruppe")
        install_test_spans(
            [
                {
                    "start": group_at,
                    "end": group_at + len("die Gruppe"),
                    "text": "die Gruppe",
                    "entity_type": "PERSON",
                }
            ]
        )

        def mixed_surface(messages, _policy):
            egress = messages[0].get("content") or ""
            match = re.search(r"\[\[PERSON:\d+\]\]", egress)
            token = match.group(0) if match else "[[PERSON:01]]"
            return ChatResult(
                content=f"Die Gruppe verstand kein Wort. {token} blieb. Später die Gruppe wieder.",
                model="fake",
                usage={},
                context_compression="not_applicable",
            )

        try:
            with patch("privacy_gateway.complete_model", mixed_surface):
                surfaced = _run_gateway(profile_id, prompt_group)
        finally:
            reset_detect_test_hooks()
        expect(surfaced.diagnostics.get("generate_calls") == 1, "mixed surface still uses one generate call")
        expect("Die Gruppe verstand" in ((surfaced.trace or {}).get("raw") or ""), "trace raw keeps model casing")
        expect("[[" in ((surfaced.trace or {}).get("raw") or ""), "trace raw keeps the emitted placeholder")
        expect(surfaced.content == (surfaced.trace or {}).get("reply"), "visible content matches trace reply")
        expect((surfaced.content or "").startswith("Die Gruppe verstand"), "gateway keeps sentence-initial model casing")
        expect("die Gruppe blieb" in (surfaced.content or ""), "gateway rehydrates the emitted placeholder")
        expect("Später die Gruppe wieder" in (surfaced.content or ""), "gateway keeps later plaintext unchanged")
        expect(
            surfaced.diagnostics.get("response_normalization") == "active_cleartext_normalized",
            "check copy still records cleartext normalization",
        )

        inactive_calls = {"n": 0}

        def keep_inactive(messages, _policy):
            inactive_calls["n"] += 1
            return ChatResult(
                content="Clarissa blieb heute unerwähnt am Hafen.",
                model="fake",
                usage={},
                context_compression="not_applicable",
            )

        with patch("privacy_gateway.complete_model", keep_inactive):
            kept_inactive = _run_gateway(profile_id, "Heute nur Markt und Kirschen am Hafen.")
        expect(inactive_calls["n"] == 1, "inactive mapping does not start a retry")
        expect(kept_inactive.trace.get("response_validation_retry") in (None, 0), "inactive mapping does not mark a retry")
        expect("Clarissa blieb heute unerwähnt" in (kept_inactive.content or ""), "inactive mapping keeps the model text")

        retry_fail = []

        def always_leak(messages, _policy):
            retry_fail.append(messages[0].get("content") or "")
            return ChatResult(
                content="Anna stand den ganzen Nachmittag am Markt.",
                model="fake",
                usage={"prompt_tokens": 4, "completion_tokens": 6, "total_tokens": 10, "cost": 0.002},
                context_compression="not_applicable",
            )

        with patch("privacy_gateway.complete_model", always_leak):
            accepted = _run_gateway(profile_id, "Ich war mit Anna am Markt.")
        expect(len(retry_fail) == 1, "cleartext in the reply does not start a second model call")
        expect("Anna" in (accepted.content or ""), "source-attested name is demasked after local normalization")
        expect(accepted.diagnostics.get("generate_calls") == 1, "still exactly one generate call")
        expect(accepted.trace.get("model") == "fake" or accepted.diagnostics.get("model") == "fake", "error-free trace keeps the model")
        expect(accepted.diagnostics.get("completion_tokens") == 6, "token usage of the only call remains")

    print("All privacy manifest tests passed.")


if __name__ == "__main__":
    main()