import os
import json
from dotenv import load_dotenv
from web3 import Web3
from web3.exceptions import Web3Exception
from langchain.tools import tool
from langchain_openai import ChatOpenAI
from langchain.output_parsers import StructuredOutputParser, ResponseSchema
from langchain.prompts import PromptTemplate

# ✅ Load environment variables
USER_PROJECT_DIR = os.getcwd()  # Get the directory where the script is run
USER_ENV_PATH = os.path.join(USER_PROJECT_DIR, ".env")

if os.path.exists(USER_ENV_PATH):
    load_dotenv(USER_ENV_PATH)
else:
    print("⚠️ Warning: No `.env` file found in your project directory! Transactions requiring signing may fail.")

# ✅ Load Private Keys from User's `.env`
PRIVATE_KEY = os.getenv("PRIVATE_KEY")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
if not PRIVATE_KEY:
    print("⚠️ Warning: PRIVATE_KEY not set. Please create a `.env` file in your project with your keys.")
    PRIVATE_KEY = None  # Allow package to be installed but prevent transactions

RPC_URL = "https://rpc.test.mezo.org"

# ✅ Initialize Web3 instance
web3_instance = Web3(Web3.HTTPProvider(RPC_URL))

# ✅ Check if connection is successful
if not web3_instance.is_connected():
    raise ConnectionError("❌ Failed to connect to Mezo Matsnet RPC.")

# ✅ Create Account Object
account = Web3.eth.account.from_key(PRIVATE_KEY)
sender_address = account.address

MUSD_ADDRESS = "0x637e22A1EBbca50EA2d34027c238317fD10003eB"  
ERC20_ABI = json.loads('[{"constant": false, "inputs": [{"name": "recipient", "type": "address"},{"name": "amount", "type": "uint256"}],"name": "transfer","outputs": [{"name": "", "type": "bool"}],"stateMutability": "nonpayable","type":"function"}]')
musd_contract = Web3.eth.contract(address=MUSD_ADDRESS, abi=ERC20_ABI)

# ✅ Define Structured Output Parser for Transaction Details
response_schemas = [
    ResponseSchema(name="amount", description="The amount of cryptocurrency to transfer."),
    ResponseSchema(name="currency", description="The cryptocurrency (BTC or mUSD)."),
    ResponseSchema(name="recipient", description="The recipient's wallet address."),
]

output_parser = StructuredOutputParser.from_response_schemas(response_schemas)

# ✅ Define LLM Prompt for Parsing Transactions
llm = ChatOpenAI(temperature=0, openai_api_key=os.getenv("OPENAI_API_KEY"))

prompt_template = PromptTemplate(
    template="Extract transaction details from this request:\n{input}\n{format_instructions}",
    input_variables=["input"],
    partial_variables={"format_instructions": output_parser.get_format_instructions()},
)

def extract_transaction_details(prompt: str):
    """
    Uses LLM to extract structured transaction details from user input.
    """
    formatted_prompt = prompt_template.format(input=prompt)
    response = llm.invoke(formatted_prompt)

    try:
        return output_parser.parse(response.content)
    except Exception as e:
        return f"❌ Failed to extract transaction details: {str(e)}"

@tool
def mezo_agent_transaction_btc(transaction_prompt: str) -> str:
    """
    Sends BTC on Mezo Matsnet. Parses a transaction request and executes the transfer.

    Example usage:
    ```
    mezo_agent_transaction_btc("Send 0.05 BTC to 0x93a1Eadb069A791d23aAeDF3C272E2905bb63240")
    ```

    :param transaction_prompt: Natural language transaction request.
    :return: Transaction hash or failure message.
    """
    transaction_details = extract_transaction_details(transaction_prompt)

    if isinstance(transaction_details, str):  # Handle errors
        return transaction_details

    amount = float(transaction_details["amount"])
    currency = transaction_details["currency"].lower()
    recipient = transaction_details["recipient"]

    if currency != "btc":
        return "❌ This function only supports BTC transactions."

    # ✅ Convert amount to Wei (BTC uses 18 decimals on Mezo Matsnet)
    amount_wei = Web3.to_wei(amount, "ether")

    # ✅ Check sender's BTC balance
    sender_balance = Web3.eth.get_balance(sender_address)
    sender_balance_btc = Web3.from_wei(sender_balance, "ether")

    if sender_balance < amount_wei:
        return f"❌ Insufficient BTC balance! You have {sender_balance_btc} BTC but need {amount} BTC."

    # ✅ Fetch nonce and gas price
    nonce = Web3.eth.get_transaction_count(sender_address)
    gas_price = Web3.eth.gas_price
    gas_limit = Web3.eth.estimate_gas({"to": recipient, "value": amount_wei, "from": sender_address})

    tx = {
        "to": recipient,
        "value": amount_wei,
        "gas": gas_limit,
        "gasPrice": gas_price,
        "nonce": nonce,
        "chainId": 31611,  # Mezo Testnet Chain ID
    }

    try:
        # ✅ Sign and send the transaction
        signed_tx = account.sign_transaction(tx)
        tx_hash = Web3.eth.send_raw_transaction(signed_tx.raw_transaction)
        return f"✅ BTC Transaction Successful! Hash: {tx_hash.hex()}"

    except Web3Exception as e:
        return f"❌ BTC Transaction Failed: {str(e)}"

@tool
def mezo_agent_transaction_musd(transaction_prompt: str) -> str:
    """
    Sends mUSD on Mezo Matsnet. Parses a transaction request and executes the transfer.

    Example usage:
    ```
    mezo_agent_transaction_musd("Transfer 100 mUSD to 0xABC123...")
    ```

    :param transaction_prompt: Natural language transaction request.
    :return: Transaction hash or failure message.
    """
    transaction_details = extract_transaction_details(transaction_prompt)

    if isinstance(transaction_details, str):  # Handle errors
        return transaction_details

    amount = float(transaction_details["amount"])
    currency = transaction_details["currency"].lower()
    recipient = transaction_details["recipient"]

    if currency != "musd":
        return "❌ This function only supports mUSD transactions."

    # ✅ Convert amount to Wei (mUSD uses 18 decimals)
    amount_musd_wei = int(amount * 10**18)
    nonce = Web3.eth.get_transaction_count(sender_address)
    gas_price = Web3.eth.gas_price

    try:
        # ✅ Build the transaction for mUSD transfer
        txn = musd_contract.functions.transfer(recipient, amount_musd_wei).build_transaction({
            "from": sender_address,
            "nonce": nonce,
            "gas": 50000,
            "gasPrice": gas_price,
        })

        signed_txn = Web3.eth.account.sign_transaction(txn, PRIVATE_KEY)
        tx_hash = Web3.eth.send_raw_transaction(signed_txn.raw_transaction)

        return f"✅ mUSD Transaction Successful! Hash: {tx_hash.hex()}"

    except Web3Exception as e:
        return f"❌ mUSD Transaction Failed: {str(e)}"