import streamlit as st
from groq import Groq
import requests
import re
import numpy as np
from bs4 import BeautifulSoup
import PyPDF2
import docx
from io import StringIO
import csv
import json
from datetime import datetime
import pytz
import time
import hashlib
from collections import defaultdict
import os

# ============================================
# DOCUMENT GENERATION FUNCTIONS
# ============================================

from pptx import Presentation
from pptx.util import Inches, Pt
from pptx.enum.text import PP_ALIGN
from pptx.dml.color import RGBColor
from docx import Document as WordDocument
from docx.shared import Inches as DocInches, Pt as DocPt
from docx.enum.text import WD_ALIGN_PARAGRAPH
import io
import base64

def create_ppt_from_content(title, content, filename="presentation"):
    """Create a PowerPoint presentation from content - properly split across slides"""
    try:
        prs = Presentation()
        
        # Title slide
        title_slide_layout = prs.slide_layouts[0]
        slide = prs.slides.add_slide(title_slide_layout)
        slide.shapes.title.text = title[:100]
        slide.placeholders[1].text = f"Created by MozeAI\n{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"
        
        # Content slides layout
        content_slide_layout = prs.slide_layouts[1]
        
        # Split content into slides based on headings or paragraphs
        lines = content.split('\n')
        
        current_slide = None
        current_text_frame = None
        current_title = None
        
        for line in lines:
            line = line.strip()
            if not line:
                continue
            
            # Check if this line looks like a slide title
            is_title = False
            if len(line) < 60 and (line.endswith(':') or line.isupper() or re.match(r'^\d+\.', line) or line[0].isupper() and len(line) < 40):
                is_title = True
            
            if is_title:
                # Create new slide for this title
                current_slide = prs.slides.add_slide(content_slide_layout)
                clean_title = line.rstrip(':')
                current_slide.shapes.title.text = clean_title[:100]
                content_box = current_slide.placeholders[1]
                current_text_frame = content_box.text_frame
                current_text_frame.text = ""
                current_title = clean_title
            else:
                # If no slide exists yet, create one
                if current_slide is None:
                    current_slide = prs.slides.add_slide(content_slide_layout)
                    current_slide.shapes.title.text = "Information"
                    content_box = current_slide.placeholders[1]
                    current_text_frame = content_box.text_frame
                    current_text_frame.text = ""
                
                # Add as bullet point
                if current_text_frame:
                    p = current_text_frame.add_paragraph()
                    p.text = line[:150]
                    p.font.size = Pt(18)
                    p.level = 0
                    p.space_after = Pt(6)
        
        # If no content slides were created, add a default one
        if len(prs.slides) == 1:
            slide = prs.slides.add_slide(content_slide_layout)
            slide.shapes.title.text = "Content Summary"
            content_box = slide.placeholders[1]
            text_frame = content_box.text_frame
            text_frame.text = content[:500]
        
        ppt_bytes = io.BytesIO()
        prs.save(ppt_bytes)
        ppt_bytes.seek(0)
        return ppt_bytes
    except Exception as e:
        print(f"PPT creation error: {e}")
        return None

