"""
Obsideo Cloud Drive client — core crypto + API functions.
CLI commands (cmd_*, main) are intentionally omitted; use server.py for MCP.

Reads config from:
  OBSIDEO_API_KEY              — required
  OBSIDEO_ENCRYPTION_KEY       — optional hex key (falls back to ~/.config/obsideo/key)
  OBSIDEO_WALLET_MNEMONIC      — optional BIP39 mnemonic (falls back to ~/.config/obsideo/jackal-mnemonic)
"""

import base64
import hashlib
import hmac
import json
import os
import pathlib
import struct
import sys
import threading
import time
import urllib.error
import urllib.request

BASE_URL = "https://web-production-5cce7.up.railway.app"

_KEY_FILE    = pathlib.Path.home() / ".config" / "obsideo" / "key"
_WALLET_FILE = pathlib.Path.home() / ".config" / "obsideo" / "jackal-mnemonic"
_CACHE_FILE  = pathlib.Path.home() / ".config" / "obsideo" / "cache.json"


# ── Local cache ───────────────────────────────────────────────────────────────

class LocalCache:
    """
    Thread-safe write-through cache backed by a local JSON file.

    Saves write here first (instant) then sync to Jackal in the background.
    Loads hit here first; fall back to Jackal only on a cache miss.
    Structure: { key: { content: str, synced: bool, failed: bool, ts: float } }
    """

    def __init__(self):
        self._lock = threading.Lock()
        self._data: dict = {}
        self._load()

    def _load(self) -> None:
        if _CACHE_FILE.exists():
            try:
                self._data = json.loads(_CACHE_FILE.read_text())
            except Exception:
                self._data = {}

    def _flush(self) -> None:
        _CACHE_FILE.parent.mkdir(parents=True, exist_ok=True)
        _CACHE_FILE.write_text(json.dumps(self._data, indent=2))

    def get(self, key: str) -> str | None:
        """Return the encrypted content for key, or None on miss."""
        with self._lock:
            entry = self._data.get(key)
            return entry["content"] if entry else None

    def set(self, key: str, encrypted: str) -> None:
        """Store encrypted content locally, marked as pending sync."""
        with self._lock:
            self._data[key] = {
                "content": encrypted,
                "synced":  False,
                "failed":  False,
                "ts":      time.time(),
            }
            self._flush()

    def mark_synced(self, key: str) -> None:
        with self._lock:
            if key in self._data:
                self._data[key]["synced"] = True
                self._data[key]["failed"] = False
                self._flush()

    def mark_failed(self, key: str) -> None:
        with self._lock:
            if key in self._data:
                self._data[key]["failed"] = True
                self._flush()

    def status(self) -> dict:
        with self._lock:
            pending = [k for k, v in self._data.items() if not v["synced"] and not v["failed"]]
            failed  = [k for k, v in self._data.items() if v["failed"]]
            synced  = [k for k, v in self._data.items() if v["synced"]]
            return {"pending": pending, "failed": failed, "synced": synced}


cache = LocalCache()


# ── Encryption ────────────────────────────────────────────────────────────────

def _encryption_key() -> bytes:
    key_hex = os.environ.get("OBSIDEO_ENCRYPTION_KEY", "").strip()
    if key_hex:
        return bytes.fromhex(key_hex)

    if _KEY_FILE.exists():
        return bytes.fromhex(_KEY_FILE.read_text().strip())

    key_hex = os.urandom(32).hex()
    _KEY_FILE.parent.mkdir(parents=True, exist_ok=True)
    _KEY_FILE.write_text(key_hex)
    print(
        "\n[obsideo-cloud-drive] Generated a new encryption key and saved it to:\n"
        f"  {_KEY_FILE}\n\n"
        "Your memories are encrypted with this key. Back it up:\n"
        f"  export OBSIDEO_ENCRYPTION_KEY={key_hex}\n",
        file=sys.stderr,
    )
    return bytes.fromhex(key_hex)


def _encrypt(plaintext: str) -> str:
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    key   = _encryption_key()
    nonce = os.urandom(12)
    ct    = AESGCM(key).encrypt(nonce, plaintext.encode(), None)
    return base64.b64encode(nonce + ct).decode()


