import os
import streamlit as st
import pandas as pd
from datetime import datetime, timedelta
import time
from huggingface_hub import InferenceClient
import dotenv

# Load environment variables
dotenv.load_dotenv()

# Set page configuration
st.set_page_config(
    page_title="Debt Collection AI",
    page_icon="💰",
    layout="wide"
)

# Custom CSS for styling
st.markdown("""
<style>
    .main-header {
        font-size: 2.5rem;
        color: #1E3A8A;
        text-align: center;
        margin-bottom: 2rem;
        font-weight: bold;
        text-shadow: 1px 1px 2px rgba(0,0,0,0.1);
    }
    .agent-button {
        background-color: #3B82F6;
        color: white;
        padding: 1rem;
        border-radius: 10px;
        text-align: center;
        margin: 1rem 0;
        cursor: pointer;
        transition: all 0.3s;
    }
    .agent-button:hover {
        background-color: #1E40AF;
        transform: translateY(-2px);
        box-shadow: 0 4px 6px rgba(0,0,0,0.1);
    }
    .result-container {
        background-color: #F3F4F6;
        padding: 1.5rem;
        border-radius: 10px;
        margin-top: 1.5rem;
        border-left: 5px solid #3B82F6;
    }
    .sidebar-item {
        padding: 0.75rem;
        border-radius: 5px;
        margin-bottom: 0.5rem;
        cursor: pointer;
        transition: background-color 0.2s;
    }
    .sidebar-item:hover {
        background-color: #E5E7EB;
    }
    .sidebar-item.active {
        background-color: #DBEAFE;
        font-weight: bold;
    }
    .metric-card {
        background-color: #EFF6FF;
        padding: 1rem;
        border-radius: 8px;
        text-align: center;
        box-shadow: 0 2px 4px rgba(0,0,0,0.05);
    }
    .metric-value {
        font-size: 1.5rem;
        font-weight: bold;
        color: #1E40AF;
    }
    .metric-label {
        font-size: 0.9rem;
        color: #6B7280;
    }
</style>
""", unsafe_allow_html=True)

# Configuration
MODEL_NAME = "meta-llama/Meta-Llama-3-8B-Instruct"
HF_API_KEY = os.getenv("HF_API_KEY")
COMPLIANCE_RULES = {"prohibited_words": ["sue", "arrest", "threaten"], "max_installment_ratio": 0.3}

# Initialize Hugging Face client
@st.cache_resource
def get_client():
    return InferenceClient(
        model=MODEL_NAME,
        token=HF_API_KEY
    )

# Sample data (Mock Database)
DEBTORS = {
    1000: {  
        "debtor_id": 1000,
        "name": "Rami",
        "age": 35,
        "preferred_language": "English",
        "preferred_channel": "Email",
        "income": 60000,
        "outstanding_balance": 1500,
        "installments": 15,
        "last_payment_date": "2023-11-01",
        "dispute_history": ["Billing error on 2023-10-01"]
    },
    1001: {
        "debtor_id": 1001,
        "name": "Jane Doe",
        "age": 35,
        "preferred_language": "English",
        "preferred_channel": "Email",
        "income": 45000,
        "outstanding_balance": 2500,
        "installments": 12,
        "last_payment_date": "2023-10-15",
        "dispute_history": ["Billing error on 2023-09-01"]
    }
}

# Business Logic Functions
def optimize_channel(debtor_data: dict) -> dict:
    age = debtor_data["age"]
    return {
        "channel": "SMS" if age < 35 else "Email",
        "best_time": "18:00-20:00" if age < 35 else "09:00-12:00"
    }

def generate_message(debtor_data: dict) -> str:
    try:
        client = get_client()
        
        prompt = f"""<|begin_of_text|>
        <|start_header_id|>system<|end_header_id|>
        Generate a friendly debt collection message in {debtor_data['preferred_language']}.
        Avoid these words: {COMPLIANCE_RULES['prohibited_words']}.<|eot_id|>
        
        <|start_header_id|>user<|end_header_id|>
        Debtor: {debtor_data['name']}
        Last Payment: {debtor_data['last_payment_date']}
        Balance: ${debtor_data['outstanding_balance']}<|eot_id|>
        
        <|start_header_id|>assistant<|end_header_id|>
        """

        response = client.text_generation(
            prompt,
            max_new_tokens=200,
            temperature=0.7
        )
        print(response)
        return response.strip()
        
    except Exception as e:
        st.error(f"Error generating message: {str(e)}")
        return "Hello, this is a friendly reminder about your outstanding balance. Please contact us to discuss payment options."