def create_word_from_content(title, content, filename="document"):
    """Create a Word document from content"""
    try:
        doc = WordDocument()
        
        title_heading = doc.add_heading(title, 0)
        title_heading.alignment = WD_ALIGN_PARAGRAPH.CENTER
        
        doc.add_paragraph(f"Generated by MozeAI on {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        doc.add_paragraph()
        
        paragraphs = content.split('\n\n')
        for para in paragraphs:
            if para.strip():
                p = doc.add_paragraph(para.strip())
                p.style.font.size = DocPt(12)
        
        word_bytes = io.BytesIO()
        doc.save(word_bytes)
        word_bytes.seek(0)
        return word_bytes
    except Exception as e:
        return None

def create_real_excel_file(title, data_rows):
    """Create a REAL .xlsx Excel file with proper formatting"""
    try:
        from openpyxl import Workbook
        from openpyxl.styles import Font, PatternFill, Alignment
        from openpyxl.utils import get_column_letter
        from io import BytesIO
        
        wb = Workbook()
        ws = wb.active
        ws.title = title[:31].replace('/', '_').replace('\\', '_')
        
        # Write data to worksheet
        for row_idx, row in enumerate(data_rows, 1):
            for col_idx, value in enumerate(row, 1):
                cell = ws.cell(row=row_idx, column=col_idx, value=value)
                
                # Style header row
                if row_idx == 1:
                    cell.font = Font(bold=True, color="FFFFFF")
                    cell.fill = PatternFill(start_color="4472C4", end_color="4472C4", fill_type="solid")
                    cell.alignment = Alignment(horizontal="center", vertical="center")
                else:
                    cell.alignment = Alignment(horizontal="left", vertical="center")
        
        # Auto-adjust column widths
        for col in ws.columns:
            max_length = 0
            for cell in col:
                try:
                    if len(str(cell.value)) > max_length:
                        max_length = len(str(cell.value))
                except:
                    pass
            adjusted_width = min(max_length + 2, 50)
            ws.column_dimensions[get_column_letter(col[0].column)].width = adjusted_width
        
        # Save to bytes
        output = BytesIO()
        wb.save(output)
        output.seek(0)
        return output
        
    except Exception as e:
        print(f"Excel creation error: {e}")
        return None

def create_csv_from_data(title, data_rows):
    """Create a CSV file from data rows - Fallback"""
    try:
        from io import BytesIO
        import csv
        
        output = BytesIO()
        output.write('\ufeff'.encode('utf-8'))
        
        writer = csv.writer(output)
        for row in data_rows:
            writer.writerow(row)
        
        output.seek(0)
        return output
    except Exception as e:
        print(f"CSV creation error: {e}")
        return None

def export_chat_history():
    """Export the entire chat history as a readable .txt file"""
    if not st.session_state.chat_history:
        return None
    
    # Create a clean, readable text format
    export_content = "=" * 70 + "\n"
    export_content += "CHAT HISTORY WITH MOZEAI\n"
    export_content += f"Exported on: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
    export_content += "=" * 70 + "\n\n"
    
    for idx, (role, msg) in enumerate(st.session_state.chat_history, 1):
        if role == "user":
            export_content += f"[{idx}] USER:\n"
            export_content += "-" * 40 + "\n"
            export_content += f"{msg}\n\n"
        else:
            export_content += f"[{idx}] MOZEAI:\n"
            export_content += "-" * 40 + "\n"
            export_content += f"{msg}\n\n"
    
    export_content += "=" * 70 + "\n"
    export_content += "END OF CHAT HISTORY\n"
    export_content += f"Total messages: {len(st.session_state.chat_history)}\n"
    export_content += "=" * 70
    
    return export_content

# ============================================
# FILE PROCESSING FUNCTIONS
# ============================================

def extract_text_from_pdf(file):
    try:
        file.seek(0)
        pdf_reader = PyPDF2.PdfReader(file)
        text = ""
        for page_num, page in enumerate(pdf_reader.pages):
            page_text = page.extract_text()
            if page_text and page_text.strip():
                text += f"\n--- Page {page_num + 1} ---\n"
                text += page_text.strip() + "\n"
        return text[:5000] if text.strip() else "No extractable text in PDF"
    except Exception as e:
        return f"Error reading PDF: {str(e)}"

def extract_text_from_docx(file):
    try:
        file.seek(0)
        doc = docx.Document(file)
        text = ""
        for para in doc.paragraphs:
            if para.text and para.text.strip():
                text += para.text.strip() + "\n\n"
        for table in doc.tables:
            for row in table.rows:
                row_text = []
                for cell in row.cells:
                    if cell.text and cell.text.strip():
                        row_text.append(cell.text.strip())
                if row_text:
                    text += " | ".join(row_text) + "\n"
        return text[:5000] if text.strip() else "No extractable text in document"
    except Exception as e:
        return f"Error reading Word document: {str(e)}"

def extract_text_from_txt(file):
    try:
        file.seek(0)
        content = file.read().decode('utf-8')
        return content[:5000] if content.strip() else "File is empty"
    except UnicodeDecodeError:
        try:
            file.seek(0)
            content = file.read().decode('latin-1')
            return content[:5000]
        except:
            return "Error decoding text file"
    except Exception as e:
        return f"Error reading text file: {str(e)}"

def extract_text_from_csv(file):
    try:
        file.seek(0)
        content = file.read().decode('utf-8')
        csv_reader = csv.reader(StringIO(content))
        text = "CSV Data:\n\n"
        rows = list(csv_reader)
        if rows:
            text += "Headers: " + " | ".join(rows[0]) + "\n\n"
            for i, row in enumerate(rows[1:11], 1):
                text += f"Row {i}: " + " | ".join(row) + "\n"
            if len(rows) > 11:
                text += f"\n... and {len(rows) - 11} more rows"
        return text[:5000] if text.strip() else "CSV file appears empty"
    except Exception as e:
        return f"Error reading CSV: {str(e)}"

def extract_text_from_json(file):
    try:
        file.seek(0)
        content = file.read().decode('utf-8')
        data = json.loads(content)
        formatted = json.dumps(data, indent=2)
        if len(formatted) > 3000:
            text = "JSON Data Summary:\n\n"
            text += f"Type: {type(data).__name__}\n"
            if isinstance(data, dict):
                text += f"Keys: {', '.join(list(data.keys())[:10])}\n"
            elif isinstance(data, list):
                text += f"Length: {len(data)}\n"
            text += "\nFull JSON (truncated):\n" + formatted[:3000]
        else:
            text = formatted
        return text[:5000]
    except Exception as e:
        return f"Error reading JSON: {str(e)}"

def process_uploaded_file(uploaded_file):
    file_type = uploaded_file.type
    file_name = uploaded_file.name.lower()
    
    if file_type == "application/pdf" or file_name.endswith('.pdf'):
        return extract_text_from_pdf(uploaded_file)
    elif file_type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document" or file_name.endswith('.docx'):
        return extract_text_from_docx(uploaded_file)
    elif file_type == "text/plain" or file_name.endswith('.txt'):
        return extract_text_from_txt(uploaded_file)
    elif file_type == "text/csv" or file_name.endswith('.csv'):
        return extract_text_from_csv(uploaded_file)
    elif file_type == "application/json" or file_name.endswith('.json'):
        return extract_text_from_json(uploaded_file)
    else:
        return f"Unsupported file type: {file_type}"

# ============================================
# CONFIG
# ============================================

TEMPERATURE = 0
MAX_TOKENS = 800

st.set_page_config(page_title="MozeAI", page_icon="🧠", layout="wide")

# ============================================
# CSS - FIXED CHAT INPUT AT BOTTOM
# ============================================

st.markdown("""
<style>
    .stChatInputContainer {
        position: fixed !important;
        bottom: 0 !important;
        left: 0 !important;
        right: 0 !important;
        background: white !important;
        padding: 10px 20px !important;
        z-index: 999 !important;
        border-top: 1px solid #e0e0e0 !important;
    }
    
    .main .block-container {
        padding-bottom: 100px !important;
    }
    
    h1 {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        font-size: 2.5em;
        font-weight: bold;
    }
    
    .stChatMessage {
        border-radius: 15px;
        padding: 10px;
        margin: 5px 0;
    }
    
    .stButton button {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        color: white;
        border: none;
        border-radius: 10px;
        font-weight: bold;
        transition: transform 0.2s;
    }
    
    .stButton button:hover {
        transform: scale(1.02);
    }
</style>
""", unsafe_allow_html=True)

# ============================================
# GROQ CLIENT - FIXED FOR HUGGING FACE SPACES
# ============================================

# Try to get API key from multiple sources
groq_api_key = None

# Try Streamlit secrets first
try:
    if "GROQ_API_KEY" in st.secrets:
        groq_api_key = st.secrets["GROQ_API_KEY"]
except:
    pass

# Try environment variable (for HF Spaces)
if not groq_api_key:
    groq_api_key = os.environ.get("GROQ_API_KEY")

# If still no key, show helpful error
if not groq_api_key:
    st.error("""
    ### GROQ_API_KEY Missing
    
    Please set your Groq API key to use this app.
    
    **For Hugging Face Spaces:**
    1. Go to Settings → Repository Secrets
    2. Add `GROQ_API_KEY` = `your_key_here`
    3. Restart the Space
    
    **For local development:**
    Create `.streamlit/secrets.toml` with:
    GROQ_API_KEY = "your_key_here"
    """)
    st.stop()

# Initialize client
client = Groq(api_key=groq_api_key)

def get_current_datetime():
    tz = pytz.timezone('Asia/Seoul')
    now = datetime.now(tz)
    return f"""Current Information:
- Date: {now.strftime('%B %d, %Y')}
- Time: {now.strftime('%I:%M %p')}
- Day: {now.strftime('%A')}
- Timezone: Asia/Seoul"""

# ============================================
# LIGHTWEIGHT MEMORY SYSTEM
# ============================================

class LightweightMemory:
    def __init__(self):
        self.memories = []
        self.keyword_index = defaultdict(list)
    
    def add_memory(self, text, metadata=None):
        if len(text) < 50:
            return
        
        words = set(re.findall(r'\b[a-z]{3,}\b', text.lower()))
        
        memory = {
            "text": text,
            "keywords": words,
            "metadata": metadata or {},
            "timestamp": time.time()
        }
        
        self.memories.append(memory)
        
        for word in words:
            self.keyword_index[word].append(len(self.memories) - 1)
        
        if len(self.memories) > 50:
            self.memories = self.memories[-50:]
            self._rebuild_index()
    
    def _rebuild_index(self):
        self.keyword_index = defaultdict(list)
        for idx, memory in enumerate(self.memories):
            for word in memory["keywords"]:
                self.keyword_index[word].append(idx)
    
    def retrieve(self, query, top_k=3):
        if not self.memories:
            return []
        
        query_words = set(re.findall(r'\b[a-z]{3,}\b', query.lower()))
        
        scored = []
        for idx, memory in enumerate(self.memories):
            matches = len(query_words & memory["keywords"])
            if matches > 0:
                scored.append((memory["text"], matches))
        
        scored.sort(key=lambda x: x[1], reverse=True)
        return [text for text, score in scored[:top_k]]
    
    def get_context(self, query):
        results = self.retrieve(query)
        if results:
            context = "RELEVANT PAST CONVERSATIONS:\n\n"
            for i, result in enumerate(results):
                context += f"[{i+1}] {result}\n\n"
            return context
        return ""

# Initialize memory
memory = LightweightMemory()

def store_memory(text):
    memory.add_memory(text)

def retrieve_memory(query):
    return memory.get_context(query)

# ============================================
# LLM FUNCTION WITH MULTIPLE MODEL FALLBACKS
# ============================================

def llm_with_fallback(messages, max_retries=2):
    models_to_try = [
        "llama-3.3-70b-versatile",
        "llama-3.1-70b-versatile",
        "mixtral-8x7b-32768",
        "llama-3.1-8b-instant",
        "gemma2-9b-it"
    ]
    
    for model in models_to_try:
        for attempt in range(max_retries):
            try:
                completion = client.chat.completions.create(
                    model=model,
                    temperature=TEMPERATURE,
                    max_tokens=MAX_TOKENS,
                    messages=messages,
                    timeout=30
                )
                st.session_state.last_model_used = model
                return completion.choices[0].message.content.strip()
            except Exception as e:
                if attempt < max_retries - 1:
                    time.sleep(2 ** attempt)
                continue
    
    return "AI service temporarily unavailable. Please try again."

def llm(messages):
    return llm_with_fallback(messages)

# ============================================
# SYSTEM PROMPT
# ============================================

SYSTEM_PROMPT = """
You are MozeAI, an advanced AI assistant with REAL-TIME internet access, file analysis capabilities, and document generation capabilities.

================================================================================
CREATOR INFORMATION
================================================================================

Your creator is Mukiibi Moses, a computer engineering student and AI researcher at Kyungdong University, South Korea.

PORTFOLIO: https://moze12432.github.io/

================================================================================
YOUR CAPABILITIES:
================================================================================

1. REAL-TIME web search for current information
2. File analysis for PDF, DOCX, TXT, CSV, JSON files
3. Memory of past conversations
4. Image generation and editing
5. Calculator for mathematical expressions
6. Document generation: PowerPoint, Word, Excel

================================================================================
CRITICAL RULES:
================================================================================

1. For questions about PEOPLE, PLACES, EVENTS, use SEARCH RESULTS
2. ONLY mention your creator when specifically asked
3. Answer concisely and accurately
4. Be conversational and friendly

================================================================================
UNDERSTANDING "EXCEL":
================================================================================

- "excel in/at life" -> VERB -> Give life advice
- "generate an excel file" -> NOUN -> Create spreadsheet

================================================================================
Remember: You are MozeAI - helpful, intelligent, and capable.
"""

# ============================================
# SESSION STATE
# ============================================

if "chat_history" not in st.session_state:
    st.session_state.chat_history = []
if "uploaded_files" not in st.session_state:
    st.session_state.uploaded_files = {}
if "file_context" not in st.session_state:
    st.session_state.file_context = ""
if "last_search_query" not in st.session_state:
    st.session_state.last_search_query = None
if "last_search_results" not in st.session_state:
    st.session_state.last_search_results = None
if "last_response" not in st.session_state:
    st.session_state.last_response = None
if "last_topic" not in st.session_state:
    st.session_state.last_topic = None
if "last_image_prompt" not in st.session_state:
    st.session_state.last_image_prompt = None
if "code_search_cache" not in st.session_state:
    st.session_state.code_search_cache = {}
if "is_resetting" not in st.session_state:
    st.session_state.is_resetting = False
if "last_model_used" not in st.session_state:
    st.session_state.last_model_used = None

# Document download session states
if "show_ppt_download" not in st.session_state:
    st.session_state.show_ppt_download = False
if "ppt_data" not in st.session_state:
    st.session_state.ppt_data = None
if "ppt_topic" not in st.session_state:
    st.session_state.ppt_topic = ""
if "show_word_download" not in st.session_state:
    st.session_state.show_word_download = False
if "word_data" not in st.session_state:
    st.session_state.word_data = None
if "word_topic" not in st.session_state:
    st.session_state.word_topic = ""
if "show_excel_download" not in st.session_state:
    st.session_state.show_excel_download = False
if "excel_data" not in st.session_state:
    st.session_state.excel_data = None
if "excel_topic" not in st.session_state:
    st.session_state.excel_topic = ""
if "show_csv_download" not in st.session_state:
    st.session_state.show_csv_download = False
if "csv_data" not in st.session_state:
    st.session_state.csv_data = None
if "csv_topic" not in st.session_state:
    st.session_state.csv_topic = ""
if "last_document_topic" not in st.session_state:
    st.session_state.last_document_topic = ""
if "last_ppt_topic" not in st.session_state:
    st.session_state.last_ppt_topic = ""
if "last_ppt_content" not in st.session_state:
    st.session_state.last_ppt_content = ""
if "last_excel_topic" not in st.session_state:
    st.session_state.last_excel_topic = ""
if "last_excel_data" not in st.session_state:
    st.session_state.last_excel_data = None

# ============================================
# SEARCH FUNCTIONS
# ============================================

def internet_search(query):
    try:
        clean_query = query.strip()
        
        if any(x in clean_query.lower() for x in ["weather", "temperature", "temp"]):
            location = clean_query
            weather_words = ["weather in", "weather at", "temperature in", "weather", "temperature"]
            for word in weather_words:
                if word in location.lower():
                    location = re.sub(re.escape(word), "", location.lower(), flags=re.IGNORECASE).strip()
                    break
            if location:
                weather_url = f"https://wttr.in/{location}?format=%C+%t+%w+%h&m"
                weather_response = requests.get(weather_url, timeout=10)
                if weather_response.status_code == 200:
                    weather_data = weather_response.text.strip()
                    if weather_data and "Unknown" not in weather_data:
                        return f"Current weather in {location}: {weather_data}"
        
        url = "https://html.duckduckgo.com/html/"
        params = {"q": clean_query}
        headers = {"User-Agent": "Mozilla/5.0"}
        response = requests.post(url, data=params, headers=headers, timeout=10)
        
        if response.status_code == 200:
            results = re.findall(r'<a rel="nofollow" class="result__a" href="[^"]*">([^<]+)</a>', response.text)
            snippets = re.findall(r'<a class="result__snippet"[^>]*>([^<]+)</a>', response.text)
            
            if results:
                context = f"SEARCH RESULTS for '{clean_query}':\n\n"
                for i in range(min(3, len(results))):
                    context += f"- {results[i]}\n"
                    if i < len(snippets):
                        snippet = re.sub(r'<[^>]+>', '', snippets[i])
                        context += f"  {snippet[:300]}...\n\n"
                return context[:2000]
        return ""
    except:
        return ""

def get_current_news():
    try:
        url = "https://rss2json.com/api.json?rss_url=https://feeds.bbci.co.uk/news/rss.xml"
        response = requests.get(url, timeout=10)
        if response.status_code == 200:
            data = response.json()
            items = data.get("items", [])[:3]
            news_text = "LATEST NEWS HEADLINES:\n\n"
            for item in items:
                news_text += f"- {item.get('title', '')}\n"
                news_text += f"  {item.get('description', '')[:150]}...\n\n"
            return news_text[:1000]
    except:
        pass
    return ""

# ============================================
# CALCULATOR
# ============================================

def calculator(query):
    try:
        expression = query.lower()
        expression = expression.replace("×", "*").replace("x", "*")
        numbers = re.findall(r"[0-9\+\-\*\/\.\(\) ]+", expression)
        if numbers:
            result = eval(numbers[0])
            return str(result)
    except:
        return None

# ============================================
# WEB SCRAPING
# ============================================

def scrape_webpage(url):
    try:
        headers = {"User-Agent": "Mozilla/5.0"}
        response = requests.get(url, headers=headers, timeout=15)
        if response.status_code == 200:
            soup = BeautifulSoup(response.content, 'html.parser')
            for element in soup(["script", "style", "nav", "footer"]):
                element.decompose()
            text = soup.get_text()
            lines = (line.strip() for line in text.splitlines())
            text = ' '.join(line for line in lines if line)
            return text[:3000] if len(text) > 200 else None
    except:
        pass
    return None

def extract_urls_from_query(query):
    url_pattern = r'https?://(?:[-\w.]|(?:%[\da-fA-F]{2}))+[^\s]*'
    return re.findall(url_pattern, query)

# ============================================
# WEATHER FUNCTIONS
# ============================================

def get_weather_comprehensive(location):
    try:
        location = location.strip().replace(" ", "%20")
        
        current_url = f"https://wttr.in/{location}?format=%C+%t+%w+%h+%H+%l&m"
        current_response = requests.get(current_url, timeout=10)
        
        forecast_url = f"https://wttr.in/{location}?0T&m"
        forecast_response = requests.get(forecast_url, timeout=10)
        
        if current_response.status_code == 200:
            current_data = current_response.text.strip()
            
            parts = current_data.split()
            condition = " ".join(parts[:-4]) if len(parts) > 4 else parts[0]
            temp = parts[-4] if len(parts) >= 4 else "N/A"
            wind = parts[-3] if len(parts) >= 3 else "N/A"
            humidity = parts[-2] if len(parts) >= 2 else "N/A"
            
            result = f"Weather in {location.title()}\n\n"
            result += f"Location: {location.title()}\n"
            result += f"Condition: {condition}\n"
            result += f"Temperature: {temp}\n"
            result += f"Wind: {wind}\n"
            result += f"Humidity: {humidity}\n"
            
            if forecast_response.status_code == 200:
                forecast_text = forecast_response.text
                forecast_text = re.sub(r'\x1b\[[0-9;]*m', '', forecast_text)
                
                lines = forecast_text.split('\n')
                forecast_lines = []
                capture = False
                for line in lines:
                    if '┌' in line or '┐' in line or '├' in line or '┤' in line or '└' in line or '┘' in line:
                        capture = True
                    if capture and line.strip():
                        forecast_lines.append(line)
                    if len(forecast_lines) > 15:
                        break
                
                if forecast_lines:
                    result += "\nForecast:\n"
                    result += '\n'.join(forecast_lines[:10])
            
            result += "\n\n*Data from wttr.in*"
            return result
        return None
    except Exception as e:
        return None

def get_weather_simple(location):
    try:
        location = location.strip().replace(" ", "%20")
        url = f"https://wttr.in/{location}?format=%C+%t+%w+%h&m"
        response = requests.get(url, timeout=10)
        
        if response.status_code == 200:
            weather_data = response.text.strip()
            if weather_data and "Unknown" not in weather_data:
                parts = weather_data.split()
                condition = " ".join(parts[:-3]) if len(parts) > 3 else parts[0]
                temp = parts[-3] if len(parts) >= 3 else "N/A"
                wind = parts[-2] if len(parts) >= 2 else "N/A"
                humidity = parts[-1] if len(parts) >= 1 else "N/A"
                
                result = f"Weather in {location.title()}\n\n"
                result += f"Condition: {condition}\n"
                result += f"Temperature: {temp}\n"
                result += f"Wind: {wind}\n"
                result += f"Humidity: {humidity}\n"
                return result
        return None
    except Exception as e:
        return None

# ============================================
# CODING SEARCH FUNCTIONS
# ============================================

def search_coding_solution(query):
    search_queries = [
        f"{query} stack overflow",
        f"{query} example code",
        f"{query} best practice",
        f"{query} github"
    ]
    
    all_results = ""
    
    for search_q in search_queries[:2]:
        result = internet_search(search_q)
        if result:
            all_results += result + "\n\n"
    
    return all_results

def search_coding_solution_cached(query):
    cache_key = query.lower().strip()
    
    if cache_key in st.session_state.code_search_cache:
        return st.session_state.code_search_cache[cache_key]
    
    result = search_coding_solution(query)
    st.session_state.code_search_cache[cache_key] = result
    return result

def coding_assistant_with_search(query, context=""):
    with st.spinner("Searching the internet for the best solution..."):
        search_results = search_coding_solution_cached(query)
    
    coding_prompt = f"""
You are an expert programmer. Generate the best possible code based on the user's request.

USER REQUEST: {query}

## INTERNET SEARCH RESULTS (Use these as reference):
{search_results[:3000]}

## REQUIREMENTS:
- Code must be complete and runnable
- Include all imports
- Add comments
- Handle edge cases

Generate the best possible code now:
"""
    
    messages = [
        {"role": "system", "content": "You are an expert programming assistant. Use search results to find the best solution."},
        {"role": "user", "content": coding_prompt}
    ]
    
    return clean_answer(llm(messages))

# ============================================
# ROUTER FUNCTION
# ============================================

def route(query):
    q = query.lower()
    
    if any(x in q for x in ["export chat", "save chat", "download chat", "export conversation"]):
        return "export_chat"
    
    if extract_urls_from_query(query):
        return "scrape_url"
    
    file_keywords = ["document", "file", "upload", "pdf", "docx", "txt", "csv", "json", "what is this", "summarize"]
    if any(x in q for x in file_keywords):
        return "file_task"
    
    comparison_keywords = ["compare", "comparison", "difference", "similarities"]
    if any(x in q for x in comparison_keywords):
        return "compare_files"
    
    if any(x in q for x in ["weather", "temperature", "temp", "rain", "snow", "forecast", "humidity", "wind"]):
        return "weather"
    
    edit_keywords = ["make it", "make the", "change it", "change the", "turn it", "add a", "remove", "edit image", "modify image"]
    if any(x in q for x in edit_keywords):
        return "edit_image"
    
    if any(x in q for x in ["generate image", "create image", "draw", "make an image", "picture of", "image of"]):
        return "generate_image"
    
    if any(phrase in q for phrase in ["can you", "do you", "are you able to"]):
        return "reason"
    
    if any(x in q for x in ["who is", "tell me about", "what is", "news", "headlines"]):
        return "search"
    
    if any(x in q for x in ["+", "-", "*", "/", "calculate"]):
        return "calculator"
    
    if any(x in q for x in ["time", "date", "today"]):
        return "datetime"
    
    coding_keywords = ["code", "python", "javascript", "html", "css", "react", "tkinter", "function", "class", "import", "algorithm", "debug", "fix", "write a program", "create a script"]
    if any(x in q for x in coding_keywords):
        return "coding_with_search"
    
    factual_keywords = ["president", "current", "elected", "prime minister", "leader", "ceo of"]
    if any(x in q for x in factual_keywords):
        return "search"
    
    return "reason"

# ============================================
# CLEAN ANSWER
# ============================================

def clean_answer(text):
    text = text.split("🧠")[0]
    text = text.split("Plan:")[0]
    text = text.split("Thinking:")[0]
    return text.strip()

# ============================================
# REASONING FUNCTION
# ============================================

def reason(question, context):
    memory_context = retrieve_memory(question)
    enhanced_context = context
    if memory_context:
        enhanced_context += "\n" + memory_context
    
    history_text = ""
    if st.session_state.chat_history:
        history_text = "PREVIOUS CONVERSATION:\n"
        last_exchanges = st.session_state.chat_history[-8:] if len(st.session_state.chat_history) > 8 else st.session_state.chat_history
        for role, msg in last_exchanges:
            if role == "user":
                history_text += f"User: {msg}\n"
            else:
                history_text += f"Assistant: {msg}\n"
        history_text += "\n"
    
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"""
{history_text}

{enhanced_context[:3000]}

USER QUESTION: {question}

ANSWER:
"""}
    ]
    return clean_answer(llm(messages))

