"""
Data Analyst Agent - Streamlit Version
Beautiful UI with Voice Input
"""

import streamlit as st
import asyncio
import os
import base64
import tempfile
from fastapi_poe.client import get_bot_response
from fastapi_poe.types import ProtocolMessage, Attachment

# Page config
st.set_page_config(
    page_title="📊 Data Analyst Agent",
    page_icon="📊",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for beautiful UI
st.markdown("""
<style>
    .main-header {
        text-align: center;
        padding: 1.5rem;
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        border-radius: 16px;
        margin-bottom: 2rem;
        color: white;
    }
    .main-header h1 {
        margin: 0;
        font-size: 2.5rem;
    }
    .main-header p {
        margin: 0.5rem 0 0 0;
        opacity: 0.9;
    }
    .sidebar-card {
        background: linear-gradient(135deg, #f5f7fa 0%, #e4e8ec 100%);
        border-radius: 12px;
        padding: 1rem;
        margin-bottom: 1rem;
    }
    .voice-card {
        background: linear-gradient(135deg, #a8edea 0%, #fed6e3 100%);
        border-radius: 12px;
        padding: 1rem;
        margin-bottom: 1rem;
    }
    .stButton > button {
        border-radius: 12px;
        font-weight: 600;
    }
    .chat-message {
        padding: 1rem;
        border-radius: 12px;
        margin-bottom: 1rem;
    }
    .user-message {
        background: #e3f2fd;
        border-left: 4px solid #2196f3;
    }
    .assistant-message {
        background: #f5f5f5;
        border-left: 4px solid #667eea;
    }
</style>
""", unsafe_allow_html=True)

# Get API key
POE_API_KEY = os.environ.get("POE_API_KEY", "")

SYSTEM_PROMPT = """You are a powerful data analysis agent. You can:

1. **Analyze Data**: Read CSV, Excel (.xlsx, .xls) files and perform comprehensive analysis
2. **Generate Visualizations**: Create charts using matplotlib, seaborn, plotly
3. **Execute Python Code**: Run any Python code for data manipulation and analysis
4. **Install Libraries**: Use `pip install` if a library is not available
5. **Create Reports**: Generate Word (.docx) or PowerPoint (.pptx) documents with charts and summaries

## Guidelines
- When user uploads a file, first explore the data (shape, columns, dtypes, sample rows, missing values)
- Create meaningful visualizations based on the data types and relationships
- Always display charts inline so the user can see them
- For reports, save charts as images first, then embed them in Word/PowerPoint
- Be proactive: suggest insights and additional analyses the user might find valuable
- When creating downloadable files, provide the download link"""


async def transcribe_audio(audio_bytes: bytes) -> str:
    """Transcribe audio using Whisper via Poe API."""
    if not audio_bytes:
        return ""
    try:
        b64_data = base64.b64encode(audio_bytes).decode("utf-8")
        data_url = f"data:audio/wav;base64,{b64_data}"
        attachment = Attachment(url=data_url, name="voice.wav", content_type="audio/wav")
        messages = [ProtocolMessage(role="user", content="Transcribe this audio accurately.", attachments=[attachment])]
        text = ""
        async for partial in get_bot_response(messages=messages, bot_name="Whisper-V3-Large-T", api_key=POE_API_KEY):
            text += partial.text
        return text.strip()
    except Exception as e:
        return f"Transcription error: {str(e)}"


async def call_analyst(message: str, file_bytes: bytes = None, filename: str = None, history: list = None):
    """Call Claude-Code for analysis."""
    if not POE_API_KEY:
        return "❌ **Error**: POE_API_KEY not set! Go to Settings → Secrets and add your Poe API key."

    messages = [ProtocolMessage(role="system", content=SYSTEM_PROMPT)]

    # Add history
    if history:
        for item in history:
            if item["role"] == "user":
                messages.append(ProtocolMessage(role="user", content=item["content"]))
            elif item["role"] == "assistant":
                messages.append(ProtocolMessage(role="assistant", content=item["content"]))

    # Handle file
    attachments = []
    if file_bytes and filename:
        if filename.endswith(".csv"):
            ctype = "text/csv"
        elif filename.endswith(".xlsx"):
            ctype = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
        else:
            ctype = "application/octet-stream"

        b64 = base64.b64encode(file_bytes).decode("utf-8")
        attachments.append(Attachment(url=f"data:{ctype};base64,{b64}", name=filename, content_type=ctype))

    messages.append(ProtocolMessage(role="user", content=message, attachments=attachments))

    response = ""
    try:
        async for partial in get_bot_response(messages=messages, bot_name="Claude-Code", api_key=POE_API_KEY):
            response += partial.text
    except Exception as e:
        response = f"❌ **API Error**: {str(e)}"

    return response


def run_async(coro):
    """Run async function."""
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        return loop.run_until_complete(coro)
    finally:
        loop.close()


# Initialize session state
if "messages" not in st.session_state:
    st.session_state.messages = []
if "file_bytes" not in st.session_state:
    st.session_state.file_bytes = None
if "filename" not in st.session_state:
    st.session_state.filename = None

# Header
st.markdown("""
<div class="main-header">
    <h1>📊 Data Analyst Agent</h1>
    <p>Upload your data • Ask questions • Get insights, charts & reports</p>
</div>
""", unsafe_allow_html=True)

# Sidebar
with st.sidebar:
    st.markdown("### 📁 Upload Your Data")
    uploaded_file = st.file_uploader(
        "Drop CSV or Excel file here",
        type=["csv", "xlsx", "xls"],
        help="Upload a CSV or Excel file to analyze"
    )

    if uploaded_file:
        st.session_state.file_bytes = uploaded_file.read()
        st.session_state.filename = uploaded_file.name
        st.success(f"✅ Loaded: {uploaded_file.name}")
        uploaded_file.seek(0)  # Reset for potential re-read

    st.markdown("---")

    st.markdown("### 💡 Example Prompts")
    st.markdown("""
    - Analyze this data and show key insights
    - Create a bar chart of sales by category
    - Show me a correlation heatmap
    - Find the top 10 records by value
    - Generate a PowerPoint summary
    - Create a Word report with charts
    """)

    st.markdown("---")

    st.markdown("### 🎤 Voice Input")
    audio_file = st.file_uploader(
        "Upload audio recording",
        type=["wav", "mp3", "m4a", "ogg"],
        help="Record audio on your device, then upload it here"
    )

    if audio_file:
        audio_bytes = audio_file.read()
        st.audio(audio_bytes, format=f"audio/{audio_file.type.split('/')[-1]}")

        if st.button("🎯 Transcribe Audio", use_container_width=True):
            with st.spinner("Transcribing..."):
                transcription = run_async(transcribe_audio(audio_bytes))
                if transcription:
                    st.session_state.voice_text = transcription
                    st.success("✅ Transcribed!")
                    st.info(f"📝 {transcription}")

    st.markdown("---")

    if st.button("🗑️ Clear Chat History", use_container_width=True):
        st.session_state.messages = []
        st.rerun()

# Main chat area
st.markdown("### 💬 Analysis Chat")

# Display chat history
for msg in st.session_state.messages:
    if msg["role"] == "user":
        st.markdown(f"""
        <div class="chat-message user-message">
            <strong>👤 You:</strong><br>{msg["content"]}
        </div>
        """, unsafe_allow_html=True)
    else:
        st.markdown(f"""
        <div class="chat-message assistant-message">
            <strong>🤖 Agent:</strong><br>{msg["content"]}
        </div>
        """, unsafe_allow_html=True)

# Input area
col1, col2 = st.columns([5, 1])

with col1:
    # Check if we have voice text to use
    default_text = st.session_state.get("voice_text", "")
    user_input = st.text_area(
        "Your message",
        value=default_text,
        placeholder="Ask about your data, request charts, or generate reports...",
        height=100,
        label_visibility="collapsed"
    )
    # Clear voice text after using
    if default_text and user_input == default_text:
        st.session_state.voice_text = ""

with col2:
    st.write("")  # Spacer
    st.write("")  # Spacer
    analyze_clicked = st.button("🔍 Analyze", use_container_width=True, type="primary")

# Process input
if analyze_clicked and user_input.strip():
    # Add user message
    display_msg = user_input
    if st.session_state.filename:
        display_msg = f"📎 **{st.session_state.filename}**\n\n{user_input}"

    st.session_state.messages.append({"role": "user", "content": display_msg})

    # Get response
    with st.spinner("🔍 Analyzing..."):
        response = run_async(call_analyst(
            user_input,
            st.session_state.file_bytes,
            st.session_state.filename,
            st.session_state.messages[:-1]  # Exclude the message we just added
        ))

    # Add assistant response
    st.session_state.messages.append({"role": "assistant", "content": response})

    # Clear the file after first use (optional - comment out if you want to keep it)
    # st.session_state.file_bytes = None
    # st.session_state.filename = None

    st.rerun()

# Footer
st.markdown("---")
st.markdown(
    "<p style='text-align: center; color: #666;'>Powered by <a href='https://poe.com'>Poe API</a> • Claude-Code for intelligent analysis</p>",
    unsafe_allow_html=True
)