#!/usr/bin/env python3
"""
Headless Secret Service D-Bus Daemon
Implements org.freedesktop.secrets with Plain and DH crypto support.
Persists credentials to ~/.local/share/antigravity-keyring.json
Bidirectionally syncs with ~/.gemini/antigravity-cli/antigravity-oauth-token
"""

import os
import sys
import json
import time
import uuid
import hmac
import math
import dbus
import dbus.service
import dbus.mainloop.glib
from hashlib import sha256
from gi.repository import GLib
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend

BUS_NAME = "org.freedesktop.secrets"
SS_PATH = "/org/freedesktop/secrets"
SS_SERVICE_IFACE = "org.freedesktop.Secret.Service"
SS_COLLECTION_IFACE = "org.freedesktop.Secret.Collection"
SS_ITEM_IFACE = "org.freedesktop.Secret.Item"
SS_SESSION_IFACE = "org.freedesktop.Secret.Session"
SS_PROMPT_IFACE = "org.freedesktop.Secret.Prompt"
DBUS_PROP_IFACE = "org.freedesktop.DBus.Properties"

DATA_FILE = os.path.expanduser("~/.local/share/antigravity-keyring.json")
AGY_TOKEN_FILE = os.path.expanduser("~/.gemini/antigravity-cli/antigravity-oauth-token")

# Standard DH 1024-bit prime
DH_PRIME_1024_BYTES = (
    0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xC9, 0x0F, 0xDA, 0xA2, 0x21, 0x68,
    0xC2, 0x34, 0xC4, 0xC6, 0x62, 0x8B, 0x80, 0xDC, 0x1C, 0xD1, 0x29, 0x02, 0x4E, 0x08,
    0x8A, 0x67, 0xCC, 0x74, 0x02, 0x0B, 0xBE, 0xA6, 0x3B, 0x13, 0x9B, 0x22, 0x51, 0x4A,
    0x08, 0x79, 0x8E, 0x34, 0x04, 0xDD, 0xEF, 0x95, 0x19, 0xB3, 0xCD, 0x3A, 0x43, 0x1B,
    0x30, 0x2B, 0x0A, 0x6D, 0xF2, 0x5F, 0x14, 0x37, 0x4F, 0xE1, 0x35, 0x6D, 0x6D, 0x51,
    0xC2, 0x45, 0xE4, 0x85, 0xB5, 0x76, 0x62, 0x5E, 0x7E, 0xC6, 0xF4, 0x4C, 0x42, 0xE9,
    0xA6, 0x37, 0xED, 0x6B, 0x0B, 0xFF, 0x5C, 0xB6, 0xF4, 0x06, 0xB7, 0xED, 0xEE, 0x38,
    0x6B, 0xFB, 0x5A, 0x89, 0x9F, 0xA5, 0xAE, 0x9F, 0x24, 0x11, 0x7C, 0x4B, 0x1F, 0xE6,
    0x49, 0x28, 0x66, 0x51, 0xEC, 0xE6, 0x53, 0x81, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
    0xFF, 0xFF
)
DH_PRIME_1024 = int.from_bytes(DH_PRIME_1024_BYTES, "big")


def int_to_bytes(number: int) -> bytes:
    return number.to_bytes(math.ceil(number.bit_length() / 8), "big")


def sync_from_agy_token_file(data):
    if os.path.exists(AGY_TOKEN_FILE):
        try:
            with open(AGY_TOKEN_FILE, "r") as f:
                token_content = f.read().strip()
            if token_content:
                items = data.setdefault("collections", {}).setdefault("login", {}).setdefault("items", {})
                found = False
                for item_id, idata in items.items():
                    attrs = idata.get("attributes", {})
                    if attrs.get("service") == "antigravity" and attrs.get("username") == "token":
                        if idata.get("secret") != token_content:
                            idata["secret"] = token_content
                            idata["modified"] = int(time.time())
                        found = True
                        break
                if not found:
                    items["antigravity_oauth_token"] = {
                        "label": "Antigravity CLI OAuth Token",
                        "attributes": {
                            "service": "antigravity",
                            "username": "token"
                        },
                        "secret": token_content,
                        "created": int(time.time()),
                        "modified": int(time.time())
                    }
        except Exception as e:
            print(f"Error syncing from agy token file: {e}", file=sys.stderr)
    return data