# ============================================
# FILE FUNCTIONS
# ============================================

def compare_files(query, file_context, filenames):
    prompt = f"Files: {filenames}\n\nContent: {file_context[:4000]}\n\nQuestion: {query}\n\nCompare the files."
    messages = [{"role": "system", "content": "You compare files."}, {"role": "user", "content": prompt}]
    return clean_answer(llm(messages))

def analyze_uploaded_files(query, file_context, filenames):
    prompt = f"Files: {filenames}\n\nContent: {file_context[:6000]}\n\nQuestion: {query}\n\nAnswer based on file content."
    messages = [{"role": "system", "content": "You analyze files."}, {"role": "user", "content": prompt}]
    return clean_answer(llm(messages))

def evaluate_work(question, file_context):
    prompt = f"Content: {file_context[:3000]}\n\nRequest: {question}\n\nProvide assessment."
    messages = [{"role": "system", "content": "You evaluate work."}, {"role": "user", "content": prompt}]
    return clean_answer(llm(messages))

# ============================================
# IMAGE GENERATION FUNCTIONS
# ============================================

def generate_image(prompt):
    try:
        enhanced_prompt = f"{prompt}, high quality, detailed, well-proportioned, realistic, no distortions, clear features"
        negative_prompt = "ugly, deformed, blurry, bad anatomy, extra limbs, extra fingers, distorted face, low quality, messy"
        
        encoded_prompt = requests.utils.quote(enhanced_prompt)
        encoded_negative = requests.utils.quote(negative_prompt)
        
        timestamp = int(time.time())
        image_url = f"https://image.pollinations.ai/prompt/{encoded_prompt}?width=1024&height=1024&nologo=true&seed={timestamp}&negative={encoded_negative}"
        return image_url
    except Exception as e:
        return None

