"""
ECP Identity — Agent DID generation and signing.

DID format: did:ecp:{sha256(pubkey)[:32]}
  - sha256(pubkey) = 64 hex chars
  - [:32] = first 16 bytes (32 hex chars)

Signature format: ed25519:{hex}
"""

import hashlib
import json
import time
from pathlib import Path

try:
    from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
    from cryptography.hazmat.primitives.serialization import (
        Encoding, PublicFormat, PrivateFormat, NoEncryption
    )
    HAS_CRYPTO = True
except ImportError:
    HAS_CRYPTO = False


import os as _os
ECP_DIR = Path(_os.environ.get("ATLAST_ECP_DIR", _os.environ.get("ECP_DIR", _os.path.expanduser("~/.ecp"))))
IDENTITY_FILE = ECP_DIR / "identity.json"


def _now_ms() -> int:
    """Current Unix timestamp in milliseconds."""
    return int(time.time() * 1000)


def _resolve_ecp_dir() -> Path:
    """Resolve ECP_DIR at call time (not import time) for env var changes."""
    return Path(_os.environ.get("ATLAST_ECP_DIR", _os.environ.get("ECP_DIR", _os.path.expanduser("~/.ecp"))))


def get_or_create_identity(agent_name: str | None = None, ecp_dir: str | None = None) -> dict:
    """Load existing identity or generate a new one."""
    edir = Path(ecp_dir) if ecp_dir else _resolve_ecp_dir()
    ifile = edir / "identity.json"
    edir.mkdir(exist_ok=True)
    if ifile.exists():
        try:
            import json as _json
            identity = _json.loads(ifile.read_text())
            # Auto-migrate: if crypto is now available but identity was
            # generated with fallback (pub = sha256(priv) instead of Ed25519)
            identity = _maybe_migrate_identity(identity, ifile)
            return identity
        except Exception:
            pass
    return _create_identity(edir)


def _maybe_migrate_identity(identity: dict, ifile: Path) -> dict:
    """
    Detect identities created without cryptography package (fallback mode).
    If pub_key == sha256(priv_key_bytes), it's a fallback identity.

    Migration strategy (SAFE):
    - Generate a NEW Ed25519 keypair for signing
    - Store as crypto_pub_key / crypto_priv_key (for signing only)
    - PRESERVE the original DID (did:ecp:...) — it may be registered on LLaChat backend
    - sign() uses crypto_priv_key; verify uses crypto_pub_key
    - DID stays the same → no backend breakage
    """
    if not HAS_CRYPTO or not identity.get("priv_key"):
        return identity

    # Already migrated?
    if identity.get("crypto_pub_key"):
        return identity

    priv_hex = identity["priv_key"]
    stored_pub = identity.get("pub_key", "")

    # Check if pub was generated by fallback: sha256(priv_bytes)
    fallback_pub = hashlib.sha256(bytes.fromhex(priv_hex)).hexdigest()
    if stored_pub != fallback_pub:
        return identity  # Already a proper Ed25519 identity — no migration needed

    # Migration: generate signing keypair, but KEEP the original DID
    private_key = Ed25519PrivateKey.generate()
    pub_bytes = private_key.public_key().public_bytes(Encoding.Raw, PublicFormat.Raw)
    priv_bytes = private_key.private_bytes(Encoding.Raw, PrivateFormat.Raw, NoEncryption())

    identity["crypto_pub_key"] = pub_bytes.hex()
    identity["crypto_priv_key"] = priv_bytes.hex()
    # DID stays the same! No breakage for registered agents.
    identity["migrated_from_fallback"] = True
    identity["migrated_at"] = _now_ms()

    ifile.write_text(json.dumps(identity, indent=2))
    return identity


def _create_identity(edir: Path | None = None) -> dict:
    """
    Create a new agent identity with BIP39 recovery support.
    
    New flow (with cryptography):
      1. Generate 16 bytes random entropy
      2. Entropy → BIP39 mnemonic (12 words)
      3. Entropy → HKDF → Ed25519 seed → keypair
      4. Save identity.json (WITHOUT mnemonic — security)
      5. Return identity dict WITH transient '_mnemonic' key (for display only)
    
    Legacy flow (without cryptography):
      Random hex fallback, no mnemonic support.
    """
    edir = Path(edir) if edir else _resolve_ecp_dir()
    ifile = edir / "identity.json"
    edir.mkdir(exist_ok=True)  # ensure dir exists before writing
    
    mnemonic_words = None
    
    if HAS_CRYPTO:
        try:
            from .recovery import generate_mnemonic, entropy_to_ed25519_seed
            # New BIP39 path
            mnemonic_words, entropy = generate_mnemonic()
            seed = entropy_to_ed25519_seed(entropy)
            private_key = Ed25519PrivateKey.from_private_bytes(seed)
            pub_bytes = private_key.public_key().public_bytes(Encoding.Raw, PublicFormat.Raw)
            priv_bytes = private_key.private_bytes(Encoding.Raw, PrivateFormat.Raw, NoEncryption())
            pub_hex = pub_bytes.hex()
            priv_hex = priv_bytes.hex()
            entropy_hash = hashlib.sha256(entropy).hexdigest()[:32]
        except Exception:
            # Fallback to random if recovery module fails
            private_key = Ed25519PrivateKey.generate()
            pub_bytes = private_key.public_key().public_bytes(Encoding.Raw, PublicFormat.Raw)
            priv_bytes = private_key.private_bytes(Encoding.Raw, PrivateFormat.Raw, NoEncryption())
            pub_hex = pub_bytes.hex()
            priv_hex = priv_bytes.hex()
            entropy_hash = None
    else:
        # Fallback: random hex (records marked unverified)
        import secrets
        priv_hex = secrets.token_hex(32)
        pub_hex = hashlib.sha256(bytes.fromhex(priv_hex)).hexdigest()
        entropy_hash = None

    # DID identifier = first 32 chars of sha256(pubkey hex) = 16 bytes
    agent_id = hashlib.sha256(pub_hex.encode()).hexdigest()[:32]
    did = f"did:ecp:{agent_id}"

    identity = {
        "did": did,
        "pub_key": pub_hex,
        "priv_key": priv_hex,   # local only, NEVER transmitted
        "created_at": _now_ms(),
        "verified": HAS_CRYPTO,
    }
    
    # Store entropy hash for recovery verification (NOT the entropy itself)
    if entropy_hash:
        identity["recovery_version"] = 1
        identity["entropy_hash"] = entropy_hash

    ifile.write_text(json.dumps(identity, indent=2))
    
    # Attach mnemonic transiently (NOT persisted to disk)
    if mnemonic_words:
        identity["_mnemonic"] = mnemonic_words
    
    return identity


def sign(identity: dict, data: str) -> str:
    """
    Sign data with agent's private key.
    Returns: "ed25519:{hex}" or "unverified"
    """
    if not HAS_CRYPTO:
        return "unverified"
    # Use crypto_priv_key (migrated identity) or priv_key (native Ed25519)
    priv_hex = identity.get("crypto_priv_key") or identity.get("priv_key")
    if not priv_hex:
        return "unverified"
    try:
        priv_bytes = bytes.fromhex(priv_hex)
        private_key = Ed25519PrivateKey.from_private_bytes(priv_bytes)
        sig_bytes = private_key.sign(data.encode())
        return f"ed25519:{sig_bytes.hex()}"
    except Exception:
        return "unverified"
