import json
import os
import sys
import tempfile
import time
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest import mock

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

os.environ["LABS_FALLBACK_SIGNING_KEY"] = "unit-signing-key"
os.environ.pop("SPACE_ID", None)

import access_limits as al  # noqa: E402
import quota_fallback as qf  # noqa: E402

PARENT = qf.EMBED_PARENTS[0]
SPACE_HOST = "broyang-3dai.hf.space"


def request(referer="", forwarded="", host=SPACE_HOST, client="10.0.0.9"):
    headers = {"host": host}
    if referer:
        headers["referer"] = referer
    if forwarded:
        headers["x-forwarded-for"] = forwarded
    return SimpleNamespace(headers=headers, client=SimpleNamespace(host=client))


def embed_request(**kwargs):
    return request(referer=f"https://{SPACE_HOST}/?__theme=light&__labs_parent={PARENT}", **kwargs)


def public_request(forwarded="203.0.113.7"):
    return request(referer="https://3dai.skyebrowse.com/", forwarded=forwarded)


def opened(ticket):
    return json.loads(qf._b64url_decode(ticket.split(".")[0]))


def grant_for(nonce, mode="visitor", **overrides):
    payload = {"v": 1, "purpose": "grant", "space": qf.space_identifier(), "nonce": nonce,
               "mode": mode, "exp": int(time.time()) + 180}
    payload.update(overrides)
    return qf.sign_document(payload)


class GrantBookTests(unittest.TestCase):
    def setUp(self):
        self.book = al.GrantBook()

    def issue(self, session="session-a"):
        receipt = self.book.request(session)
        return receipt, opened(receipt["ticket"])["nonce"]

    def test_grant_requests_are_signed_grant_operations_that_cannot_unlock_the_fallback(self):
        receipt, nonce = self.issue()
        claim = opened(receipt["ticket"])
        self.assertEqual((claim["operation"], claim["space"]), (qf.GRANT, qf.space_identifier()))
        self.assertRegex(nonce, r"^[0-9a-f]{32}$")
        self.assertIsNone(qf.verify_receipt(receipt["ticket"]))

    def test_verified_grants_run_natively_once(self):
        receipt, nonce = self.issue()
        self.book.settle("session-a", {"ticket": receipt["ticket"], "grant": grant_for(nonce)}, receipt)
        access = self.book.take("session-a")
        self.assertEqual((access.mode, access.nonce, access.notice), ("native", nonce, None))
        self.assertEqual(self.book.take("session-a").mode, "denied")

    def test_forged_mismatched_or_missing_grants_deny_the_generation(self):
        receipt, nonce = self.issue()
        forged = [
            None,
            "not-a-grant",
            grant_for("b" * 32),
            grant_for(nonce, purpose="execute"),
            grant_for(nonce, mode="unlimited"),
            grant_for(nonce, mode="shared"),
            grant_for(nonce, space="broyang/other"),
            grant_for(nonce, exp=int(time.time()) - 1),
        ]
        for grant in forged:
            receipt, nonce = self.issue()
            self.book.settle("session-a", {"ticket": receipt["ticket"], "grant": grant}, receipt)
            self.assertEqual(self.book.take("session-a").mode, "denied")
        receipt, nonce = self.issue()
        self.book.settle("session-b", {"ticket": receipt["ticket"], "grant": grant_for(nonce)}, receipt)
        self.assertEqual(self.book.take("session-a").mode, "denied")
        self.assertEqual(self.book.take("session-b").mode, "denied")

    def test_limit_refusal_shows_the_upgrade_message(self):
        receipt, _nonce = self.issue()
        message = "You have used your 3 free 3D AI generations. Upgrade to a Commercial plan to keep generating."
        self.book.settle("session-a", {"ticket": receipt["ticket"], "reason": al.FREE_LIMIT_REASON,
                                       "error": message}, receipt)
        access = self.book.take("session-a")
        self.assertEqual(access.mode, "denied")
        self.assertIn("3 free 3D AI generations", access.notice)

    def test_stale_grant_requests_expire(self):
        receipt, nonce = self.issue()
        with mock.patch.object(al.time, "time", return_value=time.time() + al.GRANT_TTL_SECONDS + 1):
            self.book.settle("session-a", {"ticket": receipt["ticket"], "grant": grant_for(nonce)}, receipt)
            self.assertEqual(self.book.take("session-a").mode, "denied")

    def test_grant_requests_are_bounded_per_session_and_in_total(self):
        for _ in range(al.MAX_GRANTS_PER_SESSION):
            self.assertIsNotNone(self.book.request("session-a"))
        self.assertIsNone(self.book.request("session-a"))
        self.assertIsNotNone(self.book.request("session-b"))
        book = al.GrantBook(max_entries=1)
        self.assertIsNotNone(book.request("one"))
        self.assertIsNone(book.request("two"))

    def settle(self, receipt, nonce, session="session-a"):
        self.book.settle(session, {"ticket": receipt["ticket"], "grant": grant_for(nonce)}, receipt)

    def test_a_settled_grant_survives_a_long_gpu_queue(self):
        receipt, nonce = self.issue()
        self.settle(receipt, nonce)
        with mock.patch.object(al.time, "time", return_value=time.time() + al.GRANT_TTL_SECONDS + 60):
            access = self.book.take("session-a")
        self.assertEqual((access.mode, access.nonce), ("native", nonce))

    def test_a_lapsed_grant_is_denied_with_its_nonce_for_a_refund(self):
        receipt, nonce = self.issue()
        self.settle(receipt, nonce)
        with mock.patch.object(al.time, "time", return_value=time.time() + al.SETTLED_GRANT_TTL_SECONDS + 1):
            access = self.book.take("session-a")
        self.assertEqual((access.mode, access.nonce), ("denied", nonce))
        self.assertIsNone(self.book.take("session-a").nonce)

    def test_overlapping_runs_in_one_session_each_keep_their_grant(self):
        first, first_nonce = self.issue()
        second, second_nonce = self.issue()
        self.settle(first, first_nonce)
        self.settle(second, second_nonce)
        self.assertEqual(self.book.take("session-a").nonce, first_nonce)
        self.assertEqual(self.book.take("session-a").nonce, second_nonce)
        self.assertEqual(self.book.take("session-a").mode, "denied")

    def test_overlapping_runs_keep_their_own_refusals(self):
        first, first_nonce = self.issue()
        second, _nonce = self.issue()
        self.settle(first, first_nonce)
        self.book.settle("session-a", {"ticket": second["ticket"], "reason": al.FREE_LIMIT_REASON, "error": ""}, second)
        self.assertEqual(self.book.take("session-a").nonce, first_nonce)
        self.assertIn(al.EMBED_LIMIT_MESSAGE, self.book.take("session-a").notice)

    def test_an_ack_for_an_unknown_nonce_marks_no_other_grant(self):
        receipt, nonce = self.issue()
        stranger = {"ticket": qf.issue_ticket("e" * 32, qf.GRANT)}
        self.book.settle("session-a", {"ticket": stranger["ticket"], "grant": grant_for("e" * 32)}, stranger)
        self.settle(receipt, nonce)
        self.assertEqual(self.book.take("session-a").nonce, nonce)


