tools.py

from typing import List, Dict, Optional
from langchain.tools import StructuredTool
import json
# Now import from payments module
from src.payments.agent import process_purchase_request
# Load datasets
with open("data/supplements.json", "r") as f:
    SUPPLEMENTS = json.load(f)

with open("data/prescription_drugs.json", "r") as f:
    PRESCRIPTION_DRUGS = json.load(f)

# Combine both datasets for unified searching
ALL_PRODUCTS = {**SUPPLEMENTS, **PRESCRIPTION_DRUGS}

# Supplement search and retrieval functions
def search_supplements_by_condition(query: str) -> List[Dict]:
    """
    Search for supplements that are recommended for a specific health condition,
    well-being goal, or enhancement purpose.
    
    Args:
        query: The condition, goal, or purpose to search supplements for 
        (e.g., 'hair_loss', 'cognitive_function', 'muscle_growth')
    
    Returns:
        A list of supplements that match the query
    """
    matching_supplements = []
    query = query.lower()
    
    for supp_id, supp_data in SUPPLEMENTS.items():
        # Check conditions
        conditions = [c.lower() for c in supp_data.get("conditions", [])]
        # Check benefits/enhancements
        benefits = [b.lower() for b in supp_data.get("benefits", [])]
        # Check categories
        categories = [cat.lower() for cat in supp_data.get("categories", [])]
        
        # Match against all fields
        if (query in conditions or 
            query in benefits or 
            query in categories or 
            query in supp_data["name"].lower() or 
            query in supp_data["description"].lower()):
            
            matching_supplements.append({
                "id": supp_id,
                "name": supp_data["name"],
                "description": supp_data["description"],
                "link": supp_data["affiliate_link"],
                "conditions": supp_data.get("conditions", []),
                "benefits": supp_data.get("benefits", []),
                "categories": supp_data.get("categories", []),
                "disclaimers": supp_data.get("disclaimers", [])
            })
    
    return matching_supplements

def search_supplements_by_name(name: str) -> Optional[Dict]:
    """
    Search for a supplement by its name.
    
    Args:
        name: The name of the supplement to search for (e.g., 'Vitamin D3', 'Magnesium')
    
    Returns:
        Details of the supplement if found, None otherwise
    """
    name = name.lower()
    for supp_id, supp_data in SUPPLEMENTS.items():
        if name in supp_data["name"].lower():
            return {
                "id": supp_id,
                "name": supp_data["name"],
                "description": supp_data["description"],
                "link": supp_data["affiliate_link"],
                "conditions": supp_data.get("conditions", []),
                "benefits": supp_data.get("benefits", []),
                "categories": supp_data.get("categories", []),
                "disclaimers": supp_data.get("disclaimers", [])
            }
    return None

def get_supplement_details(supplement_id: str) -> Optional[Dict]:
    """
    Get detailed information about a specific supplement.
    
    Args:
        supplement_id: The ID of the supplement to retrieve
    
    Returns:
        Detailed information about the supplement or None if not found
    """
    if supplement_id in SUPPLEMENTS:
        supp_data = SUPPLEMENTS[supplement_id]
        return {
            "id": supplement_id,
            "name": supp_data["name"],
            "description": supp_data["description"],
            "link": supp_data["affiliate_link"],
            "conditions": supp_data.get("conditions", []),
            "disclaimers": supp_data.get("disclaimers", [])
        }
    return None

def list_all_supplements() -> List[Dict]:
    """
    List all available supplements in the database.
    
    Returns:
        A list of all supplements with their basic information
    """
    return [
        {
            "id": supp_id,
            "name": supp_data["name"],
            "description": supp_data["description"],
            "conditions": supp_data.get("conditions", [])
        }
        for supp_id, supp_data in SUPPLEMENTS.items()
    ]