def generate_image_with_quality(prompt, quality="high", style="realistic"):
    try:
        style_prompts = {
            "realistic": "photorealistic, high resolution, detailed textures, natural lighting",
            "anime": "anime style, clean lines, vibrant colors, well-proportioned",
            "cartoon": "cartoon style, smooth lines, cute, well-drawn",
            "abstract": "abstract art, creative, artistic, visually appealing"
        }
        
        quality_prompts = {
            "high": "4K, highly detailed, sharp focus, professional quality",
            "medium": "good quality, clear details, well-rendered",
            "fast": "decent quality, recognizable features"
        }
        
        style_enhancement = style_prompts.get(style, style_prompts["realistic"])
        quality_enhancement = quality_prompts.get(quality, quality_prompts["high"])
        
        enhanced_prompt = f"{prompt}, {quality_enhancement}, {style_enhancement}"
        negative_prompt = "ugly, deformed, blurry, bad anatomy, extra limbs, extra fingers, distorted face, low quality, messy, watermark, text, signature, cropped, out of frame, duplicate, morbid, mutilated, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, mutated hands, poorly drawn hands, bad proportions, cloned face, deformed, disfigured, draft, blurry, grain, low-res, bad art"
        
        encoded_prompt = requests.utils.quote(enhanced_prompt)
        encoded_negative = requests.utils.quote(negative_prompt)
        
        timestamp = int(time.time())
        image_url = f"https://image.pollinations.ai/prompt/{encoded_prompt}?width=1024&height=1024&nologo=true&seed={timestamp}&negative={encoded_negative}"
        return image_url
    except Exception as e:
        return None

