import os
import sys
import tempfile
import time
import unittest

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.dirname(HERE))

from graylog_ops import watch  # noqa: E402

NOW = "2026-10-17T16:00:00Z"


def _behavior(ip, verdict="建议封禁", score=75):
    return {
        "ip": ip, "score": score, "verdict": verdict,
        "signals": ["高失败率"], "reason_text": "24 小时内扫描了大量路径。",
        "share_note": "", "behavior_type": "目录扫描",
        "asn": "4538", "asn_org": "教育网", "shared_type": "education",
        "total": 50, "pct4": 60.0, "uri_card": 20, "uri_templates": 15,
        "ua_card": 1, "host_card": 2, "coverage": 100.0, "peak_share": 50.0,
        "top_uri": "/.env", "top_ua": "curl/8.0",
        "country_code": "CN", "country_name": "中国", "region": "北京",
        "city": "Beijing",
    }


class FusedTests(unittest.TestCase):
    def test_feature_and_behavior_merge(self):
        with tempfile.TemporaryDirectory() as directory:
            store = watch.Store(os.path.join(directory, "ips.db"))
            try:
                store.upsert({"ip": "203.0.113.9", "tier": watch.TIER_MALICIOUS,
                              "score": 91.0, "total": 120, "window": 3600,
                              "rules": {"probe_java": 10},
                              "reasons": ["探测 10 次"]}, NOW)
                store.sync_fused(NOW)
                fused = store.fused_items()
                self.assertEqual(len(fused), 1)
                self.assertEqual(fused[0]["sources"], "feature")
                self.assertEqual(fused[0]["first_banned_at"], NOW)
                self.assertIn("特征命中", fused[0]["fused_reason"])

                # behaviour-only ban is merged
                store.replace_behavior_profiles([_behavior("198.51.100.9")])
                store.sync_fused("2026-10-17T16:10:00Z")
                merged = {item["ip"]: item for item in store.fused_items()}
                self.assertIn("198.51.100.9", merged)
                self.assertEqual(merged["198.51.100.9"]["sources"], "behavior")
                self.assertEqual(merged["203.0.113.9"]["sources"], "feature")

                # same IP on both sides becomes "both"
                store.upsert({"ip": "198.51.100.9", "tier": watch.TIER_MALICIOUS,
                              "score": 80.0, "total": 50, "window": 3600,
                              "rules": {"probe_java": 6},
                              "reasons": ["探测 6 次"]}, "2026-10-17T16:11:00Z")
                store.sync_fused("2026-10-17T16:12:00Z")
                merged = {item["ip"]: item for item in store.fused_items()}
                self.assertEqual(merged["198.51.100.9"]["sources"], "both")

                # watch-only behaviour must not enter the block list
                store.replace_behavior_profiles([
                    _behavior("192.0.2.44", verdict="仅告警（共享出口）")])
                store.sync_fused("2026-10-17T16:20:00Z")
                ips = {item["ip"] for item in store.fused_items()}
                self.assertNotIn("192.0.2.44", ips)
                self.assertIn("203.0.113.9", ips)  # feature rows stay forever
            finally:
                store.close()

    def test_push_throttle_24h(self):
        with tempfile.TemporaryDirectory() as directory:
            store = watch.Store(os.path.join(directory, "ips.db"))
            try:
                now = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
                store.replace_behavior_profiles([_behavior("198.51.100.7")])
                store.sync_fused(now)
                # first sync suppresses historical rows; a fresh detection is due
                with store.lock, store.conn:
                    store.conn.execute(
                        "UPDATE ip_fused SET last_pushed_at=NULL WHERE ip=?",
                        ("198.51.100.7",))
                due = store.due_pushes(active_hours=24, cooldown_hours=24)
                self.assertEqual([row["ip"] for row in due], ["198.51.100.7"])

                store.mark_pushed(["198.51.100.7"], now)
                self.assertEqual(store.due_pushes(24, 24), [])  # throttled

                old = time.strftime("%Y-%m-%dT%H:%M:%SZ",
                                    time.gmtime(time.time() - 25 * 3600))
                with store.lock, store.conn:
                    store.conn.execute(
                        "UPDATE ip_fused SET last_pushed_at=? WHERE ip=?",
                        (old, "198.51.100.7"))
                due = store.due_pushes(active_hours=24, cooldown_hours=24)
                self.assertEqual([row["ip"] for row in due], ["198.51.100.7"])
            finally:
                store.close()


class HistoricalBehaviourTests(unittest.TestCase):
    def test_historical_profile_survives_live_refresh(self):
        with tempfile.TemporaryDirectory() as directory:
            store = watch.Store(os.path.join(directory, "ips.db"))
            try:
                store.upsert_historical_profiles([{
                    "ip": "198.51.100.9", "score": 80,
                    "verdict": "建议封禁", "signals": ["历史补测"],
                    "last_ts": "2026-10-18T10:00:00Z", "total": 1000,
                }])
                store.replace_behavior_profiles([{
                    "ip": "203.0.113.9", "score": 70,
                    "verdict": "仅告警（共享出口）"}])
                rows = {r["ip"]: dict(r) for r in store.conn.execute(
                    "SELECT * FROM behavior_profiles")}
                self.assertEqual(rows["198.51.100.9"]["historical"], 1)
                self.assertEqual(rows["203.0.113.9"]["historical"], 0)

                store.sync_fused()
                fused = {r["ip"]: dict(r) for r in store.conn.execute(
                    "SELECT * FROM ip_fused")}
                self.assertEqual(fused["198.51.100.9"]["first_banned_at"],
                                 "2026-10-18T10:00:00Z")
            finally:
                store.close()


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