from enum import IntEnum
import time
from typing import Optional, Tuple, TypedDict
from datetime import datetime
from web3 import Web3
from eth_account import Account
from acp_plugin_gamesdk.acp_token_abi import ACP_TOKEN_ABI
import requests
from eth_account.messages import encode_defunct
import json

class MemoType(IntEnum):
    MESSAGE = 0
    CONTEXT_URL = 1
    IMAGE_URL = 2
    VOICE_URL = 3
    OBJECT_URL = 4
    TXHASH = 5

class IMemo(TypedDict):
    content: str
    memoType: MemoType
    isSecured: bool
    nextPhase: int
    jobId: int
    numApprovals: int
    sender: str

class IJob(TypedDict):
    id: int
    client: str
    provider: str
    budget: int
    amountClaimed: int
    phase: int
    memoCount: int
    expiredAt: int
    evaluatorCount: int

JobResult = Tuple[int, str, str, str, str, str, str, str, int]

class AcpToken:
    def __init__(
        self,
        wallet_private_key: str,
        agent_wallet_address: str,
        network_url: str,
        acp_base_url: Optional[str] = None,
        contract_address: str = "0x2422c1c43451Eb69Ff49dfD39c4Dc8C5230fA1e6",
        virtuals_token_address: str = "0xbfAB80ccc15DF6fb7185f9498d6039317331846a",
    ):
        self.web3 = Web3(Web3.HTTPProvider(network_url))
        self.account = Account.from_key(wallet_private_key)
        self.agent_wallet_address = agent_wallet_address
        self.contract_address = Web3.to_checksum_address(contract_address)
        self.virtuals_token_address = Web3.to_checksum_address(virtuals_token_address)
        self.contract = self.web3.eth.contract(
            address=self.contract_address,
            abi=ACP_TOKEN_ABI
        )
        self.virtuals_token_contract = self.web3.eth.contract(
            address=self.virtuals_token_address,
            abi=[{
                "inputs": [
                    {
                        "internalType": "address",
                        "name": "spender",
                        "type": "address"
                    },
                    {
                        "internalType": "uint256",
                        "name": "amount", 
                        "type": "uint256"
                    }
                ],
                "name": "approve",
                "outputs": [
                    {
                        "internalType": "bool",
                        "name": "",
                        "type": "bool"
                    }
                ],
                "stateMutability": "nonpayable",
                "type": "function"
            }]
        )
        self.acp_base_url = acp_base_url if acp_base_url else "https://acpx-staging.virtuals.io/api"
    def get_agent_wallet_address(self) -> str:
        return self.agent_wallet_address
        
    def get_contract_address(self) -> str:
        return self.contract_address

    def validate_transaction(self, hash_value: str) -> object:
        try:
            response = requests.post(f"{self.acp_base_url}/acp-agent-wallets/trx-result", json={"userOpHash": hash_value})
            return response.json()
        except Exception as error:
            print(f"Error getting job_id: {error}")
            raise Exception("Failed to get job_id")

    def create_job(
        self,
        provider_address: str,
        evaluator_address: str,
        expire_at: datetime
    ) -> dict:
        try:
            provider_address = Web3.to_checksum_address(provider_address)
            expire_timestamp = int(expire_at.timestamp())
        
            # Sign the transaction
            trx_data, signature = self._sign_transaction(
                "createJob", 
                [provider_address, evaluator_address, expire_timestamp]
            )
                
            # Prepare payload
            payload = {
                "agentWallet": self.get_agent_wallet_address(),
                "trxData": trx_data,
                "signature": signature
            }
            
            # Submit to custom API
            api_url = f"{self.acp_base_url}/acp-agent-wallets/transactions"
            response = requests.post(api_url, json=payload)
                
            # Return transaction hash or response ID
            return { "txHash": response.json().get("data", {}).get("userOpHash", "")}
        
        except Exception as error:
            print(f"Error creating job: {error}")
            raise Exception("Failed to create job")

    def approve_allowance(self, price_in_wei: int) -> str:
        try:
            trx_data, signature = self._sign_transaction(
                "approve", 
                [self.contract_address, price_in_wei],
                self.virtuals_token_address
            )
            
            payload = {
                "agentWallet": self.get_agent_wallet_address(),
                "trxData": trx_data,
                "signature": signature
            }
            
            api_url = f"{self.acp_base_url}/acp-agent-wallets/transactions"
            response = requests.post(api_url, json=payload)
            
            if (response.status_code != 200):
                raise Exception("Failed to approve allowance")
            
            return response.json()
        except Exception as error:
            print(f"Error approving allowance: {error}")
            raise Exception("Failed to approve allowance")

    def create_memo(
        self,
        job_id: int,
        content: str,
        memo_type: MemoType,
        is_secured: bool,
        next_phase: int
    ) -> dict:
        retries = 3
        while retries > 0:
            try:
                trx_data, signature = self._sign_transaction(
                    "createMemo", 
                    [job_id, content, memo_type, is_secured, next_phase]
                )
                
                payload = {
                    "agentWallet": self.get_agent_wallet_address(),
                    "trxData": trx_data,
                    "signature": signature
                }

                api_url = f"{self.acp_base_url}/acp-agent-wallets/transactions"
                response = requests.post(api_url, json=payload)
                
                if (response.status_code != 200):
                    raise Exception("Failed to create memo")
                
                return { "txHash": response.json().get("txHash", response.json().get("id", "")), "memoId": response.json().get("memoId", "")}
            except Exception as error:
                print(f"Error creating memo: {error}")
                retries -= 1
                time.sleep(2 * (3 - retries))
                
        raise Exception("Failed to create memo")

    def _sign_transaction(self, method_name: str, args: list, contract_address: Optional[str] = None) -> Tuple[dict, str]:
        if contract_address:
            encoded_data = self.virtuals_token_contract.encode_abi(method_name, args=args)
        else:
            encoded_data = self.contract.encode_abi(method_name, args=args)
        
        trx_data = {
            "target": contract_address if contract_address else self.get_contract_address(),
            "value": "0",
            "data": encoded_data
        }
        
        message_json = json.dumps(trx_data, separators=(",", ":"), sort_keys=False)
        message_bytes = message_json.encode()
        
        # Sign the transaction
        message = encode_defunct(message_bytes)
        signature = "0x" + self.account.sign_message(message).signature.hex()
        
        return trx_data, signature

    def sign_memo(
        self,
        memo_id: int,
        is_approved: bool,
        reason: Optional[str] = ""
    ) -> str:
        retries = 3
        while retries > 0:
            try:
                trx_data, signature = self._sign_transaction(
                    "signMemo", 
                    [memo_id, is_approved, reason]
                )
                
                payload = {
                    "agentWallet": self.get_agent_wallet_address(),
                    "trxData": trx_data,
                    "signature": signature
                }
                
                api_url = f"{self.acp_base_url}/acp-agent-wallets/transactions"
                response = requests.post(api_url, json=payload)
                
                if (response.status_code != 200):
                    raise Exception("Failed to sign memo")
                
                return response.json()
                
            except Exception as error:
                print(f"Error signing memo: {error}")
                retries -= 1
                time.sleep(2 * (3 - retries))
                
        raise Exception("Failed to sign memo")

    def set_budget(self, job_id: int, budget: int) -> str:
        try:
            trx_data, signature = self._sign_transaction(
                "setBudget", 
                [job_id, budget]
            )
            
            payload = {
                "agentWallet": self.get_agent_wallet_address(),
                "trxData": trx_data,
                "signature": signature
            }
            
            api_url = f"{self.acp_base_url}/acp-agent-wallets/transactions"
            response = requests.post(api_url, json=payload)
            
            if (response.status_code != 200):
                raise Exception("Failed to set budget")
            
            return response.json()
        except Exception as error:
            print(f"Error setting budget: {error}")
            raise Exception("Failed to set budget")

    def get_job(self, job_id: int) -> Optional[IJob]:
        try:
            job_data = self.contract.functions.jobs(job_id).call()
            
            if not job_data:
                return None
                
            return {
                'id': job_data[0],
                'client': job_data[1],
                'provider': job_data[2],
                'budget': int(job_data[3]),
                'amountClaimed': int(job_data[4]),
                'phase': int(job_data[5]),
                'memoCount': int(job_data[6]),
                'expiredAt': int(job_data[7]),
                'evaluatorCount': int(job_data[8])
            }
        except Exception as error:
            print(f"Error getting job: {error}")
            raise Exception("Failed to get job")

    def get_memo_by_job(
        self,
        job_id: int,
        memo_type: Optional[MemoType] = None
    ) -> Optional[IMemo]:
        try:
            memos = self.contract.functions.getAllMemos(job_id).call()
            
            if memo_type is not None:
                filtered_memos = [m for m in memos if m['memoType'] == memo_type]
                return filtered_memos[-1] if filtered_memos else None
            else:
                return memos[-1] if memos else None
        except Exception as error:
            print(f"Error getting memo: {error}")
            raise Exception("Failed to get memo")

    def get_memos_for_phase(
        self,
        job_id: int,
        phase: int,
        target_phase: int
    ) -> Optional[IMemo]:
        try:
            memos = self.contract.functions.getMemosForPhase(job_id, phase).call()
            
            target_memos = [m for m in memos if m['nextPhase'] == target_phase]
            return target_memos[-1] if target_memos else None
        except Exception as error:
            print(f"Error getting memos: {error}")
            raise Exception("Failed to get memos")
