import base64
import unittest

from sbp.state import ProtocolState, RealityState, State, TlsState, User, migrate

USER_UUID = "11111111-2222-3333-4444-555555555555"
USER_SS_KEY = base64.b64encode(b"0123456789abcdef").decode()
SERVER_SS_KEY = base64.b64encode(b"abcdef0123456789").decode()
CERT_SHA256 = "ab" * 32


def sample_state() -> State:
  return State(
    schema=2,
    created="2026-01-01T00:00:00Z",
    host="203.0.113.10",
    sni="www.bing.com",
    tls=TlsState(
      mode="self-signed",
      cert="/tmp/cert.pem",
      key="/tmp/key.pem",
      sni="www.bing.com",
      self_signed=True,
      spki_pin="SPKI",
      cert_sha256=CERT_SHA256,
    ),
    reality=RealityState(private_key="PRIVATE", public_key="PUBLIC"),
    users=[User("u1", "PASSWORD", USER_UUID, USER_SS_KEY, "2026-01-01T00:00:00Z")],
    protocols={
      "hysteria2": ProtocolState(
        True, 21001, "::", {"obfs": True, "obfs_password": "OBFS", "masquerade": ""}
      ),
      "vless-reality": ProtocolState(
        True, 21002, "::", {"server_name": "www.bing.com", "short_id": "", "flow": "xtls-rprx-vision"}
      ),
      "trojan": ProtocolState(True, 21003, "::", {}),
      "shadowsocks": ProtocolState(
        True, 21004, "::", {"method": "2022-blake3-aes-128-gcm", "server_key": SERVER_SS_KEY}
      ),
      "tuic": ProtocolState(True, 21005, "::", {"congestion_control": "bbr"}),
      "anytls": ProtocolState(True, 21006, "::", {}),
    },
    singbox_version="v1.14.4",
  )


class StateRoundTripTest(unittest.TestCase):
  def test_round_trip(self) -> None:
    state = sample_state()
    restored = State.from_dict(state.to_dict())
    self.assertEqual(restored.host, state.host)
    self.assertEqual(restored.sni, state.sni)
    self.assertEqual(restored.users[0].uuid, USER_UUID)
    self.assertEqual(restored.protocols["hysteria2"].port, 21001)
    self.assertEqual(restored.reality.public_key, "PUBLIC")

  def test_migrate_schema_one(self) -> None:
    raw = {
      "schema": 1,
      "created": "2025-01-01T00:00:00Z",
      "host": "198.51.100.1",
      "sni": "www.example.com",
      "listen": "::",
      "port": 23456,
      "tls": {"mode": "self-signed", "self_signed": True, "cert": "/c", "key": "/k",
              "sni": "www.example.com", "spki_pin": "S", "cert_sha256": "C"},
      "obfs": {"enabled": True, "password": "O"},
      "users": [{"name": "u1", "password": "P", "created": "2025-01-01T00:00:00Z"}],
      "version": "v1.10.0",
    }
    migrated = migrate(raw)
    state = State.from_dict(migrated)
    self.assertEqual(state.schema, 2)
    self.assertEqual(state.protocols["hysteria2"].port, 23456)
    options = state.protocols["hysteria2"].options
    self.assertIs(options["obfs"], True)
    self.assertEqual(options["obfs_password"], "O")
    self.assertEqual(state.users[0].name, "u1")
    self.assertTrue(state.users[0].uuid)
    self.assertTrue(state.users[0].ss_key)
    self.assertEqual(state.singbox_version, "v1.10.0")

  def test_migrated_state_builds_config_and_uri(self) -> None:
    from sbp.config import build_server_config
    from sbp.share import share_uri

    raw = {
      "schema": 1,
      "host": "198.51.100.1",
      "sni": "www.example.com",
      "listen": "::",
      "port": 23456,
      "tls": {"mode": "self-signed", "self_signed": True, "cert": "/c", "key": "/k",
              "sni": "www.example.com", "spki_pin": "S", "cert_sha256": "C"},
      "obfs": {"enabled": True, "password": "O"},
      "users": [{"name": "u1", "password": "P"}],
    }
    state = State.from_dict(migrate(raw))
    inbound = build_server_config(state)["inbounds"][0]
    self.assertEqual(inbound["obfs"]["password"], "O")
    self.assertIn("obfs-password=O", share_uri(state, "hysteria2", state.users[0]))

  def test_obfs_without_key_is_disabled(self) -> None:
    from sbp.config import build_server_config

    state = sample_state()
    state.protocols["hysteria2"].options = {"obfs": True}
    inbound = next(
      i for i in build_server_config(state)["inbounds"] if i["tag"] == "hysteria2-in"
    )
    self.assertNotIn("obfs", inbound)


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