def sync_to_agy_token_file(secret_str):
    try:
        if secret_str:
            os.makedirs(os.path.dirname(AGY_TOKEN_FILE), exist_ok=True)
            tmp_file = AGY_TOKEN_FILE + ".tmp"
            with open(tmp_file, "w") as f:
                f.write(secret_str)
            os.replace(tmp_file, AGY_TOKEN_FILE)
            os.chmod(AGY_TOKEN_FILE, 0o600)
    except Exception as e:
        print(f"Error syncing to agy token file: {e}", file=sys.stderr)


def load_data():
    data = {"collections": {"login": {"label": "Login", "items": {}}}}
    if os.path.exists(DATA_FILE):
        try:
            with open(DATA_FILE, "r") as f:
                data = json.load(f)
        except Exception as e:
            print(f"Error loading data file: {e}", file=sys.stderr)
    return sync_from_agy_token_file(data)


def save_data(data):
    os.makedirs(os.path.dirname(DATA_FILE), exist_ok=True)
    temp_file = DATA_FILE + ".tmp"
    with open(temp_file, "w") as f:
        json.dump(data, f, indent=2)
    os.replace(temp_file, DATA_FILE)
    os.chmod(DATA_FILE, 0o600)


class Session(dbus.service.Object):
    def __init__(self, bus, service, session_id, algorithm, client_key_bytes=None):
        self.bus = bus
        self.service = service
        self.session_id = session_id
        self.algorithm = algorithm
        self.path = f"{SS_PATH}/session/{session_id}"
        self.aes_key = None
        self.server_public_key_bytes = b""

        if algorithm == "dh-ietf1024-sha256-aes128-cbc-pkcs7" and client_key_bytes:
            my_private_key = int.from_bytes(os.urandom(0x80), "big")
            my_public_key = pow(2, my_private_key, DH_PRIME_1024)
            self.server_public_key_bytes = int_to_bytes(my_public_key)
            self.server_public_key_bytes = b"\x00" * (0x80 - len(self.server_public_key_bytes)) + self.server_public_key_bytes

            client_pub_int = int.from_bytes(bytes(client_key_bytes), "big")
            common_secret_int = pow(client_pub_int, my_private_key, DH_PRIME_1024)
            common_secret = int_to_bytes(common_secret_int)
            common_secret = b"\x00" * (0x80 - len(common_secret)) + common_secret

            salt = b"\x00" * 0x20
            pseudo_random_key = hmac.new(salt, common_secret, sha256).digest()
            output_block = hmac.new(pseudo_random_key, b"\x01", sha256).digest()
            self.aes_key = output_block[:0x10]

        super().__init__(bus, self.path)

    def decrypt_secret(self, secret_struct):
        _, iv, encrypted_bytes, _ = secret_struct
        iv = bytes(iv)
        encrypted_bytes = bytes(encrypted_bytes)
        if not self.aes_key or not iv or not encrypted_bytes:
            return encrypted_bytes
        try:
            cipher = Cipher(algorithms.AES(self.aes_key), modes.CBC(iv), default_backend())
            decryptor = cipher.decryptor()
            padded = decryptor.update(encrypted_bytes) + decryptor.finalize()
            pad_len = padded[-1]
            if 1 <= pad_len <= 16 and padded.endswith(bytes([pad_len]) * pad_len):
                return padded[:-pad_len]
            return padded
        except Exception as e:
            print(f"Decryption error: {e}", file=sys.stderr)
            return encrypted_bytes

    def encrypt_secret(self, raw_bytes):
        if not self.aes_key:
            return (dbus.ObjectPath(self.path), dbus.ByteArray(b""), dbus.ByteArray(raw_bytes), dbus.String("text/plain"))
        try:
            iv = os.urandom(16)
            pad_len = 16 - (len(raw_bytes) % 16)
            padded = raw_bytes + bytes([pad_len]) * pad_len
            cipher = Cipher(algorithms.AES(self.aes_key), modes.CBC(iv), default_backend())
            encryptor = cipher.encryptor()
            encrypted = encryptor.update(padded) + encryptor.finalize()
            return (dbus.ObjectPath(self.path), dbus.ByteArray(iv), dbus.ByteArray(encrypted), dbus.String("text/plain"))
        except Exception as e:
            print(f"Encryption error: {e}", file=sys.stderr)
            return (dbus.ObjectPath(self.path), dbus.ByteArray(b""), dbus.ByteArray(raw_bytes), dbus.String("text/plain"))

    @dbus.service.method(SS_SESSION_IFACE, in_signature="", out_signature="")
    def Close(self):
        self.service.remove_session(self.session_id)
        self.remove_from_connection()


