from langchain.output_parsers import StructuredOutputParser, ResponseSchema
from langchain.prompts import PromptTemplate
from langchain_openai import ChatOpenAI
from .config import OPENAI_API_KEY

# Define the expected response schema
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."),
]

# Initialize the structured output parser
output_parser = StructuredOutputParser.from_response_schemas(response_schemas)

# Instantiate the ChatOpenAI LLM
llm = ChatOpenAI(temperature=0, openai_api_key=OPENAI_API_KEY)

# Prepare the prompt template
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)}"
    

swap_response_schemas = [
    ResponseSchema(name="amount", description="The amount of mUSD to swap."),
    ResponseSchema(name="from_currency", description="The token to swap from (should always be 'mUSD')."),
    ResponseSchema(name="to_currency", description="The token to receive (should always be 'BTC')."),
    ResponseSchema(name="router_address", description="The Dumpy Swap router address for executing the swap."),
]
swap_output_parser = StructuredOutputParser.from_response_schemas(swap_response_schemas)

swap_prompt_template = PromptTemplate(
    template=(
        "Extract swap transaction details from this request:\n{input}\n\n"
        "- The token to swap from should always be 'mUSD'.\n"
        "- The token to receive should always be 'BTC'.\n"
        "- The router address should always be '0xC2E61936a542D78b9c3AA024fA141c4C632DF6c1'.\n\n"
        "{format_instructions}"
    ),
    input_variables=["input"],
    partial_variables={"format_instructions": swap_output_parser.get_format_instructions()},
)

def extract_swap_details(prompt: str):
    formatted_prompt = swap_prompt_template.format(input=prompt)
    response = llm.invoke(formatted_prompt)
    try:
        return swap_output_parser.parse(response.content)
    except Exception as e:
        return f"Failed to extract swap details: {str(e)}"