# Prescription drug search and retrieval functions
def search_prescription_drugs_by_condition(query: str) -> List[Dict]:
    """
    Search for prescription drugs that are used for a specific health condition,
    well-being goal, or therapeutic purpose.
    
    Args:
        query: The condition, goal, or purpose to search drugs for 
        (e.g., 'hypertension', 'depression', 'hair_loss')
    
    Returns:
        A list of prescription drugs that match the query
    """
    matching_drugs = []
    query = query.lower()
    
    for drug_id, drug_data in PRESCRIPTION_DRUGS.items():
        # Check conditions
        conditions = [c.lower() for c in drug_data.get("conditions", [])]
        # Check benefits
        benefits = [b.lower() for b in drug_data.get("benefits", [])]
        # Check categories
        categories = [cat.lower() for cat in drug_data.get("categories", [])]
        
        # Match against all fields
        if (query in conditions or 
            query in benefits or 
            query in categories or 
            query in drug_data["name"].lower() or 
            query in drug_data["description"].lower()):
            
            matching_drugs.append({
                "id": drug_id,
                "name": drug_data["name"],
                "description": drug_data["description"],
                "link": drug_data["affiliate_link"],
                "conditions": drug_data.get("conditions", []),
                "benefits": drug_data.get("benefits", []),
                "categories": drug_data.get("categories", []),
                "disclaimers": drug_data.get("disclaimers", []),
                "requires_prescription": True
            })
    
    return matching_drugs

def search_prescription_drugs_by_name(name: str) -> Optional[Dict]:
    """
    Search for a prescription drug by its name.
    
    Args:
        name: The name of the drug to search for (e.g., 'Metformin', 'Atorvastatin')
    
    Returns:
        Details of the prescription drug if found, None otherwise
    """
    name = name.lower()
    for drug_id, drug_data in PRESCRIPTION_DRUGS.items():
        if name in drug_data["name"].lower():
            return {
                "id": drug_id,
                "name": drug_data["name"],
                "description": drug_data["description"],
                "link": drug_data["affiliate_link"],
                "conditions": drug_data.get("conditions", []),
                "benefits": drug_data.get("benefits", []),
                "categories": drug_data.get("categories", []),
                "disclaimers": drug_data.get("disclaimers", []),
                "requires_prescription": True
            }
    return None

def get_prescription_drug_details(drug_id: str) -> Optional[Dict]:
    """
    Get detailed information about a specific prescription drug.
    
    Args:
        drug_id: The ID of the prescription drug to retrieve
    
    Returns:
        Detailed information about the drug or None if not found
    """
    if drug_id in PRESCRIPTION_DRUGS:
        drug_data = PRESCRIPTION_DRUGS[drug_id]
        return {
            "id": drug_id,
            "name": drug_data["name"],
            "description": drug_data["description"],
            "link": drug_data["affiliate_link"],
            "conditions": drug_data.get("conditions", []),
            "benefits": drug_data.get("benefits", []),
            "categories": drug_data.get("categories", []),
            "disclaimers": drug_data.get("disclaimers", []),
            "requires_prescription": True
        }
    return None

def list_all_prescription_drugs() -> List[Dict]:
    """
    List all available prescription drugs in the database.
    
    Returns:
        A list of all prescription drugs with their basic information
    """
    return [
        {
            "id": drug_id,
            "name": drug_data["name"],
            "description": drug_data["description"],
            "conditions": drug_data.get("conditions", []),
            "requires_prescription": True
        }
        for drug_id, drug_data in PRESCRIPTION_DRUGS.items()
    ]

# Combined search functions
def search_all_products_by_condition(query: str) -> List[Dict]:
    """
    Search for both supplements and prescription drugs that are recommended for a specific health condition,
    well-being goal, or enhancement purpose.
    
    Args:
        query: The condition, goal, or purpose to search for
        (e.g., 'hair_loss', 'cognitive_function', 'depression')
    
    Returns:
        A list of supplements and prescription drugs that match the query
    """
    supplements = search_supplements_by_condition(query)
    prescription_drugs = search_prescription_drugs_by_condition(query)
    
    return supplements + prescription_drugs

