import streamlit as st
import requests
import json
import time
import re
import random
import string
from datetime import datetime
import html

# Set page configuration
st.set_page_config(
    page_title="Temp Mail Client",
    page_icon="✉️",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for styling
st.markdown("""
<style>
    .main {
        background-color: #f9f9f9;
    }
    .email-list {
        border: 1px solid #ddd;
        border-radius: 5px;
        padding: 10px;
        margin-bottom: 10px;
        cursor: pointer;
        transition: background-color 0.3s;
    }
    .email-list:hover {
        background-color: #f0f0f0;
    }
    .email-selected {
        background-color: #e6f7ff;
        border-left: 3px solid #1890ff;
    }
    .email-header {
        font-weight: bold;
        margin-bottom: 5px;
    }
    .email-subject {
        font-size: 16px;
        white-space: nowrap;
        overflow: hidden;
        text-overflow: ellipsis;
    }
    .email-from {
        font-size: 14px;
        color: #666;
    }
    .email-time {
        font-size: 12px;
        color: #999;
        text-align: right;
    }
    .email-preview {
        font-size: 14px;
        color: #666;
        white-space: nowrap;
        overflow: hidden;
        text-overflow: ellipsis;
    }
    .unread {
        font-weight: bold;
    }
    .email-content {
        border: 1px solid #ddd;
        border-radius: 5px;
        padding: 20px;
        background-color: white;
    }
    .email-actions {
        display: flex;
        justify-content: flex-end;
        margin-bottom: 10px;
    }
    .stButton button {
        padding: 2px 10px;
        font-size: 14px;
    }
    .attachment {
        border: 1px solid #ddd;
        border-radius: 5px;
        padding: 10px;
        margin: 5px 0;
        display: flex;
        align-items: center;
        background-color: #f5f5f5;
    }
    .attachment-icon {
        margin-right: 10px;
        color: #1890ff;
    }
    .attachment-name {
        flex-grow: 1;
    }
    .attachment-size {
        color: #999;
        margin-right: 10px;
    }
    .sidebar-content {
        padding: 10px;
    }
    .provider-select {
        margin-bottom: 20px;
    }
    .refresh-button {
        width: 100%;
        margin-bottom: 10px;
    }
    .copy-button {
        width: 100%;
    }
    .header-container {
        display: flex;
        justify-content: space-between;
        align-items: center;
        margin-bottom: 20px;
    }
    .app-header {
        font-size: 24px;
        font-weight: bold;
        color: #1890ff;
    }
    .loading-spinner {
        display: flex;
        justify-content: center;
        align-items: center;
        height: 100px;
    }
</style>
""", unsafe_allow_html=True)

# Initialize session state variables
if 'provider' not in st.session_state:
    st.session_state.provider = "mail.tm"
if 'domains' not in st.session_state:
    st.session_state.domains = []
if 'email' not in st.session_state:
    st.session_state.email = ""
if 'password' not in st.session_state:
    st.session_state.password = ""
if 'token' not in st.session_state:
    st.session_state.token = ""
if 'account_id' not in st.session_state:
    st.session_state.account_id = ""
if 'messages' not in st.session_state:
    st.session_state.messages = []
if 'selected_message' not in st.session_state:
    st.session_state.selected_message = None
if 'message_content' not in st.session_state:
    st.session_state.message_content = None
if 'last_refresh' not in st.session_state:
    st.session_state.last_refresh = datetime.now()

# API configuration
API_ENDPOINTS = {
    "mail.tm": {
        "base_url": "https://api.mail.tm",
        "mercure_url": "https://mercure.mail.tm/.well-known/mercure"
    },
    "mail.gw": {
        "base_url": "https://api.mail.gw",
        "mercure_url": "https://api.mail.gw/.well-known/mercure"
    }
}

def get_base_url():
    return API_ENDPOINTS[st.session_state.provider]["base_url"]

def get_mercure_url():
    return API_ENDPOINTS[st.session_state.provider]["mercure_url"]

def get_domains():
    """Fetch available domains from the API"""
    try:
        response = requests.get(f"{get_base_url()}/domains")
        if response.status_code == 200:
            data = response.json()
            st.session_state.domains = [domain["domain"] for domain in data["hydra:member"]]
            return st.session_state.domains
        else:
            st.error(f"Failed to fetch domains: {response.status_code}")
            return []
    except Exception as e:
        st.error(f"Error fetching domains: {str(e)}")
        return []

def create_account(email, password):
    """Create a new email account"""
    try:
        payload = {
            "address": email,
            "password": password
        }
        response = requests.post(
            f"{get_base_url()}/accounts",
            json=payload
        )
        
        if response.status_code == 201:
            data = response.json()
            st.session_state.account_id = data["id"]
            return True
        else:
            st.error(f"Failed to create account: {response.status_code} - {response.text}")
            return False
    except Exception as e:
        st.error(f"Error creating account: {str(e)}")
        return False

def get_token(email, password):
    """Get authentication token"""
    try:
        payload = {
            "address": email,
            "password": password
        }
        response = requests.post(
            f"{get_base_url()}/token",
            json=payload
        )
        
        if response.status_code == 200:
            data = response.json()
            st.session_state.token = data["token"]
            st.session_state.account_id = data["id"]
            return True
        else:
            st.error(f"Failed to get token: {response.status_code} - {response.text}")
            return False
    except Exception as e:
        st.error(f"Error getting token: {str(e)}")
        return False

def get_account_info():
    """Get account information"""
    if not st.session_state.token:
        return None
    
    try:
        headers = {"Authorization": f"Bearer {st.session_state.token}"}
        response = requests.get(
            f"{get_base_url()}/me",
            headers=headers
        )
        
        if response.status_code == 200:
            return response.json()
        else:
            st.error(f"Failed to get account info: {response.status_code}")
            return None
    except Exception as e:
        st.error(f"Error getting account info: {str(e)}")
        return None

def get_messages():
    """Get messages for the current account"""
    if not st.session_state.token:
        return []
    
    try:
        headers = {"Authorization": f"Bearer {st.session_state.token}"}
        response = requests.get(
            f"{get_base_url()}/messages",
            headers=headers
        )
        
        if response.status_code == 200:
            data = response.json()
            st.session_state.messages = data["hydra:member"]
            return st.session_state.messages
        else:
            st.error(f"Failed to get messages: {response.status_code}")
            return []
    except Exception as e:
        st.error(f"Error getting messages: {str(e)}")
        return []

def get_message_content(message_id):
    """Get detailed content of a specific message"""
    if not st.session_state.token:
        return None
    
    try:
        headers = {"Authorization": f"Bearer {st.session_state.token}"}
        response = requests.get(
            f"{get_base_url()}/messages/{message_id}",
            headers=headers
        )
        
        if response.status_code == 200:
            return response.json()
        else:
            st.error(f"Failed to get message content: {response.status_code}")
            return None
    except Exception as e:
        st.error(f"Error getting message content: {str(e)}")
        return None

def mark_as_read(message_id):
    """Mark a message as read"""
    if not st.session_state.token:
        return False
    
    try:
        headers = {"Authorization": f"Bearer {st.session_state.token}"}
        response = requests.patch(
            f"{get_base_url()}/messages/{message_id}",
            headers=headers
        )
        
        if response.status_code == 200:
            # Update the message in the session state
            for i, msg in enumerate(st.session_state.messages):
                if msg["id"] == message_id:
                    st.session_state.messages[i]["seen"] = True
            return True
        else:
            st.error(f"Failed to mark message as read: {response.status_code}")
            return False
    except Exception as e:
        st.error(f"Error marking message as read: {str(e)}")
        return False

def delete_message(message_id):
    """Delete a message"""
    if not st.session_state.token:
        return False
    
    try:
        headers = {"Authorization": f"Bearer {st.session_state.token}"}
        response = requests.delete(
            f"{get_base_url()}/messages/{message_id}",
            headers=headers
        )
        
        if response.status_code == 204:
            # Remove the message from the session state
            st.session_state.messages = [msg for msg in st.session_state.messages if msg["id"] != message_id]
            if st.session_state.selected_message == message_id:
                st.session_state.selected_message = None
                st.session_state.message_content = None
            return True
        else:
            st.error(f"Failed to delete message: {response.status_code}")
            return False
    except Exception as e:
        st.error(f"Error deleting message: {str(e)}")
        return False

def delete_account():
    """Delete the current account"""
    if not st.session_state.token or not st.session_state.account_id:
        return False
    
    try:
        headers = {"Authorization": f"Bearer {st.session_state.token}"}
        response = requests.delete(
            f"{get_base_url()}/accounts/{st.session_state.account_id}",
            headers=headers
        )
        
        if response.status_code == 204:
            # Clear session state
            st.session_state.email = ""
            st.session_state.password = ""
            st.session_state.token = ""
            st.session_state.account_id = ""
            st.session_state.messages = []
            st.session_state.selected_message = None
            st.session_state.message_content = None
            return True
        else:
            st.error(f"Failed to delete account: {response.status_code}")
            return False
    except Exception as e:
        st.error(f"Error deleting account: {str(e)}")
        return False

def format_date(date_str):
    """Format date string to a more readable format"""
    try:
        dt = datetime.fromisoformat(date_str.replace('Z', '+00:00'))
        now = datetime.now()
        
        # If today, show only time
        if dt.date() == now.date():
            return dt.strftime("%H:%M")
        # If this year, show month and day
        elif dt.year == now.year:
            return dt.strftime("%b %d")
        # Otherwise show full date
        else:
            return dt.strftime("%Y-%m-%d")
    except:
        return date_str

def generate_random_username(length=10):
    """Generate a random username for email"""
    letters = string.ascii_lowercase + string.digits
    return ''.join(random.choice(letters) for i in range(length))

def format_size(size_bytes):
    """Format file size in bytes to human-readable format"""
    if size_bytes < 1024:
        return f"{size_bytes} B"
    elif size_bytes < 1024 * 1024:
        return f"{size_bytes/1024:.1f} KB"
    else:
        return f"{size_bytes/(1024*1024):.1f} MB"

def sanitize_html(html_content):
    """Basic sanitization of HTML content"""
    # Remove potentially dangerous tags and attributes
    # This is a very basic implementation - in production, use a proper HTML sanitizer
    dangerous_tags = ['script', 'iframe', 'object', 'embed']
    for tag in dangerous_tags:
        html_content = re.sub(f'<{tag}.*?</{tag}>', '', html_content, flags=re.DOTALL)
        html_content = re.sub(f'<{tag}.*?>', '', html_content, flags=re.DOTALL)
    
    # Remove on* attributes
    html_content = re.sub(r'on\w+=".*?"', '', html_content)
    html_content = re.sub(r"on\w+='.*?'", '', html_content)
    
    return html_content

# Main application layout
def main():
    # Sidebar for account management
    with st.sidebar:
        st.markdown("<div class='sidebar-content'>", unsafe_allow_html=True)
        
        # Provider selection
        st.markdown("<div class='provider-select'>", unsafe_allow_html=True)
        provider = st.selectbox(
            "Select Provider",
            options=["mail.tm", "mail.gw"],
            index=0 if st.session_state.provider == "mail.tm" else 1,
            key="provider_select"
        )
        st.markdown("</div>", unsafe_allow_html=True)
        
        if provider != st.session_state.provider:
            st.session_state.provider = provider
            st.session_state.domains = []
            st.session_state.email = ""
            st.session_state.password = ""
            st.session_state.token = ""
            st.session_state.account_id = ""
            st.session_state.messages = []
            st.session_state.selected_message = None
            st.session_state.message_content = None
            get_domains()
        
        # Account creation/login section
        if not st.session_state.token:
            st.subheader("Create Email Account")
            
            if not st.session_state.domains:
                get_domains()
            
            if st.session_state.domains:
                username = st.text_input("Username (optional)", 
                                        value=generate_random_username(),
                                        help="Leave blank to generate random username")
                
                domain = st.selectbox("Domain", options=st.session_state.domains)
                
                if not username:
                    username = generate_random_username()
                
                email = f"{username}@{domain}"
                st.session_state.email = email
                
                password = st.text_input("Password", 
                                        value=''.join(random.choices(string.ascii_letters + string.digits, k=12)),
                                        type="password")
                st.session_state.password = password
                
                col1, col2 = st.columns(2)
                with col1:
                    if st.button("Create Account", use_container_width=True):
                        with st.spinner("Creating account..."):
                            if create_account(email, password):
                                if get_token(email, password):
                                    st.success("Account created successfully!")
                                    st.experimental_rerun()
                
                with col2:
                    if st.button("Login", use_container_width=True):
                        with st.spinner("Logging in..."):
                            if get_token(email, password):
                                st.success("Logged in successfully!")
                                st.experimental_rerun()
            else:
                st.warning("Unable to fetch domains. Please try again later.")
        
        # Account info and management
        else:
            account_info = get_account_info()
            if account_info:
                st.subheader("Account Information")
                st.markdown(f"**Email:** {account_info['address']}")
                
                # Copy email button
                if st.button("Copy Email", key="copy_email", use_container_width=True):
                    st.write("Email copied to clipboard!")
                    # Note: This doesn't actually copy to clipboard in Streamlit,
                    # but we're simulating the UI behavior
                
                # Display quota information
                quota_used = account_info.get('used', 0)
                quota_total = account_info.get('quota', 0)
                if quota_total > 0:
                    quota_percentage = (quota_used / quota_total) * 100
                    st.progress(quota_percentage / 100)
                    st.markdown(f"Storage: {quota_used}/{quota_total} ({quota_percentage:.1f}%)")
                
                # Refresh button
                if st.button("Refresh Inbox", key="refresh_inbox", use_container_width=True):
                    with st.spinner("Refreshing..."):
                        get_messages()
                        st.session_state.last_refresh = datetime.now()
                        st.experimental_rerun()
                
                # Logout and delete account buttons
                col1, col2 = st.columns(2)
                with col1:
                    if st.button("Logout", use_container_width=True):
                        st.session_state.token = ""
                        st.session_state.account_id = ""
                        st.session_state.messages = []
                        st.session_state.selected_message = None
                        st.session_state.message_content = None
                        st.experimental_rerun()
                
                with col2:
                    if st.button("Delete Account", use_container_width=True):
                        if delete_account():
                            st.success("Account deleted successfully!")
                            st.experimental_rerun()
        
        st.markdown("</div>", unsafe_allow_html=True)
    
    # Main content area
    st.markdown("<div class='header-container'>", unsafe_allow_html=True)
    st.markdown("<div class='app-header'>✉️ Temp Mail Client</div>", unsafe_allow_html=True)
    st.markdown(f"<div>Last refreshed: {st.session_state.last_refresh.strftime('%H:%M:%S')}</div>", unsafe_allow_html=True)
    st.markdown("</div>", unsafe_allow_html=True)
    
    # If logged in, show messages
    if st.session_state.token:
        # Get messages if not already loaded
        if not st.session_state.messages:
            with st.spinner("Loading messages..."):
                get_messages()
        
        # Split the screen into two columns for message list and content
        col1, col2 = st.columns([1, 2])
        
        with col1:
            st.subheader("Inbox")
            
            # Display message list
            if st.session_state.messages:
                for message in st.session_state.messages:
                    # Determine if message is selected
                    is_selected = st.session_state.selected_message == message["id"]
                    # Determine if message is unread
                    is_unread = not message.get("seen", False)
                    
                    # Create a clickable message container
                    message_class = "email-list"
                    if is_selected:
                        message_class += " email-selected"
                    
                    st.markdown(f"<div class='{message_class}' onclick='null' id='message-{message['id']}'>", unsafe_allow_html=True)
                    
                    # Message header (from and time)
                    st.markdown("<div class='email-header'>", unsafe_allow_html=True)
                    
                    # From name/address
                    from_name = message.get("from", {}).get("name", "")
                    from_address = message.get("from", {}).get("address", "Unknown")
                    display_from = from_name if from_name else from_address
                    
                    # Format date
                    date_str = format_date(message.get("createdAt", ""))
                    
                    col_from, col_date = st.columns([3, 1])
                    with col_from:
                        st.markdown(f"<div class='email-from{' unread' if is_unread else ''}'>{display_from}</div>", unsafe_allow_html=True)
                    with col_date:
                        st.markdown(f"<div class='email-time'>{date_str}</div>", unsafe_allow_html=True)
                    
                    st.markdown("</div>", unsafe_allow_html=True)
                    
                    # Subject
                    subject = message.get("subject", "(No subject)")
                    st.markdown(f"<div class='email-subject{' unread' if is_unread else ''}'>{subject}</div>", unsafe_allow_html=True)
                    
                    # Preview
                    preview = message.get("intro", "")
                    st.markdown(f"<div class='email-preview'>{preview}</div>", unsafe_allow_html=True)
                    
                    # Close the message container
                    st.markdown("</div>", unsafe_allow_html=True)
                    
                    # Handle click on message
                    if st.button("View", key=f"view_{message['id']}", help="View this message"):
                        st.session_state.selected_message = message["id"]
                        st.session_state.message_content = get_message_content(message["id"])
                        if not message.get("seen", False):
                            mark_as_read(message["id"])
                        st.experimental_rerun()
            else:
                st.info("No messages yet. Emails sent to your address will appear here.")
        
        with col2:
            # Display message content if a message is selected
            if st.session_state.selected_message and st.session_state.message_content:
                message = st.session_state.message_content
                
                # Message actions
                st.markdown("<div class='email-actions'>", unsafe_allow_html=True)
                
                col1, col2 = st.columns([1, 1])
                with col1:
                    if st.button("Back to Inbox", key="back_to_inbox"):
                        st.session_state.selected_message = None
                        st.session_state.message_content = None
                        st.experimental_rerun()
                
                with col2:
                    if st.button("Delete", key="delete_message"):
                        if delete_message(st.session_state.selected_message):
                            st.success("Message deleted!")
                            st.session_state.selected_message = None
                            st.session_state.message_content = None
                            st.experimental_rerun()
                
                st.markdown("</div>", unsafe_allow_html=True)
                
                # Message content container
                st.markdown("<div class='email-content'>", unsafe_allow_html=True)
                
                # Subject
                subject = message.get("subject", "(No subject)")
                st.markdown(f"<h2>{subject}</h2>", unsafe_allow_html=True)
                
                # From
                from_name = message.get("from", {}).get("name", "")
                from_address = message.get("from", {}).get("address", "Unknown")
                display_from = f"{from_name} <{from_address}>" if from_name else from_address
                st.markdown(f"<p><strong>From:</strong> {display_from}</p>", unsafe_allow_html=True)
                
                # To
                to_addresses = []
                for to in message.get("to", []):
                    name = to.get("name", "")
                    address = to.get("address", "")
                    if name:
                        to_addresses.append(f"{name} <{address}>")
                    else:
                        to_addresses.append(address)
                
                st.markdown(f"<p><strong>To:</strong> {', '.join(to_addresses)}</p>", unsafe_allow_html=True)
                
                # Date
                date_str = message.get("createdAt", "")
                if date_str:
                    try:
                        dt = datetime.fromisoformat(date_str.replace('Z', '+00:00'))
                        formatted_date = dt.strftime("%a, %d %b %Y %H:%M:%S")
                        st.markdown(f"<p><strong>Date:</strong> {formatted_date}</p>", unsafe_allow_html=True)
                    except:
                        st.markdown(f"<p><strong>Date:</strong> {date_str}</p>", unsafe_allow_html=True)
                
                # Attachments
                attachments = message.get("attachments", [])
                if attachments:
                    st.markdown("<div><strong>Attachments:</strong></div>", unsafe_allow_html=True)
                    for attachment in attachments:
                        filename = attachment.get("filename", "Unknown")
                        size = attachment.get("size", 0)
                        content_type = attachment.get("contentType", "application/octet-stream")
                        download_url = attachment.get("downloadUrl", "")
                        
                        st.markdown(f"""
                        <div class='attachment'>
                            <div class='attachment-icon'>📎</div>
                            <div class='attachment-name'>{filename}</div>
                            <div class='attachment-size'>{format_size(size)}</div>
                        </div>
                        """, unsafe_allow_html=True)
                
                # Message body
                st.markdown("<hr>", unsafe_allow_html=True)
                
                # Try to use HTML content first, fall back to text
                html_content = message.get("html", [""])[0] if message.get("html") else ""
                text_content = message.get("text", "")
                
                if html_content:
                    # Sanitize HTML content
                    safe_html = sanitize_html(html_content)
                    st.markdown(f"<div>{safe_html}</div>", unsafe_allow_html=True)
                elif text_content:
                    # Display plain text with line breaks preserved
                    st.markdown(f"<pre style='white-space: pre-wrap;'>{html.escape(text_content)}</pre>", unsafe_allow_html=True)
                else:
                    st.markdown("<p>(No content)</p>", unsafe_allow_html=True)
                
                st.markdown("</div>", unsafe_allow_html=True)
            else:
                st.info("Select a message from the inbox to view its content.")
    else:
        # If not logged in, show welcome message
        st.markdown("""
        <div style="text-align: center; padding: 50px 20px;">
            <h1>Welcome to Temp Mail Client</h1>
            <p style="font-size: 18px; margin: 20px 0;">
                Create a temporary email address to protect your privacy.
            </p>
            <div style="max-width: 600px; margin: 0 auto; text-align: left; background-color: #f5f5f5; padding: 20px; border-radius: 10px;">
                <h3>How it works:</h3>
                <ol style="font-size: 16px;">
                    <li>Create a temporary email account using the sidebar</li>
                    <li>Use this email address for sign-ups or verification</li>
                    <li>Receive emails instantly in this interface</li>
                    <li>Delete the account when you're done</li>
                </ol>
            </div>
        </div>
        """, unsafe_allow_html=True)

if __name__ == "__main__":
    main()