import hashlib
import json
from time import time

class Blockchain:
    def __init__(self):
        self.chain = []
        self.current_transactions = []
        self.user_wallets = {}
        self.user_gpus = {}

        # Genesis 블록 생성
        self.new_block(previous_hash="1")

    def new_block(self, previous_hash=None):
        """
        블록체인에 새로운 블록 추가
        :param previous_hash: 이전 블록의 해시 값
        :return: 새로 생성된 블록
        """
        for id, mem in self.user_gpus.items():
            self.user_wallets[id] += mem//2
        
        block = {
            'index': len(self.chain) + 1,
            'timestamp': time(),
            'transactions': self.current_transactions,
            'previous_hash': previous_hash or self.hash(self.chain[-1]),
        }
        
        # 현재 트랜잭션 초기화
        self.current_transactions = []
        
        # 블록을 체인에 추가
        self.chain.append(block)
        return block

    def new_transaction(self, id, kind, data):
        """
        새로운 트랜잭션 생성 및 추가
        id: 요청자
        kind: 요청 종류 (inference, add, out)
        data: inference 시 [입력 prompt, output], peer add 시 gpu mem, out 시 또한 gpu mem
        
        return: 해당 트랜잭션을 포함한 블록의 인덱스
        """
        
        transaction = {
            'id': id,
            'kind': kind,
            'data': data,
        }
        self.current_transactions.append(transaction)

        # 각 유저의 행동력 업데이트
        if kind == "inference":
            if id not in self.user_wallets:
                # 새로운 유저인 경우 초기 행동력 설정
                self.user_wallets[id] = 10  # 초기 행동력을 10으로 설정
                self.user_wallets[id] -= 1
            else:
                # inference 요청 시 차감
                self.user_wallets[id] -= 1
                if self.user_wallets[id]<0:
                    self.user_wallets[id] = 0
        elif kind == "add":
            if id not in self.user_gpus:
                self.user_gpus[id] = int(data)
            else:
                self.user_gpus[id] += int(data)
        elif kind == "out":
            if id in self.user_gpus:
                del(self.user_gpus[id])

        return self.last_block['index'] + 1

    def get_user_balance(self, id):
        # 특정 유저의 행동력을 조회
        return self.user_wallets.get(id, 10)

    def get_user_gpu_mem(self, id):
        # 특정 유저의 기여 gpu를 조회
        return self.user_gpus.get(id, 0)

    def get_total_gpu_mem(self):
        # 전체 메모리 반환
        result = 0
        for mem in self.user_gpus.values():
            result += int(mem)
        return result

    def get_total_coin(self):
        # 전체 코인(행동력) 수 반환
        result = 0
        for coin in self.user_wallets.values():
            result += int(coin)
        return result
    @property
    def last_block(self):
        """
        체인의 마지막 블록 반환
        """
        return self.chain[-1]
    
    @staticmethod
    def hash(block):
        """
        블록을 SHA-256으로 해싱
        :param block: 블록
        :return: 해시 값
        """
        block_string = json.dumps(block, sort_keys=True).encode()
        return hashlib.sha256(block_string).hexdigest()