def search_all_products_by_name(name: str) -> List[Dict]:
    """
    Search for both supplements and prescription drugs by name.
    
    Args:
        name: The name to search for (e.g., 'Vitamin D3', 'Metformin')
    
    Returns:
        List of products matching the name query
    """
    results = []
    
    supplement = search_supplements_by_name(name)
    if supplement:
        results.append(supplement)
        
    prescription = search_prescription_drugs_by_name(name)
    if prescription:
        results.append(prescription)
        
    return results

def get_product_details(product_id: str) -> Optional[Dict]:
    """
    Get detailed information about a specific product (supplement or prescription drug).
    
    Args:
        product_id: The ID of the product to retrieve
    
    Returns:
        Detailed information about the product or None if not found
    """
    supplement = get_supplement_details(product_id)
    if supplement:
        return supplement
        
    prescription = get_prescription_drug_details(product_id)
    if prescription:
        return prescription
        
    return None

def list_all_products() -> List[Dict]:
    """
    List all available products (supplements and prescription drugs) in the database.
    
    Returns:
        A list of all products with their basic information
    """
    supplements = list_all_supplements()
    prescription_drugs = list_all_prescription_drugs()
    
    return supplements + prescription_drugs

def create_payment_link(request: str) -> str:
    """
    Create a payment link for products and consultations based on a natural language request.
    
    Args:
        request: A string describing what products and quantities to purchase
        (e.g., "I want to buy 1 tretinoin and 1 consultation")
    
    Returns:
        The payment link URL
    """
    
    # Process the request and get payment link
    response = process_purchase_request(request)
    #response['messages'][0].content
    return response

# Create and export Langchain tools
def get_tools():
    """Return all the tools needed for the agent"""
    return [
        StructuredTool.from_function(search_supplements_by_condition),
        StructuredTool.from_function(search_prescription_drugs_by_condition),
        StructuredTool.from_function(search_all_products_by_condition),
        StructuredTool.from_function(search_supplements_by_name),
        StructuredTool.from_function(search_prescription_drugs_by_name),
        StructuredTool.from_function(search_all_products_by_name),
        StructuredTool.from_function(get_supplement_details),
        StructuredTool.from_function(get_prescription_drug_details),
        StructuredTool.from_function(get_product_details),
        StructuredTool.from_function(list_all_supplements),
        StructuredTool.from_function(list_all_prescription_drugs),
        StructuredTool.from_function(list_all_products),
        StructuredTool.from_function(create_payment_link)
    ]

# Export data for use in main.py
def get_supplements():
    return SUPPLEMENTS

def get_prescription_drugs():
    return PRESCRIPTION_DRUGS 


def get_available_products() -> Dict:
    """
    Get a list of all available products in the catalog.
    
    Returns:
        A dictionary containing available products categorized by type
    """
    # Load the products catalog
    with open("src/payments/products.json", "r") as f:
        products_catalog = json.load(f)
    
    # Separate into supplements, prescription drugs, and services
    supplements = []
    prescription_drugs = []
    services = []
    
    for product_id in products_catalog.keys():
        if product_id == "NP_consultation":
            services.append(product_id)  # Add consultation as a service
        elif product_id in PRESCRIPTION_DRUGS:
            prescription_drugs.append(product_id)
        elif product_id in SUPPLEMENTS:
            supplements.append(product_id)
    
    return {
        "supplements": supplements,
        "prescription_drugs": prescription_drugs,
        "services": services,
        "all_products": list(products_catalog.keys())
    }

if __name__ == "__main__":
    sample_request = '1 Finasteride consultation, 1 Minoxidil consultation'
    result = process_purchase_request(sample_request)
    print(result)
    
    breakpoint()

    asd = 1
    # You can test with other examples
    # result2 = process_purchase_request("I want to buy 3 items of product_i and 2 of product_k")
    # print(result2)