class PublicAllowanceTests(unittest.TestCase):
    def setUp(self):
        self.directory = tempfile.TemporaryDirectory()
        self.addCleanup(self.directory.cleanup)
        self.path = Path(self.directory.name) / "usage.json"

    def allowance(self, limit=3, window=3600):
        return al.PublicAllowance(limit, window, self.path)

    def test_three_uses_per_visitor_then_refusal_and_refunds(self):
        allowance = self.allowance()
        key = allowance.key_for(public_request())
        self.assertEqual([allowance.try_consume(key) for _ in range(4)], [True, True, True, False])
        allowance.refund(key)
        self.assertTrue(allowance.try_consume(key))
        other = allowance.key_for(public_request("198.51.100.4"))
        self.assertNotEqual(key, other)
        self.assertTrue(allowance.try_consume(other))

    def test_counts_survive_a_reload_and_never_store_the_raw_address(self):
        allowance = self.allowance()
        key = allowance.key_for(public_request())
        for _ in range(3):
            allowance.try_consume(key)
        self.assertNotIn("203.0.113.7", self.path.read_text())
        self.assertFalse(self.allowance().try_consume(key))

    def test_uses_age_out_of_the_window(self):
        allowance = self.allowance(window=60)
        key = allowance.key_for(public_request())
        for _ in range(3):
            allowance.try_consume(key)
        with mock.patch.object(al.time, "time", return_value=time.time() + 61):
            self.assertTrue(allowance.try_consume(key))

    def test_first_forwarded_address_identifies_the_visitor(self):
        allowance = self.allowance()
        self.assertEqual(allowance.key_for(public_request("203.0.113.7, 10.1.1.1")),
                         allowance.key_for(public_request("203.0.113.7")))

    def test_a_missing_forwarded_address_never_shares_one_bucket(self):
        allowance = self.allowance()
        with self.assertLogs(al.logger, "WARNING"):
            key = allowance.key_for(public_request(""))
        self.assertIsNone(key)
        self.assertTrue(all(allowance.try_consume(key) for _ in range(5)))