class Item(dbus.service.Object):
    def __init__(self, bus, service, item_id, label, attributes, secret_bytes, created=None, modified=None):
        self.bus = bus
        self.service = service
        self.item_id = item_id
        self.label = label
        self.attributes = attributes
        self.secret_bytes = secret_bytes
        now = int(time.time())
        self.created = created or now
        self.modified = modified or now
        self.path = f"{SS_PATH}/collection/login/{item_id}"
        super().__init__(bus, self.path)

    @dbus.service.method(SS_ITEM_IFACE, in_signature="", out_signature="o")
    def Delete(self):
        self.service.remove_item(self.item_id)
        self.remove_from_connection()
        return dbus.ObjectPath("/")

    @dbus.service.method(SS_ITEM_IFACE, in_signature="o", out_signature="(oayays)")
    def GetSecret(self, session_path):
        self.service.check_token_file_sync()
        session = self.service.get_session(str(session_path))
        if session:
            return session.encrypt_secret(self.secret_bytes)
        return (
            dbus.ObjectPath(session_path),
            dbus.ByteArray(b""),
            dbus.ByteArray(self.secret_bytes),
            dbus.String("text/plain"),
        )

    @dbus.service.method(SS_ITEM_IFACE, in_signature="(oayays)", out_signature="")
    def SetSecret(self, secret_struct):
        session_path, _, _, _ = secret_struct
        session = self.service.get_session(str(session_path))
        if session:
            self.secret_bytes = session.decrypt_secret(secret_struct)
        else:
            _, _, secret_bytes, _ = secret_struct
            self.secret_bytes = bytes(secret_bytes)
        self.modified = int(time.time())
        self.service.save()

    @dbus.service.method(DBUS_PROP_IFACE, in_signature="ss", out_signature="v")
    def Get(self, interface_name, property_name):
        return self.GetAll(interface_name).get(property_name, "")

    @dbus.service.method(DBUS_PROP_IFACE, in_signature="s", out_signature="a{sv}")
    def GetAll(self, interface_name):
        if interface_name == SS_ITEM_IFACE:
            return {
                "Locked": dbus.Boolean(False),
                "Attributes": dbus.Dictionary(self.attributes, signature="ss"),
                "Label": dbus.String(self.label),
                "Created": dbus.UInt64(self.created),
                "Modified": dbus.UInt64(self.modified),
            }
        return {}

    @dbus.service.method(DBUS_PROP_IFACE, in_signature="ssv", out_signature="")
    def Set(self, interface_name, property_name, value):
        if interface_name == SS_ITEM_IFACE:
            if property_name == "Label":
                self.label = str(value)
                self.modified = int(time.time())
                self.service.save()
            elif property_name == "Attributes":
                self.attributes = {str(k): str(v) for k, v in value.items()}
                self.modified = int(time.time())
                self.service.save()


