import os
from contextlib import asynccontextmanager
from datetime import datetime
from hashlib import sha256

from pytonconnect import TonConnect as ExternalLib_TonConnect
from pytonconnect.storage import DefaultStorage
from tonsdk.utils import Address

from app.core._config import PROJECT_HOST
from app.core.logger import make_log
from app.core.models.wallet_connection import WalletConnection

TON_CONNECT_MANIFEST_URI = os.getenv("TON_CONNECT_MANIFEST_URI")


def unpack_wallet_info(wallet_info) -> dict:
    return {
        'provider': wallet_info.provider,
        'device': {
            'platform': wallet_info.device.platform,
            'app_name': wallet_info.device.app_name,
            'app_version': wallet_info.device.app_version,
            'max_protocol_version': wallet_info.device.max_protocol_version,
            'features': wallet_info.device.features
        } if wallet_info.device else None,
        'account': {
            'address': wallet_info.account.address,
            'chain': wallet_info.account.chain,
            'wallet_state_init': wallet_info.account.wallet_state_init,
            'public_key': wallet_info.account.public_key,
        } if wallet_info.account else None,
        'ton_proof': {
            'timestamp': wallet_info.ton_proof.timestamp,
            'domain_len': wallet_info.ton_proof.domain_len,
            'domain_val': wallet_info.ton_proof.domain_val,
            'payload': wallet_info.ton_proof.payload,
            'signature': wallet_info.ton_proof.signature,
        } if wallet_info.ton_proof else None
    }


class TonConnect:
    def __init__(self, callback_fn: tuple = None):
        self._manifest_uri = TON_CONNECT_MANIFEST_URI
        if not self._manifest_uri:
            self._manifest_uri = f"{PROJECT_HOST}/api/tonconnect-manifest.json"

        self._sdk_client = ExternalLib_TonConnect(
            manifest_url=self._manifest_uri, storage=DefaultStorage()
        )

        def status_change_callback(status):
            status = unpack_wallet_info(status)
            make_log("TonConnect", f"Changed status (connected={self.connected}): {status}", level='debug')

            if callback_fn:
                callback_fn[0](self, status, *callback_fn[1:])

        self.set_status_change_callback(status_change_callback)

    async def new_connection(self, app_name: str):
        wallets = self._sdk_client.get_wallets()
        for wallet in wallets:
            if wallet["app_name"] == app_name:
                return await self._sdk_client.connect(wallet)

    async def restore_connection(self):
        return await self._sdk_client.restore_connection()

    def pause_connection(self):
        try:
            self._sdk_client.pause_connection()
        except BaseException as e:
            make_log("pause_connection", e, level='error')

    @asynccontextmanager
    async def connection(self):
        try:
            yield None
        finally:
            self.pause_connection()

    @property
    def connected(self):
        # make_log("is_connected", self._sdk_client._storage._cache)
        return self._sdk_client.connected

    @property
    def connection_key(self):
        return self._sdk_client._storage._cache.get(DefaultStorage.KEY_CONNECTION, None)
        # return self._sdk_client._storage.get_item(DefaultStorage.KEY_CONNECTION, None)

    def set_status_change_callback(self, callback):
        self._sdk_client.on_status_change(callback)

    @classmethod
    def by_key(cls, connection_key: str, callback_fn: tuple = None):
        ton_connect = cls(callback_fn=callback_fn)
        ton_connect._sdk_client._storage._cache[DefaultStorage.KEY_CONNECTION] = connection_key
        # ton_connect._sdk_client._storage.set_item(DefaultStorage.KEY_CONNECTION, connection_key)
        # Immediately restore connection
        return ton_connect

    @classmethod
    def by_user(cls, session, user, callback_fn: tuple = None):
        def new_callback_fn(self, status, _callback_fn):
            try:
                if not session.query(
                    WalletConnection
                ).filter(WalletConnection.connection_id == sha256(self.connection_key.encode()).hexdigest()) \
                        .count():
                    new_connection = WalletConnection(
                        user_id=user.id,
                        network='ton',
                        wallet_key=f"{status['device'].get('app_name', 'UNKNOWN_NAME')}=={status['device'].get('app_version', '1.0')}",
                        connection_id=sha256(self.connection_key.encode()).hexdigest(),
                        wallet_address=Address(status['account']['address']).to_string(1, 1, 1),
                        keys={
                            'connection_key': self.connection_key,
                        },
                        meta={
                            key: status[key] for key in status if status[key]
                        },
                        created=datetime.now(),
                        updated=datetime.now(),
                        invalidated=False
                    )
                    session.add(new_connection)
                    session.commit()
            except BaseException as e:
                make_log("TonConnect.save_connection", e, level='error')

            if _callback_fn:
                _callback_fn[0](self, status, *_callback_fn[1:])

        ton_connect = cls(callback_fn=(new_callback_fn, callback_fn))
        connections = (
            session.query(WalletConnection).filter(
                WalletConnection.user_id == user.id,
                WalletConnection.invalidated == False,
                WalletConnection.network == 'ton'
            )
        )
        if connections.count() == 0:
            return ton_connect, None

        connection = connections.first()
        ton_connect._sdk_client._storage._cache[DefaultStorage.KEY_CONNECTION] = connection.keys["connection_key"]
        # ton_connect._sdk_client._storage.set_item(DefaultStorage.KEY_CONNECTION, bytes.fromhex(connection.keys["connection_key"]))
        # Immediately restore connection
        return ton_connect, connection