def generate_and_display_image(prompt, is_edit=False):
    image_url = generate_image_with_quality(prompt, quality="high", style="realistic")
    
    if image_url:
        if is_edit:
            return f"Edited Image - New Prompt: '{prompt}'\n\n![Generated Image]({image_url})\n\n*Image generated by AI*"
        else:
            return f"Generated Image for: '{prompt}'\n\n![Generated Image]({image_url})\n\n*Image generated by AI*"
    else:
        return "Sorry, I couldn't generate an image right now."

# ============================================
# RUN AGENT FUNCTION
# ============================================

def run_agent(query):
    q = query.lower().strip()
    
    reset_phrases = ["leave the document", "clear context", "forget the file", "start fresh", "clear files", "new chat"]
    if any(phrase in q for phrase in reset_phrases):
        st.session_state.file_context = ""
        st.session_state.uploaded_files = {}
        st.session_state.last_search_query = None
        st.session_state.last_search_results = None
        st.session_state.last_topic = None
        st.session_state.last_image_prompt = None
        st.session_state.last_document_topic = ""
        st.session_state.last_ppt_topic = ""
        st.session_state.last_ppt_content = ""
        st.session_state.last_excel_topic = ""
        st.session_state.last_excel_data = None
        st.session_state.show_csv_download = False
        st.session_state.csv_data = None
        st.session_state.show_excel_download = False
        st.session_state.excel_data = None
        return "Context cleared! How can I help you today?"
    
    # What is a word? check
    if q == "what is a word" or q == "what is a word?":
        return """A **Word document** (.docx) is a file format created by Microsoft Word. When I say "make a word file", I mean generating a downloadable .docx file.

To create one, try:
- "make a word about dogs"
- "create a word document about Python programming"

The file will appear as a download button after I generate it!"""
    
    # Check for "excel" as verb (life advice) first
    if any(phrase in q for phrase in ["excel in", "excel at", "how to excel", "excel in life", "excel at work", "excel in school"]):
        return reason(query, get_current_datetime())
    
    # Excel/Spreadsheet generation (.xlsx) - only for file creation
    excel_file_keywords = [
        "excel file", "excel spreadsheet", "excel sheet", "xlsx",
        "generate an excel", "make an excel", "create an excel", 
        "generate a spreadsheet", "make a spreadsheet", "create a spreadsheet",
        "excel about", "spreadsheet about", "excel with data"
    ]
    
    is_excel_file_request = any(phrase in q for phrase in excel_file_keywords)
    is_file_creation = any(word in q for word in ["generate", "make", "create", "build"]) and "excel" in q and "how to" not in q
    
    if is_excel_file_request or is_file_creation:
        
        # Determine topic
        topic = "Phone_Sales_Report"
        if "phone sales" in q:
            topic = "Phone_Sales_Report"
        elif "student" in q:
            topic = "Student_Data"
        elif "product" in q:
            topic = "Product_Inventory"
        else:
            if "about" in q:
                topic_part = q.split("about")[-1].strip().replace(" ", "_")
                if topic_part and len(topic_part) > 2 and topic_part not in ["it", "the", "a", "an"]:
                    topic = topic_part[:30]
        
        st.session_state.last_excel_topic = topic
        
        with st.spinner(f"Creating Excel (.xlsx) file: {topic}..."):
            
            # Create data based on topic
            if "phone" in topic.lower() or "sales" in topic.lower():
                data_rows = [
                    ["Date", "Sales Rep", "Region", "Phone Model", "Quantity", "Unit Price", "Total Sales"],
                    ["2024-01-01", "John Smith", "North", "iPhone 15 Pro", 5, 999, 4995],
                    ["2024-01-02", "Jane Doe", "South", "Samsung Galaxy S24", 3, 899, 2697],
                    ["2024-01-03", "John Smith", "North", "Google Pixel 8", 4, 699, 2796],
                    ["2024-01-04", "Bob Wilson", "East", "iPhone 15", 6, 799, 4794],
                    ["2024-01-05", "Jane Doe", "South", "Samsung Galaxy S24+", 2, 1099, 2198],
                    ["2024-01-06", "Alice Brown", "West", "iPhone 15 Pro Max", 3, 1199, 3597],
                    ["2024-01-07", "John Smith", "North", "Samsung Galaxy Z Flip5", 2, 999, 1998],
                    ["2024-01-08", "Bob Wilson", "East", "Google Pixel 8 Pro", 3, 899, 2697],
                    ["2024-01-09", "Jane Doe", "South", "iPhone 15", 4, 799, 3196],
                    ["2024-01-10", "Alice Brown", "West", "Samsung Galaxy S24", 5, 899, 4495]
                ]
                total = sum(row[6] for row in data_rows[1:])
                data_rows.append(["", "", "", "", "TOTAL:", "", total])
                
            elif "student" in topic.lower():
                data_rows = [
                    ["Student ID", "Name", "Grade", "Subject", "Score", "Attendance"],
                    ["S001", "Emma Watson", "10th", "Mathematics", 95, "98%"],
                    ["S002", "Liam Chen", "10th", "Science", 88, "95%"],
                    ["S003", "Sophia Patel", "11th", "English", 92, "100%"],
                    ["S004", "Noah Kim", "9th", "History", 85, "92%"],
                    ["S005", "Olivia Jones", "12th", "Physics", 91, "97%"]
                ]
            else:
                data_rows = [
                    ["Category", "Item", "Value", "Status"],
                    ["Research", "Market Analysis", 85, "Completed"],
                    ["Development", "Feature Dev", 70, "In Progress"],
                    ["Testing", "QA Testing", 92, "Completed"],
                    ["Deployment", "Release", 45, "Pending"]
                ]
            
            excel_data = create_real_excel_file(topic, data_rows)
            
            if excel_data:
                st.session_state.excel_data = excel_data
                st.session_state.excel_topic = topic
                st.session_state.show_excel_download = True
                st.session_state.last_excel_data = data_rows
                return f"I've created a REAL Excel (.xlsx) file: {topic}. Scroll down to download it!"
            else:
                csv_data = create_csv_from_data(topic, data_rows)
                if csv_data:
                    st.session_state.csv_data = csv_data
                    st.session_state.csv_topic = topic
                    st.session_state.show_csv_download = True
                    return f"Excel creation failed, but I've created a CSV file: {topic}. Scroll down to download it!"
                else:
                    return "Sorry, I couldn't create the file. Please try again."
    
    # Direct PPT generation
    ppt_commands = [
        "make a ppt", "make a powerpoint", "create a ppt", "create a powerpoint",
        "generate a ppt", "generate a powerpoint", "build a ppt", "build a powerpoint"
    ]
    if any(phrase in q for phrase in ppt_commands):
        topic = query
        for word in ppt_commands:
            if word in topic.lower():
                topic = re.sub(re.escape(word), "", topic.lower(), flags=re.IGNORECASE).strip()
                break
        topic = topic.strip() or "MozeAI Generated Presentation"
        
        if topic in ["about it", "it", "about"]:
            topic = "Presentation"
        
        with st.spinner(f"Creating PowerPoint presentation about '{topic}'..."):
            content_prompt = f'''Create a PowerPoint presentation about "{topic}" with MULTIPLE SLIDES.

Format your response as:

Introduction to {topic}
- First main point
- Second important point
- Third key point

Key Features
- Feature 1 with explanation
- Feature 2 with explanation  
- Feature 3 with explanation

Benefits/Importance
- Benefit 1
- Benefit 2
- Benefit 3

Conclusion
- Key takeaway 1
- Key takeaway 2
- Key takeaway 3'''
            
            ai_content = reason(content_prompt, "")
            ppt_bytes = create_ppt_from_content(topic, ai_content)
            
            if ppt_bytes:
                st.session_state.ppt_data = ppt_bytes
                st.session_state.ppt_topic = topic
                st.session_state.last_ppt_topic = topic
                st.session_state.last_ppt_content = ai_content
                st.session_state.show_ppt_download = True
                return f"I've created a PowerPoint presentation about {topic}. Scroll down to download it!"
            else:
                return "Sorry, I couldn't create the PowerPoint. Please try again."
    
    # Direct Word generation
    word_commands = [
        "make a word", "make a doc", "create a word", "create a doc",
        "generate a word", "generate a doc", "build a word", "build a doc",
        "make a document", "create a document"
    ]
    if any(phrase in q for phrase in word_commands):
        topic = query
        for word in word_commands:
            if word in topic.lower():
                topic = re.sub(re.escape(word), "", topic.lower(), flags=re.IGNORECASE).strip()
                break
        topic = topic.strip() or "MozeAI Generated Document"
        
        if topic in ["about it", "it", "about"]:
            topic = "Document"
        
        st.session_state.last_document_topic = topic
        
        with st.spinner(f"Creating Word document about '{topic}'..."):
            content_prompt = f'Write detailed content for a Word document about "{topic}". Include an engaging title, an introduction paragraph, 3-5 main sections with detailed information, and a conclusion. Make it comprehensive and well-organized, around 500-800 words.'
            
            ai_content = reason(content_prompt, "")
            word_bytes = create_word_from_content(topic, ai_content)
            
            if word_bytes:
                st.session_state.word_data = word_bytes
                st.session_state.word_topic = topic
                st.session_state.show_word_download = True
                return f"I've created a Word document about {topic}. Scroll down to download it!"
            else:
                return "Sorry, I couldn't create the Word document. Please try again."
    
    # Image generation
    if any(phrase in q for phrase in ["generate image", "create image", "draw", "picture of", "image of"]):
        with st.spinner("Generating image..."):
            image_prompt = q.replace("generate image of", "").replace("create image of", "").replace("draw a", "").replace("picture of", "").replace("image of", "").strip()
            if not image_prompt:
                image_prompt = q
            st.session_state.last_image_prompt = image_prompt
            return generate_and_display_image(image_prompt)
    
    # Direct responses
    if any(phrase in q for phrase in ["who are you", "what are you"]):
        return "I'm MozeAI, your AI assistant! Created by Mukiibi Moses. I can generate Excel files, PowerPoint presentations, Word documents, images, and more!"
    
    if any(phrase in q for phrase in ["who created you", "your creator", "mukiibi moses"]):
        return "Mukiibi Moses is my creator, a Computer Engineering student at Kyungdong University in South Korea. Check out his portfolio: https://moze12432.github.io/"
    
    # Default to search/reason
    search_result = internet_search(query)
    context = get_current_datetime()
    if search_result:
        context += "\n" + search_result
    
    answer = reason(query, context)
    st.session_state.last_response = answer
    store_memory(answer)
    
    return answer