def _decrypt(ciphertext_b64: str) -> str:
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    key  = _encryption_key()
    data = base64.b64decode(ciphertext_b64)
    nonce, ct = data[:12], data[12:]
    return AESGCM(key).decrypt(nonce, ct, None).decode()


# ── Jackal wallet (BIP39/BIP44 — no extra dependencies) ──────────────────────

_BIP39_WORDS = None  # lazy-loaded


def _load_wordlist() -> list:
    global _BIP39_WORDS
    if _BIP39_WORDS is None:
        try:
            url = "https://raw.githubusercontent.com/trezor/python-mnemonic/master/src/mnemonic/wordlist/english.txt"
            with urllib.request.urlopen(url, timeout=10) as r:
                _BIP39_WORDS = r.read().decode().split()
        except Exception:
            raise RuntimeError("[obsideo-cloud-drive] Could not fetch BIP39 wordlist. Check your internet connection.")
    return _BIP39_WORDS


def _generate_mnemonic() -> str:
    """Generate a 24-word BIP39 mnemonic from 256 bits of entropy."""
    words = _load_wordlist()
    entropy = os.urandom(32)
    h = hashlib.sha256(entropy).digest()
    bits = bin(int.from_bytes(entropy, 'big'))[2:].zfill(256) + bin(h[0])[2:].zfill(8)
    return ' '.join(words[int(bits[i*11:(i+1)*11], 2)] for i in range(24))


def _mnemonic_to_seed(mnemonic: str) -> bytes:
    """BIP39: mnemonic → 64-byte seed via PBKDF2-HMAC-SHA512."""
    return hashlib.pbkdf2_hmac(
        'sha512',
        mnemonic.encode('utf-8'),
        b'mnemonic',
        2048,
    )


_SECP256K1_ORDER = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEBAAEDCE6AF48A03BBFD25E8CD0364141


def _bip32_derive(seed: bytes, path: str) -> bytes:
    """Derive a private key from seed using BIP32/BIP44 path (e.g. m/44'/118'/0'/0/0)."""
    I = hmac.new(b'Bitcoin seed', seed, hashlib.sha512).digest()
    key, chain = I[:32], I[32:]

    for part in path.lstrip('m/').split('/'):
        hardened = part.endswith("'")
        index = int(part.rstrip("'")) + (0x80000000 if hardened else 0)
        if hardened:
            data = b'\x00' + key + struct.pack('>I', index)
        else:
            data = _privkey_to_pubkey(key) + struct.pack('>I', index)
        I = hmac.new(chain, data, hashlib.sha512).digest()
        il = int.from_bytes(I[:32], 'big')
        key_int = (il + int.from_bytes(key, 'big')) % _SECP256K1_ORDER
        key = key_int.to_bytes(32, 'big')
        chain = I[32:]

    return key


def _privkey_to_pubkey(privkey: bytes) -> bytes:
    """secp256k1 compressed public key from private key."""
    from cryptography.hazmat.primitives.asymmetric.ec import (
        SECP256K1, derive_private_key,
    )
    from cryptography.hazmat.primitives.serialization import Encoding, PublicFormat
    priv = derive_private_key(int.from_bytes(privkey, 'big'), SECP256K1())
    return priv.public_key().public_bytes(Encoding.X962, PublicFormat.CompressedPoint)


def _pubkey_to_address(pubkey: bytes, hrp: str) -> str:
    """Cosmos address: RIPEMD160(SHA256(pubkey)) → bech32(hrp)."""
    sha = hashlib.sha256(pubkey).digest()
    try:
        rip = hashlib.new('ripemd160', sha).digest()
    except ValueError:
        from cryptography.hazmat.primitives import hashes
        from cryptography.hazmat.primitives.hashes import Hash
        h = Hash(hashes.RIPEMD160())
        h.update(sha)
        rip = h.finalize()
    return _bech32_encode(hrp, rip)