def calculate_payment_plan(debtor_data: dict) -> dict:
    max_payment = debtor_data["income"] * COMPLIANCE_RULES["max_installment_ratio"]
    installments = debtor_data["installments"]
    return {
        "monthly_payment": round(debtor_data["outstanding_balance"] / installments, 2),
        "duration_months": installments,
        "due_date": (datetime.now() + timedelta(days=15)).strftime("%Y-%m-%d")
    }

def generate_negotiation_strategy(debtor_data: dict) -> str:
    try:
        client = get_client()
        
        prompt = f"""<|begin_of_text|>
        <|start_header_id|>system<|end_header_id|>
        Generate a negotiation strategy for a debt collection agent. Include personalized payment options and dispute handling recommendations.<|eot_id|>
        
        <|start_header_id|>user<|end_header_id|>
        Debtor: {debtor_data['name']}
        Income: ${debtor_data['income']}
        Balance: ${debtor_data['outstanding_balance']}
        Installments: {debtor_data['installments']}
        Dispute History: {debtor_data['dispute_history'] if debtor_data['dispute_history'] else "None"}<|eot_id|>
        
        <|start_header_id|>assistant<|end_header_id|>
        """

        response = client.text_generation(
            prompt,
            max_new_tokens=300,
            temperature=0.7
        )
        print(response)
        return response.strip()
        
    except Exception as e:
        st.error(f"Error generating negotiation strategy: {str(e)}")
        return "Recommend standard payment plan with flexible options. Address any disputes promptly and professionally."

def get_message_response(debtor_id):
    debtor = DEBTORS.get(debtor_id)
    if not debtor:
        return {"error": "Debtor not found"}
    
    channel_info = optimize_channel(debtor)
    return {
        "debtor_id": debtor_id,
        "message": generate_message(debtor),
        "channel": channel_info["channel"],
        "send_time": f"{datetime.now().date()}T{channel_info['best_time']}"
    }

def get_payment_plan_response(debtor_id):
    debtor = DEBTORS.get(debtor_id)
    if not debtor:
        return {"error": "Debtor not found"}
    
    return {
        "debtor_id": debtor_id,
        "payment_plan": calculate_payment_plan(debtor),
        "negotiation_strategy": generate_negotiation_strategy(debtor)
    }

# Initialize session state
if 'selected_debtor_id' not in st.session_state:
    st.session_state.selected_debtor_id = None
if 'message_response' not in st.session_state:
    st.session_state.message_response = None
if 'payment_plan' not in st.session_state:
    st.session_state.payment_plan = None

# Calculate total outstanding balance
total_outstanding = sum(debtor["outstanding_balance"] for debtor in DEBTORS.values())

# Check for API key
if not HF_API_KEY:
    st.error("Hugging Face API key not found. Please set the HF_API_KEY environment variable.")
    st.stop()

# Sidebar - Debtor selection
with st.sidebar:
    st.markdown("<h2 style='text-align: center;'>Debtors</h2>", unsafe_allow_html=True)
    st.markdown("<hr>", unsafe_allow_html=True)
    
    # Using buttons for debtor selection
    for debtor_id, debtor in DEBTORS.items():
        if st.button(f"{debtor['name']} (ID: {debtor_id})", key=f"select_{debtor_id}"):
            st.session_state.selected_debtor_id = debtor_id
            st.session_state.message_response = None
            st.session_state.payment_plan = None
    
    st.markdown("<hr>", unsafe_allow_html=True)
    
    # Metrics
    st.markdown("<h3 style='text-align: center;'>Summary</h3>", unsafe_allow_html=True)
    st.markdown(f"""
    <div class="metric-card">
        <div class="metric-value">{len(DEBTORS)}</div>
        <div class="metric-label">Active Debtors</div>
    </div>
    """, unsafe_allow_html=True)
    st.markdown(f"""
    <div class="metric-card" style="margin-top: 1rem;">
        <div class="metric-value">${total_outstanding}</div>
        <div class="metric-label">Total Outstanding</div>
    </div>
    """, unsafe_allow_html=True)
    
    # Model info
    st.markdown("<hr>", unsafe_allow_html=True)
    st.markdown("<h3 style='text-align: center;'>Model</h3>", unsafe_allow_html=True)
    st.markdown(f"**Using:** {MODEL_NAME}")

