import streamlit as st
import requests
from typing import List, Dict
import os

API_URL = os.getenv('POND_API_URL')
API_TOKEN = os.getenv('POND_API_TOKEN')

st.title('🌊 Pond Model Demo')

st.sidebar.header('Model Selection')
model_info = {
    1: {
        'name': 'Security Model',
        'description': 'Analyze whether an account is secure by detecting and examining malicious activities within complex blockchain data structures.'
    },
    2: {
        'name': 'Sybil Model',
        'description': 'An model aims to detect on-chain "Sybil Attacks"'
    },
    3: {
        'name': 'ZORA NFT Recommendation',
        'description': 'On-Chain Recommendation System: Making Discoveries & Spread Easier for Everyone'
    }
}

selected_model = st.sidebar.selectbox(
    'Choose a model',
    list(model_info.keys()),
    format_func=lambda x: model_info[x]['name']
)

# Show model description
st.markdown(f"### {model_info[selected_model]['name']}")
st.markdown(model_info[selected_model]['description'])

# Input section
st.header('Input Wallet Addresses')
wallet_input = st.text_area(
    'Enter wallet addresses (one per line)',
    height=100,
    help='Enter wallet addresses, one per line'
)

def predict(addresses: List[str], model_id: int) -> Dict:
    """
    Make prediction using the Pond API
    """
    if not API_URL:
        st.error('API URL is not configured. Please set POND_API_URL environment variable.')
        return None
    
    headers = {
        "Content-Type": "application/json"
    }
    
    try:
        payload = {
            "req_type": "1",
            "access_token": API_TOKEN,
            "input_keys": addresses,
            "model_id": model_id
        }
        
        headers = {
            "Content-Type": "application/json"
        }
        
        # Make the API call with explicit method
        session = requests.Session()
        req = requests.Request('POST', 
                             API_URL,
                             json=payload,
                             headers=headers)
        prepped = req.prepare()
            
        response = session.send(prepped,
                              allow_redirects=True)
        
        # Check response
        if response.status_code != 200:
            st.error(f"API Error: {response.status_code}")
            st.error(f"Response: {response.text}")
            return None
            
        return response.json()
    except Exception as e:
        st.error(f"Error making prediction: {str(e)}")
        return None

def display_results(response: Dict, model_id: int):
    """
    Display the results based on model type
    """
    if not response or 'resp_items' not in response:
        return
    
    if model_id in [1, 2]:  # Security and Sybil models
        st.header('Results')
        for item in response['resp_items']:
            score = item['score']
            address = item['input_key']
            
            # Create color coding based on score
            if score < 0.3:
                color = 'green'
            elif score < 0.7:
                color = 'orange'
            else:
                color = 'red'
                
            st.markdown(f"""
            **Address**: `{address}`  
            **Score**: <span style='color: {color}'>{score:.4f}</span>
            """, unsafe_allow_html=True)
            st.markdown("")

    elif model_id == 3:  # NFT Recommendation model
        st.header('NFT Recommendations')
        for item in response['resp_items']:
            address = item['input_key']
            st.subheader(f'Recommendations for: `{address}`')
            
            if 'candidates' in item:
                for idx, candidate in enumerate(item['candidates'][:5], 1):
                    st.markdown(f"""
                    **{idx}. NFT ID**: `{candidate['item_id']}`  
                    **Score**: {candidate['score']:.4f}
                    """)
            st.markdown("")

# Process button
if st.button('Get Predictions'):
    if wallet_input:
        # Process input addresses
        addresses = [addr.strip() for addr in wallet_input.split('\n') if addr.strip()]
        
        # Validate addresses
        valid_addresses = [addr for addr in addresses if addr.startswith('0x')]
        
        if not valid_addresses:
            st.error('Please enter valid Ethereum addresses (starting with 0x)')
        else:
            with st.spinner('Making predictions...'):
                response = predict(valid_addresses, selected_model)
                if response:
                    display_results(response, selected_model)
    else:
        st.warning('Please enter at least one wallet address')

# Footer
st.markdown("---")
st.markdown("ℹ️ This is a demo interface for the Pond Model API. For production use, please refer to the official documentation.")