def _bech32_encode(hrp: str, data: bytes) -> str:
    """Standard bech32 encoding (Cosmos compatible)."""
    CHARSET = 'qpzry9x8gf2tvdw0s3jn54khce6mua7l'
    GEN = [0x3b6a57b2, 0x26508e6d, 0x1ea119fa, 0x3d4233dd, 0x2a1462b3]

    def polymod(values):
        chk = 1
        for v in values:
            b = chk >> 25
            chk = (chk & 0x1ffffff) << 5 ^ v
            for i in range(5):
                chk ^= GEN[i] if (b >> i) & 1 else 0
        return chk

    def hrp_expand(s):
        return [ord(x) >> 5 for x in s] + [0] + [ord(x) & 31 for x in s]

    def convertbits(data, frombits, tobits):
        acc = bits = 0
        ret = []
        maxv = (1 << tobits) - 1
        for v in data:
            acc = ((acc << frombits) | v) & 0xffffffff
            bits += frombits
            while bits >= tobits:
                bits -= tobits
                ret.append((acc >> bits) & maxv)
        if bits:
            ret.append((acc << (tobits - bits)) & maxv)
        return ret

    data5 = convertbits(data, 8, 5)
    chk = polymod(hrp_expand(hrp) + data5 + [0] * 6) ^ 1
    checksum = [(chk >> 5 * (5 - i)) & 31 for i in range(6)]
    return hrp + '1' + ''.join(CHARSET[d] for d in data5 + checksum)


def _jackal_mnemonic() -> str | None:
    """Return the stored Jackal wallet mnemonic, or None if not set up yet."""
    env = os.environ.get("OBSIDEO_WALLET_MNEMONIC", "").strip()
    if env:
        return env
    if _WALLET_FILE.exists():
        return _WALLET_FILE.read_text().strip()
    return None


def _mnemonic_to_jackal_address(mnemonic: str) -> str:
    """Derive jkl1... address from BIP39 mnemonic via m/44'/118'/0'/0/0."""
    seed    = _mnemonic_to_seed(mnemonic)
    privkey = _bip32_derive(seed, "m/44'/118'/0'/0/0")
    pubkey  = _privkey_to_pubkey(privkey)
    return _pubkey_to_address(pubkey, "jkl")


# ── API ───────────────────────────────────────────────────────────────────────

def _api_key() -> str:
    key = os.environ.get("OBSIDEO_API_KEY", "")
    if not key:
        raise RuntimeError("OBSIDEO_API_KEY is not set in the MCP server environment")
    return key


def _request(method: str, path: str, body: dict | None = None, _retry: bool = True) -> dict:
    url  = BASE_URL + path
    data = json.dumps(body).encode() if body else None
    req  = urllib.request.Request(
        url,
        data=data,
        method=method,
        headers={
            "Authorization": f"Bearer {_api_key()}",
            "Content-Type": "application/json",
        },
    )
    try:
        with urllib.request.urlopen(req, timeout=180) as resp:
            return json.loads(resp.read())
    except urllib.error.HTTPError as e:
        try:
            detail = json.loads(e.read()).get("detail", e.reason)
        except Exception:
            detail = e.reason
        raise RuntimeError(f"API error {e.code}: {detail}")
    except urllib.error.URLError as e:
        # On a timeout, the server was likely cold-starting on Railway.
        # Retry once — by the time we timed out, the dyno is awake.
        if _retry and "timed out" in str(e.reason).lower():
            return _request(method, path, body, _retry=False)
        raise RuntimeError(f"Network error reaching Obsideo Cloud Drive API: {e.reason}")


def _ensure_wallet_registered() -> None:
    """
    On first save: generate a Jackal wallet if none exists, then register
    the jkl1... address with the API. Idempotent — safe to call every save.
    """
    mnemonic = _jackal_mnemonic()

    if mnemonic is None:
        mnemonic = _generate_mnemonic()
        _WALLET_FILE.parent.mkdir(parents=True, exist_ok=True)
        _WALLET_FILE.write_text(mnemonic)
        address = _mnemonic_to_jackal_address(mnemonic)
        print(
            "\n[obsideo-cloud-drive] Generated your Jackal wallet and saved the mnemonic to:\n"
            f"  {_WALLET_FILE}\n\n"
            f"  Jackal address: {address}\n\n"
            "  This wallet will own your storage on-chain. Back up the mnemonic.\n",
            file=sys.stderr,
        )
    else:
        address = _mnemonic_to_jackal_address(mnemonic)

    try:
        _request("POST", "/register-wallet", {"jackal_address": address})
    except RuntimeError:
        pass  # Registration failure is non-fatal — server falls back to sidecar wallet