# Main content
st.markdown("<h1 class='main-header'>Debt Collection AI</h1>", unsafe_allow_html=True)

# Display selected debtor information
if st.session_state.selected_debtor_id:
    debtor = DEBTORS[st.session_state.selected_debtor_id]
    
    # Create columns for debtor info
    col1, col2, col3 = st.columns(3)
    with col1:
        st.markdown(f"### {debtor['name']}")
        st.markdown(f"**ID:** {debtor['debtor_id']}")
        st.markdown(f"**Age:** {debtor['age']}")
    with col2:
        st.markdown("### Contact Info")
        st.markdown(f"**Language:** {debtor['preferred_language']}")
        st.markdown(f"**Channel:** {debtor['preferred_channel']}")
    with col3:
        st.markdown("### Financial Info")
        st.markdown(f"**Balance:** ${debtor['outstanding_balance']}")
        st.markdown(f"**Installments:** {debtor['installments']}")
        st.markdown(f"**Last Payment:** {debtor['last_payment_date']}")
    
    st.markdown("<hr>", unsafe_allow_html=True)
    
    # Action buttons
    col1, col2 = st.columns(2)
    
    with col1:
        st.markdown("""
        <div class="agent-button" id="message-btn">
            <h3>Agent 1: Empathy-Driven Communicator</h3>
            <p>Create an empathetic collection message</p>
        </div>
        """, unsafe_allow_html=True)
        if st.button("Generate Message", key="gen_message"):
            with st.spinner("Agent 1 is generating a message..."):
                # Get response
                response = get_message_response(st.session_state.selected_debtor_id)
                if response and "error" not in response:
                    st.session_state.message_response = response
                    st.session_state.payment_plan = None
                else:
                    st.error(response.get("error", "Unknown error occurred"))
    
    with col2:
        st.markdown("""
        <div class="agent-button" id="payment-btn">
            <h3>Agent 2: Negotiation & Payment Planner</h3>
            <p>Create a personalized payment plan</p>
        </div>
        """, unsafe_allow_html=True)
        if st.button("Calculate Plan", key="calc_plan"):
            with st.spinner("Agent 2 is analyzing and planning..."):
                # Get response
                response = get_payment_plan_response(st.session_state.selected_debtor_id)
                if response and "error" not in response:
                    st.session_state.payment_plan = response
                    st.session_state.message_response = None
                else:
                    st.error(response.get("error", "Unknown error occurred"))
    
    # Display results
    if st.session_state.message_response:
        st.markdown("<div class='result-container'>", unsafe_allow_html=True)
        st.markdown("### Agent 1: Empathy-Driven Communication")
        st.markdown(f"**Channel:** {st.session_state.message_response['channel']}")
        st.markdown(f"**Best Send Time:** {st.session_state.message_response['send_time'].split('T')[1]}")
        st.markdown("#### Message Content:")
        st.markdown(f"```\n{st.session_state.message_response['message']}\n```")
        st.markdown("</div>", unsafe_allow_html=True)
    
    elif st.session_state.payment_plan:
        st.markdown("<div class='result-container'>", unsafe_allow_html=True)
        st.markdown("### Agent 2: Negotiation & Payment Plan")
        plan = st.session_state.payment_plan['payment_plan']
        st.markdown(f"**Monthly Payment:** ${plan['monthly_payment']}")
        st.markdown(f"**Duration:** {plan['duration_months']} months")
        st.markdown(f"**First Payment Due:** {plan['due_date']}")
        
        st.markdown("#### Negotiation Strategy:")
        st.markdown(f"{st.session_state.payment_plan['negotiation_strategy']}")
        st.markdown("</div>", unsafe_allow_html=True)

else:
    st.info("👈 Please select a debtor from the sidebar to get started.")

# Footer
st.markdown("<hr>", unsafe_allow_html=True)
st.markdown("<p style='text-align: center; color: #6B7280;'>© 2025 Debt Collection AI System</p>", unsafe_allow_html=True)