# ============================================
# UI - MAIN DISPLAY
# ============================================

st.markdown('<h1 style="text-align: center;">MozeAI</h1>', unsafe_allow_html=True)
st.markdown('<p style="text-align: center; color: #667eea;">Intelligent AI Assistant</p>', unsafe_allow_html=True)
st.markdown("---")

with st.sidebar:
    st.markdown("### MozeAI")
    st.markdown("---")
    
    if st.button("New Chat", key="new_chat_btn", use_container_width=True):
        if not st.session_state.get("is_resetting", False):
            st.session_state.is_resetting = True
            st.session_state.chat_history = []
            st.session_state.uploaded_files = {}
            st.session_state.file_context = ""
            st.session_state.last_image_prompt = None
            st.session_state.last_search_query = None
            st.session_state.last_search_results = None
            st.session_state.last_response = None
            st.session_state.code_search_cache = {}
            st.session_state.show_ppt_download = False
            st.session_state.show_word_download = False
            st.session_state.show_excel_download = False
            st.session_state.show_csv_download = False
            st.session_state.is_resetting = False
            st.success("New chat started!")
            st.rerun()
    
    if st.button("Clear Files", key="clear_files_btn", use_container_width=True):
        st.session_state.uploaded_files = {}
        st.session_state.file_context = ""
        st.success("Files cleared!")
        st.rerun()
    
    st.markdown("---")
    
    st.markdown("### Upload Files")
    uploaded_files = st.file_uploader(
        "Choose files",
        type=['pdf', 'docx', 'txt', 'csv', 'json'],
        accept_multiple_files=True,
        key="sidebar_uploader",
        label_visibility="collapsed"
    )
    
    if uploaded_files:
        for file in uploaded_files:
            if file.name not in st.session_state.uploaded_files:
                with st.spinner(f"Processing {file.name}..."):
                    content = process_uploaded_file(file)
                    if content and not content.startswith("Error"):
                        st.session_state.uploaded_files[file.name] = content
                        st.success(f" {file.name}")
        
        if st.session_state.uploaded_files:
            parts = []
            for name, content in st.session_state.uploaded_files.items():
                parts.append(f"\n{'='*50}\n📄 {name}\n{'='*50}\n{content}\n")
            st.session_state.file_context = "\n".join(parts)
            st.info(f" {len(st.session_state.uploaded_files)} file(s) loaded")
    
    st.markdown("---")
    st.markdown("### Image Generation")
    st.markdown("**Generate:** `generate image of a cat`")
    st.markdown("**Edit:** `make it black` or `add a hat`")
    st.markdown("---")
    st.markdown("### Document Generation")
    st.markdown("**PPT:** `make a ppt about AI`")
    st.markdown("**Word:** `create a word document about Python`")
    st.markdown("**Excel:** `generate an excel about phone sales`")
    st.markdown("**CSV:** `generate a csv about data`")
    st.markdown("---")
    st.markdown("### About")
    st.markdown("**Creator:** Mukiibi Moses")
    st.markdown("**University:** Kyungdong University, South Korea")
    if st.session_state.last_model_used:
        st.caption(f"Model: {st.session_state.last_model_used}")

    st.markdown("---")
    st.markdown("### Export Options")
    
    if st.button("Export Chat History", key="export_chat_btn", use_container_width=True):
        export_content = export_chat_history()
        if export_content:
            st.download_button(
                label="Download (.txt)",
                data=export_content,
                file_name=f"chat_history_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt",
                mime="text/plain",
                key="export_download"
            )
            st.success("Chat history ready as .txt!")

