import os
import sys
import tempfile
import unittest

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

from graylog_ops.geo import GeoIP, build_geo_db, country_name, flag, region_name


class GeoIPTests(unittest.TestCase):
    def test_build_and_lookup(self):
        with tempfile.TemporaryDirectory() as directory:
            csv_path = os.path.join(directory, "city.csv")
            with open(csv_path, "w", encoding="utf-8") as fh:
                fh.write("1.0.0.0,1.0.0.255,OC,AU,Queensland,\"South Brisbane\",0,0\n")
                fh.write("1.0.1.0,1.0.3.255,AS,CN,Fujian,Wenquan,0,0\n")
                fh.write("8.8.8.0,8.8.8.255,NA,US,California,\"Mountain View\",0,0\n")
                fh.write("2001:db8::,2001:db8::ffff,NA,US,California,X,0,0\n")
            db_path = os.path.join(directory, "geo.db")
            count = build_geo_db(csv_path, db_path)
            self.assertEqual(count, 3)  # IPv6 range skipped

            geo = GeoIP(db_path)
            try:
                cn = geo.lookup("1.0.2.1")
                self.assertEqual(cn["country_code"], "CN")
                self.assertEqual(cn["country_name"], "中国")
                self.assertEqual(cn["region"], "福建")
                self.assertEqual(geo.lookup("8.8.8.8")["country_name"], "美国")
                self.assertIsNone(geo.lookup("9.9.9.9"))
                self.assertIsNone(geo.lookup("not-an-ip"))
            finally:
                geo.close()

    def test_helpers(self):
        self.assertEqual(flag("CN"), "\U0001F1E8\U0001F1F3")
        self.assertEqual(flag("ZZ"), "")
        self.assertEqual(flag(""), "")
        self.assertEqual(country_name("CN"), "中国")
        self.assertEqual(country_name("US"), "美国")
        self.assertEqual(country_name("XX"), "XX")
        self.assertEqual(region_name("CN", "Guangdong"), "广东")
        self.assertEqual(region_name("US", "California"), "California")


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