class Collection(dbus.service.Object):
    def __init__(self, bus, service, path, label="Login"):
        self.bus = bus
        self.service = service
        self.path = path
        self.label = label
        super().__init__(bus, self.path)

    @dbus.service.method(SS_COLLECTION_IFACE, in_signature="", out_signature="o")
    def Delete(self):
        return dbus.ObjectPath("/")

    @dbus.service.method(SS_COLLECTION_IFACE, in_signature="a{ss}", out_signature="ao")
    def SearchItems(self, attributes):
        self.service.check_token_file_sync()
        results = []
        for item in self.service.items.values():
            match = True
            for k, v in attributes.items():
                if item.attributes.get(str(k)) != str(v):
                    match = False
                    break
            if match:
                results.append(dbus.ObjectPath(item.path))
        return results

    @dbus.service.method(SS_COLLECTION_IFACE, in_signature="a{sv}(oayays)b", out_signature="oo")
    def CreateItem(self, properties, secret_struct, replace):
        label = ""
        attrs = {}
        for k, v in properties.items():
            if k.endswith(".Label") or k == "Label":
                label = str(v)
            elif k.endswith(".Attributes") or k == "Attributes":
                attrs = {str(attr_k): str(attr_v) for attr_k, attr_v in v.items()}

        session_path, _, _, _ = secret_struct
        session = self.service.get_session(str(session_path))
        if session:
            secret_bytes = session.decrypt_secret(secret_struct)
        else:
            _, _, secret_b, _ = secret_struct
            secret_bytes = bytes(secret_b)

        existing_id = None
        if replace and attrs:
            for item_id, item in list(self.service.items.items()):
                if item.attributes == attrs:
                    existing_id = item_id
                    break

        if existing_id:
            item = self.service.items[existing_id]
            item.label = label or item.label
            item.secret_bytes = bytes(secret_bytes)
            item.modified = int(time.time())
        else:
            item_id = "i" + uuid.uuid4().hex[:12]
            item = Item(self.bus, self.service, item_id, label, attrs, bytes(secret_bytes))
            self.service.items[item_id] = item

        self.service.save()
        return (dbus.ObjectPath(item.path), dbus.ObjectPath("/"))

    @dbus.service.method(DBUS_PROP_IFACE, in_signature="ss", out_signature="v")
    def Get(self, interface_name, property_name):
        return self.GetAll(interface_name).get(property_name, "")

    @dbus.service.method(DBUS_PROP_IFACE, in_signature="s", out_signature="a{sv}")
    def GetAll(self, interface_name):
        if interface_name == SS_COLLECTION_IFACE:
            item_paths = [dbus.ObjectPath(item.path) for item in self.service.items.values()]
            return {
                "Items": dbus.Array(item_paths, signature="o"),
                "Label": dbus.String(self.label),
                "Locked": dbus.Boolean(False),
                "Created": dbus.UInt64(0),
                "Modified": dbus.UInt64(0),
            }
        return {}