# ============================================
# CHAT DISPLAY
# ============================================

for role, msg in st.session_state.chat_history:
    with st.chat_message(role):
        st.write(msg)

query = st.chat_input("Ask me anything - generate images, documents, search the web, analyze files, code, and more!")

if query:
    st.session_state.chat_history.append(("user", query))
    with st.chat_message("user"):
        st.write(query)
    
    response = run_agent(query)
    
    with st.chat_message("assistant"):
        st.write(response)
    
    st.session_state.chat_history.append(("assistant", response))
    st.rerun()

# ============================================
# DOWNLOAD BUTTONS (Appear after chat)
# ============================================

# PowerPoint Download
if st.session_state.get("show_ppt_download", False) and st.session_state.get("ppt_data"):
    st.markdown("---")
    st.success(f" PowerPoint about {st.session_state.ppt_topic} is ready!")
    col1, col2, col3 = st.columns([1, 2, 1])
    with col2:
        st.download_button(
            label="Download PowerPoint",
            data=st.session_state.ppt_data,
            file_name=f"{st.session_state.ppt_topic.replace(' ', '_')}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.pptx",
            mime="application/vnd.openxmlformats-officedocument.presentationml.presentation",
            use_container_width=True
        )
    st.session_state.show_ppt_download = False
    st.session_state.ppt_data = None