class AccessPolicyTests(unittest.TestCase):
    def setUp(self):
        self.directory = tempfile.TemporaryDirectory()
        self.addCleanup(self.directory.cleanup)
        self.policy = al.AccessPolicy(al.GrantBook(), al.PublicAllowance(3, 3600, Path(self.directory.name) / "u.json"))

    def test_public_browser_visits_get_three_generations_then_a_signup_link(self):
        modes = [self.policy.authorize(public_request(), "s").mode for _ in range(4)]
        self.assertEqual(modes, ["native", "native", "native", "denied"])
        notice = self.policy.authorize(public_request(), "s").notice
        self.assertIn('href="https://app.skyebrowse.com/register"', notice)
        self.assertIn("Sign up free", notice)

    def test_server_side_callers_without_a_browser_referer_are_not_counted(self):
        for _ in range(5):
            self.assertEqual(self.policy.authorize(request(forwarded="203.0.113.7"), "s").mode, "native")
        self.assertEqual(self.policy.public.used(self.policy.public.key_for(public_request())), 0)

    def test_embedded_generations_need_a_grant_only_when_grants_are_enabled(self):
        with mock.patch.dict(os.environ, {"LABS_EMBED_GRANTS": ""}):
            self.assertEqual(self.policy.authorize(embed_request(), "s").mode, "native")
        with mock.patch.dict(os.environ, {"LABS_EMBED_GRANTS": "1"}):
            self.assertEqual(self.policy.authorize(embed_request(), "s").mode, "denied")
            receipt = self.policy.grants.request("s")
            nonce = opened(receipt["ticket"])["nonce"]
            self.policy.grants.settle("s", {"ticket": receipt["ticket"], "grant": grant_for(nonce)}, receipt)
            self.assertEqual(self.policy.authorize(embed_request(), "s").nonce, nonce)
        self.assertEqual(self.policy.public.used(self.policy.public.key_for(embed_request())), 0)

    def test_grants_can_be_enforced_for_listed_parents_only(self):
        with mock.patch.dict(os.environ, {"LABS_EMBED_GRANTS": f"http://localhost:8080, {PARENT}"}):
            self.assertTrue(al.grants_required(embed_request()))
            self.assertFalse(al.grants_required(public_request()))
        with mock.patch.dict(os.environ, {"LABS_EMBED_GRANTS": "http://localhost:8080"}):
            self.assertFalse(al.grants_required(embed_request()))

    def test_quota_denial_refunds_the_public_use_and_points_to_signup(self):
        access = self.policy.authorize(public_request(), "s")
        notice = self.policy.quota_denied(access)
        self.assertIn("app.skyebrowse.com/register", notice)
        self.assertEqual(self.policy.public.used(access.visitor_key), 0)
        self.assertIsNone(self.policy.quota_denied(al.Access()))

    def test_failed_generations_issue_one_signed_refund_and_return_the_public_use(self):
        public = self.policy.authorize(public_request(), "p")
        self.policy.generation_failed("p", public)
        self.assertEqual(self.policy.public.used(public.visitor_key), 0)
        self.assertIsNone(self.policy.refund_receipt("p"))
        self.policy.generation_failed("e", al.Access(nonce="c" * 32))
        refund = self.policy.refund_receipt("e")
        claim = opened(refund["ticket"])
        self.assertEqual((claim["operation"], claim["nonce"]), (qf.REFUND, "c" * 32))
        self.assertIsNone(qf.verify_receipt(refund["ticket"]))
        self.assertIsNone(self.policy.refund_receipt("e"))

    def grant(self, session="e"):
        receipt = self.policy.grants.request(session)
        nonce = opened(receipt["ticket"])["nonce"]
        self.policy.grants.settle(session, {"ticket": receipt["ticket"], "grant": grant_for(nonce)}, receipt)
        return nonce

    def test_overlapping_failed_generations_are_each_refunded(self):
        with mock.patch.dict(os.environ, {"LABS_EMBED_GRANTS": "1"}):
            nonces = [self.grant(), self.grant()]
            for _ in nonces:
                access = self.policy.authorize(embed_request(), "e")
                self.assertEqual(access.mode, "native")
                self.policy.generation_failed("e", access)
        self.assertEqual([opened(self.policy.refund_receipt("e")["ticket"])["nonce"] for _ in nonces], nonces)
        self.assertIsNone(self.policy.refund_receipt("e"))

    def test_a_lapsed_grant_is_refunded_when_its_run_is_denied(self):
        with mock.patch.dict(os.environ, {"LABS_EMBED_GRANTS": "1"}):
            nonce = self.grant()
            with mock.patch.object(al.time, "time", return_value=time.time() + al.SETTLED_GRANT_TTL_SECONDS + 1):
                self.assertEqual(self.policy.authorize(embed_request(), "e").mode, "denied")
        self.assertEqual(opened(self.policy.refund_receipt("e")["ticket"])["nonce"], nonce)

    def test_notices_are_delivered_once(self):
        self.policy.remember("s", "<div>notice</div>")
        self.assertEqual(self.policy.pop_notice("s"), "<div>notice</div>")
        self.assertIsNone(self.policy.pop_notice("s"))


class EmbedParentTests(unittest.TestCase):
    def test_only_localhost_extras_extend_the_parent_allow_list(self):
        extras = "http://localhost:8109, https://evil.example, http://localhost:8109, javascript:alert(1)"
        with mock.patch.dict(os.environ, {"LABS_EXTRA_EMBED_PARENTS": extras}):
            self.assertEqual(qf.embed_parents(), qf.EMBED_PARENTS + ("http://localhost:8109",))
            for name in ("grant_bridge.js", "refund_bridge.js"):
                bridge = al.bridge_script(ROOT / name)
                self.assertIn('"http://localhost:8109"', bridge)
                self.assertNotIn("evil.example", bridge)
        self.assertNotIn("__LABS_EMBED_PARENTS__", al.bridge_script(ROOT / "fallback_bridge.js"))


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