class SecretService(dbus.service.Object):
    def __init__(self, bus):
        self.bus = bus
        self.sessions = {}
        self.items = {}
        self.collections = {}
        super().__init__(bus, SS_PATH)

        data = load_data()
        items_data = data.get("collections", {}).get("login", {}).get("items", {})
        for item_id, idata in items_data.items():
            secret_bytes = idata.get("secret", "").encode("utf-8")
            item = Item(
                bus,
                self,
                item_id,
                idata.get("label", ""),
                idata.get("attributes", {}),
                secret_bytes,
                idata.get("created"),
                idata.get("modified"),
            )
            self.items[item_id] = item

        # Register collections at all standard paths
        paths = [
            (f"{SS_PATH}/collection/login", "Login"),
            (f"{SS_PATH}/collection/default", "Default"),
            (f"{SS_PATH}/aliases/default", "Default"),
            (f"{SS_PATH}/aliases/login", "Login"),
        ]
        for p, lbl in paths:
            self.collections[p] = Collection(bus, self, p, lbl)

        self.bus_name = dbus.service.BusName(BUS_NAME, bus)

    def check_token_file_sync(self):
        if os.path.exists(AGY_TOKEN_FILE):
            try:
                with open(AGY_TOKEN_FILE, "r") as f:
                    content = f.read().strip()
                if content:
                    found = False
                    for item in self.items.values():
                        if item.attributes.get("service") == "antigravity" and item.attributes.get("username") == "token":
                            if item.secret_bytes.decode("utf-8", errors="replace") != content:
                                item.secret_bytes = content.encode("utf-8")
                                item.modified = int(time.time())
                            found = True
                            break
                    if not found:
                        item_id = "antigravity_oauth_token"
                        item = Item(
                            self.bus,
                            self,
                            item_id,
                            "Antigravity CLI OAuth Token",
                            {"service": "antigravity", "username": "token"},
                            content.encode("utf-8"),
                        )
                        self.items[item_id] = item
                        self.save()
            except Exception as e:
                print(f"Error checking token file sync: {e}", file=sys.stderr)

    def get_session(self, session_path):
        for session in self.sessions.values():
            if session.path == session_path:
                return session
        return None

    def remove_session(self, session_id):
        if session_id in self.sessions:
            del self.sessions[session_id]

    def remove_item(self, item_id):
        if item_id in self.items:
            del self.items[item_id]
            self.save()

    def save(self):
        items_dict = {}
        for item_id, item in self.items.items():
            try:
                secret_str = item.secret_bytes.decode("utf-8", errors="replace")
            except Exception:
                secret_str = ""
            items_dict[item_id] = {
                "label": item.label,
                "attributes": item.attributes,
                "secret": secret_str,
                "created": item.created,
                "modified": item.modified,
            }
            if item.attributes.get("service") == "antigravity" and item.attributes.get("username") == "token":
                sync_to_agy_token_file(secret_str)

        data = {"collections": {"login": {"label": "Login", "items": items_dict}}}
        save_data(data)

    @dbus.service.method(SS_SERVICE_IFACE, in_signature="sv", out_signature="vo")
    def OpenSession(self, algorithm, input_var):
        session_id = "s" + uuid.uuid4().hex[:12]
        client_key_bytes = None
        if algorithm == "dh-ietf1024-sha256-aes128-cbc-pkcs7":
            client_key_bytes = bytes(input_var)

        session = Session(self.bus, self, session_id, algorithm, client_key_bytes)
        self.sessions[session_id] = session

        if algorithm == "dh-ietf1024-sha256-aes128-cbc-pkcs7":
            return (dbus.ByteArray(session.server_public_key_bytes), dbus.ObjectPath(session.path))
        return (dbus.String("", variant_level=1), dbus.ObjectPath(session.path))

    @dbus.service.method(SS_SERVICE_IFACE, in_signature="a{sv}s", out_signature="oo")
    def CreateCollection(self, properties, alias):
        col_path = f"{SS_PATH}/collection/login"
        return (dbus.ObjectPath(col_path), dbus.ObjectPath("/"))

    @dbus.service.method(SS_SERVICE_IFACE, in_signature="a{ss}", out_signature="aoao")
    def SearchItems(self, attributes):
        self.check_token_file_sync()
        results = []
        for item in self.items.values():
            match = True
            for k, v in attributes.items():
                if item.attributes.get(str(k)) != str(v):
                    match = False
                    break
            if match:
                results.append(dbus.ObjectPath(item.path))
        return (dbus.Array(results, signature="o"), dbus.Array([], signature="o"))

    @dbus.service.method(SS_SERVICE_IFACE, in_signature="ao", out_signature="aoo")
    def Unlock(self, objects):
        return (objects, dbus.ObjectPath("/"))

    @dbus.service.method(SS_SERVICE_IFACE, in_signature="ao", out_signature="aoo")
    def Lock(self, objects):
        return (dbus.Array([], signature="o"), dbus.ObjectPath("/"))

    @dbus.service.method(SS_SERVICE_IFACE, in_signature="aoo", out_signature="a{o(oayays)}")
    def GetSecrets(self, item_paths, session_path):
        self.check_token_file_sync()
        results = {}
        session = self.get_session(str(session_path))
        for item in self.items.values():
            if dbus.ObjectPath(item.path) in item_paths:
                if session:
                    results[dbus.ObjectPath(item.path)] = session.encrypt_secret(item.secret_bytes)
                else:
                    results[dbus.ObjectPath(item.path)] = (
                        dbus.ObjectPath(session_path),
                        dbus.ByteArray(b""),
                        dbus.ByteArray(item.secret_bytes),
                        dbus.String("text/plain"),
                    )
        return results

    @dbus.service.method(SS_SERVICE_IFACE, in_signature="s", out_signature="o")
    def ReadAlias(self, name):
        return dbus.ObjectPath(f"{SS_PATH}/collection/login")

    @dbus.service.method(SS_SERVICE_IFACE, in_signature="so", out_signature="")
    def SetAlias(self, name, collection_path):
        pass

    @dbus.service.method(DBUS_PROP_IFACE, in_signature="ss", out_signature="v")
    def Get(self, interface_name, property_name):
        return self.GetAll(interface_name).get(property_name, "")

    @dbus.service.method(DBUS_PROP_IFACE, in_signature="s", out_signature="a{sv}")
    def GetAll(self, interface_name):
        if interface_name == SS_SERVICE_IFACE:
            col_paths = [dbus.ObjectPath(f"{SS_PATH}/collection/login")]
            return {"Collections": dbus.Array(col_paths, signature="o")}
        return {}


def main():
    dbus.mainloop.glib.DBusGMainLoop(set_as_default=True)
    bus = dbus.SessionBus()
    service = SecretService(bus)
    print("Secret Service Daemon started successfully.", flush=True)
    loop = GLib.MainLoop()
    try:
        loop.run()
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()