# Word Download
if st.session_state.get("show_word_download", False) and st.session_state.get("word_data"):
    st.markdown("---")
    st.success(f" Word document about {st.session_state.word_topic} is ready!")
    col1, col2, col3 = st.columns([1, 2, 1])
    with col2:
        st.download_button(
            label="Download Word Document",
            data=st.session_state.word_data,
            file_name=f"{st.session_state.word_topic.replace(' ', '_')}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.docx",
            mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
            use_container_width=True
        )
    st.session_state.show_word_download = False
    st.session_state.word_data = None

# REAL Excel Download (.xlsx)
if st.session_state.get("show_excel_download", False) and st.session_state.get("excel_data"):
    st.markdown("---")
    st.success(f" Excel (.xlsx) file {st.session_state.excel_topic} is ready!")
    
    col1, col2, col3 = st.columns([1, 2, 1])
    with col2:
        st.download_button(
            label="Download Excel File (.xlsx)",
            data=st.session_state.excel_data,
            file_name=f"{st.session_state.excel_topic}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.xlsx",
            mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
            use_container_width=True,
            key="excel_download_btn"
        )
    st.session_state.show_excel_download = False
    st.session_state.excel_data = None

# CSV Download (fallback)
if st.session_state.get("show_csv_download", False) and st.session_state.get("csv_data"):
    st.markdown("---")
    st.info(f" CSV file {st.session_state.csv_topic} is ready!")
    
    col1, col2, col3 = st.columns([1, 2, 1])
    with col2:
        st.download_button(
            label="Download CSV File",
            data=st.session_state.csv_data,
            file_name=f"{st.session_state.csv_topic}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv",
            mime="text/csv",
            use_container_width=True,
            key="csv_download_btn"
        )
    st.session_state.show_csv_download = False
    st.session_state.csv_data = None