import requests
import json
from typing import Dict, Any

import os
port = os.getenv('NIM_HTTP_API_PORT', '8000')
print(port)

SEQUENCE = "MTEYKLVVVGACGVGKSALTIQLIQNHFVDEYDPTIEDSYRKQVVID"

def query_boltz2_nim(
    input_data: Dict[str, Any],
    base_url: str = f"http://localhost:{port}"
) -> Dict[str, Any]:
    """
    Query the Boltz2 NIM with input data.

    Args:
        input_data: Dictionary containing the prediction request data
        base_url: Base URL of the NIM service (default: http://localhost:8000)

    Returns:
        Dictionary containing the prediction response
    """
    # Construct the full URL
    url = f"{base_url}/biology/mit/boltz2/predict"

    # Set headers
    headers = {
        "Content-Type": "application/json"
    }

    try:
        # Make the POST request
        response = requests.post(url, json=input_data, headers=headers)

        # Check if request was successful
        response.raise_for_status()

        # Return the JSON response
        return response.json()

    except requests.exceptions.RequestException as e:
        print(f"Error querying NIM: {e}")
        if hasattr(e.response, 'text'):
            print(f"Response text: {e.response.text}")
        raise

# Example usage
if __name__ == "__main__":
    # Example input data - modify this according to your BoltzPredictionRequest structure
    example_input = {"polymers":[
        {
            "id": "A",
            "molecule_type": "protein",
            "sequence": SEQUENCE
        }
    ]}

    try:
        # Query the NIM
        result = query_boltz2_nim(example_input)

        # Print the result
        print("Prediction result:")
        print(json.dumps(result, indent=2))

    except Exception as e:
        print(f"Failed to get prediction: {e}")