import streamlit as st
import time
from openai import OpenAI
import os
from dotenv import load_dotenv
import pandas as pd
import io
import re

# Load environment variables from .env file (for local development)
load_dotenv()

# Page configuration
st.set_page_config(
    page_title="Creative Professionals Job Analysis",
    page_icon="🎨",
    layout="wide"
)

# Password protection
def check_password():
    correct_password = None
    try:
        if "app_password" in st.secrets:
            correct_password = st.secrets["app_password"]
    except:
        pass
    if not correct_password:
        correct_password = os.environ.get("APP_PASSWORD") or os.getenv("APP_PASSWORD")
    if not correct_password:
        correct_password = "cpro2024"

    def password_entered():
        entered_password = st.session_state.get("password", "")
        if entered_password == correct_password:
            st.session_state["password_correct"] = True
            if "password" in st.session_state:
                del st.session_state["password"]
        else:
            st.session_state["password_correct"] = False

    if st.session_state.get("password_correct", False):
        return True

    st.markdown("## 🔒 Authentication Required")
    st.text_input(
        "Please enter the password to access this app:",
        type="password",
        on_change=password_entered,
        key="password"
    )

    if "password_correct" in st.session_state and not st.session_state["password_correct"]:
        st.error("😕 Password incorrect. Please try again.")
    return False

if not check_password():
    st.stop()

# Custom CSS
st.markdown("""
<style>
    .stButton > button {
        background-color: #B19CD9;
        color: white;
        border-radius: 5px;
    }
    .stButton > button:hover {
        background-color: #9B84C9;
        color: white;
    }
</style>
""", unsafe_allow_html=True)

# Title
st.title("🎨 Creative Professionals Job Listings Analysis")

# Sidebar
with st.sidebar:

    # --- NEW: Clear Conversation at top, no “Configuration” header ---
    if st.button("🗑️ Clear Conversation", use_container_width=True):
        st.session_state.messages = []
        st.session_state.thread_id = None
        st.rerun()

    api_key = st.secrets.get("openai_api_key", os.getenv("OPENAI_API_KEY", ""))
    assistant_id = "asst_vaqQlpvneKehFNBtk8LOIj0A"

    if not api_key:
        st.error("⚠️ OpenAI API key not configured.")
        st.stop()

    st.divider()
    st.header("Select Analysis Method")

    mode = st.radio(
        "Choose method:",
        ["Code Interpreter (CSV)", "Semantic Search", "Hybrid"],
        index=0
    )

    mode_prefix = {
        "Code Interpreter (CSV)": "/csv",
        "Semantic Search": "/semantic",
        "Hybrid": "/hybrid"
    }[mode]

    # --- UPDATED: Remove emojis from section headers ---
    st.subheader(f"About the {mode} method:")

    if mode == "Code Interpreter (CSV)":
        st.markdown("""
        Use this for precise, structured questions. The assistant loads the CSV of job listings,
        runs Python, and filters or aggregates the data.

        **Examples:**
        - *“List all job titles of Business Professionals.”*
        - *“Count the number of Design Pro listings.”*
        - *“Show me all job titles that mention Figma and are classified as Design Pro.”*

        **Best for:**
        - Exact counts  
        - Filtering by category  
        - CSV-driven analysis  
        """)

    elif mode == "Semantic Search":
        st.markdown("""
        Use this for interpretive, conceptual, or text-heavy questions. The assistant retrieves
        text from job descriptions using vector search.

        **Examples:**
        - *“What are common responsibilities of Business Professionals?”*
        - *“Which job descriptions mention generative AI skills?”*
        - *“Summarize trends in photography-related jobs.”*

        **Best for:**
        - Thematic summaries  
        - Conceptual trends  
        - Free-text exploration  
        """)

    else:
        st.markdown("""
        Use this when you want both semantic retrieval and structured filtering.
        The assistant mixes judgement and rigorous filtering.

        **Examples:**
        - *“Find jobs mentioning TikTok and report their creative category.”*
        - *“Identify job descriptions referencing UX research and list their job titles.”*

        **Best for:**
        - Hybrid text + structured queries  
        - Context-driven filtering  
        """)

    st.divider()

# Session init
if "messages" not in st.session_state:
    st.session_state.messages = []
if "thread_id" not in st.session_state:
    st.session_state.thread_id = None

client = OpenAI(api_key=api_key)

def is_mode_detection_failure(response_text):
    """
    Check if the response indicates the assistant failed to detect the mode prefix.
    This happens when the assistant asks for mode specification despite already having one.
    """
    if not response_text:
        return False

    failure_phrases = [
        "please specify /csv",
        "please specify a mode",
        "specify /csv, /semantic, or /hybrid",
        "/csv, /semantic, or /hybrid mode",
        "which mode would you like"
    ]
    response_lower = response_text.lower()
    return any(phrase in response_lower for phrase in failure_phrases)

def clean_response_text(text):
    """
    Remove source citations and sandbox links from response text.
    These citations reference internal files that cannot be downloaded.
    """
    if not text:
        return ""

    # Remove source citations like 【10:2†source】 or 【14:0†job_listings_part_27.txt】
    text = re.sub(r'【[^】]*】', '', text)

    # Remove sandbox:// links
    text = re.sub(r'\[.*?\]\(sandbox:/[^\)]*\)', '', text)
    text = re.sub(r'\[Download.*?\]\(/mnt/data/.*?\)', '', text)
    text = re.sub(r'sandbox:/[^\s\)]+', '', text)

    # Remove "Source: filename" lines
    text = re.sub(r'Source:\s*[^\n]*?【[^】]*】\s*\.', '', text)
    text = re.sub(r'^\s*[○◦•]\s*Source:.*$', '', text, flags=re.MULTILINE)

    # Clean up extra whitespace
    text = re.sub(r'\n\s*\n\s*\n', '\n\n', text)

    return text.strip()

def extract_table_from_response(response_text):
    """
    Extract tabular data from assistant's response.
    Looks for markdown tables or numbered lists that can be converted to CSV.
    """
    lines = response_text.split('\n')
    data_rows = []

    # Try to find a markdown table first
    in_table = False
    headers = None

    for line in lines:
        line = line.strip()
        if '|' in line and not line.startswith('---'):
            if headers is None:
                # First row with pipes is likely headers
                headers = [h.strip() for h in line.split('|') if h.strip()]
                in_table = True
            elif in_table:
                # Subsequent rows are data
                row = [cell.strip() for cell in line.split('|') if cell.strip()]
                if row and len(row) == len(headers):
                    data_rows.append(row)
        elif in_table and not line:
            # Empty line might end the table
            break

    # If we found a table, return it as DataFrame
    if headers and data_rows:
        return pd.DataFrame(data_rows, columns=headers)

    # Try to find numbered lists (e.g., "1. Title - 123 jobs")
    pattern = r'^\d+\.\s+(.+?)\s+[-–]\s+(\d+)'
    for line in lines:
        match = re.match(pattern, line.strip())
        if match:
            data_rows.append([match.group(1).strip(), match.group(2).strip()])

    if data_rows:
        return pd.DataFrame(data_rows, columns=['Item', 'Count'])

    return None

def query_assistant(question, api_key, assistant_id, thread_id=None, timeout=600):
    client = OpenAI(api_key=api_key)
    try:
        if thread_id:
            thread = client.beta.threads.retrieve(thread_id)
        else:
            thread = client.beta.threads.create()
            thread_id = thread.id

        client.beta.threads.messages.create(
            thread_id=thread_id,
            role="user",
            content=question
        )

        run = client.beta.threads.runs.create(
            thread_id=thread_id,
            assistant_id=assistant_id,
            temperature=0.1,
            top_p=1.0,
            timeout=timeout
        )

        start = time.time()
        while run.status in ['queued', 'in_progress', 'cancelling']:
            time.sleep(1)
            run = client.beta.threads.runs.retrieve(
                thread_id=thread_id, run_id=run.id
            )
            if time.time() - start > timeout:
                client.beta.threads.runs.cancel(
                    thread_id=thread_id, run_id=run.id
                )
                return "Query timed out.", thread_id, run.id, []

        if run.status == "completed":
            messages = client.beta.threads.messages.list(
                thread_id=thread_id, order='desc', limit=10
            )
            response_parts = []
            file_ids = []

            for msg in messages.data:
                if msg.role == "assistant":
                    for c in msg.content:
                        if c.type == "text":
                            response_parts.append(c.text.value)

                            # Check for file annotations in text content
                            # Only collect file_path annotations (generated output files)
                            # Skip file_citation annotations (source references that can't be downloaded)
                            if hasattr(c.text, 'annotations'):
                                for annotation in c.text.annotations:
                                    if hasattr(annotation, 'file_path'):
                                        # These are actual output files from Code Interpreter
                                        file_ids.append(("file", annotation.file_path.file_id))
                                    # NOTE: Intentionally NOT collecting file_citation annotations
                                    # because those reference source files with purpose='assistants'
                                    # which OpenAI does not allow downloading

                        elif c.type == "image_file":
                            file_ids.append(("image", c.image_file.file_id))

                    # Check for file attachments
                    if hasattr(msg, 'attachments') and msg.attachments:
                        for attachment in msg.attachments:
                            if hasattr(attachment, 'file_id'):
                                file_ids.append(("file", attachment.file_id))
                else:
                    break

            return "\n\n".join(reversed(response_parts)), thread_id, run.id, file_ids

        return f"Unexpected status: {run.status}", thread_id, run.id, []

    except Exception as e:
        return f"Error: {str(e)}", thread_id, None, []

# Display chat history
for m in st.session_state.messages:
    with st.chat_message(m["role"]):
        st.markdown(m["content"])

# User input
user_input = st.chat_input("Ask about job listings...")

if user_input:
    prefixed_query = f"{mode_prefix} {user_input}"

    with st.chat_message("user"):
        st.markdown(user_input)

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

    with st.chat_message("assistant"):
        # Retry logic for mode detection failures
        max_retries = 3
        retry_count = 0
        response = None
        thread_id = st.session_state.thread_id
        file_ids = []

        while retry_count <= max_retries:
            with st.spinner("Analyzing..." if retry_count == 0 else f"Retrying ({retry_count}/{max_retries})..."):
                response, thread_id, run_id, file_ids = query_assistant(
                    prefixed_query, api_key, assistant_id, thread_id=None  # Fresh thread on each attempt
                )

            # Check if mode detection failed
            if is_mode_detection_failure(response) and retry_count < max_retries:
                retry_count += 1
                st.warning(f"Mode detection issue detected. Retrying... ({retry_count}/{max_retries})")
                time.sleep(1)  # Brief pause before retry
            else:
                break

        # Show retry info if we had to retry
        if retry_count > 0 and not is_mode_detection_failure(response):
            st.info(f"Query succeeded after {retry_count} retry(ies).")

        # If still failing after all retries, show helpful message to user
        if is_mode_detection_failure(response):
            st.error(
                f"The assistant is having trouble processing this query in **{mode}** mode. "
                f"Please try one of these options:\n"
                f"- Switch to a different analysis method in the sidebar\n"
                f"- Rephrase your question\n"
                f"- Try again in a few moments"
            )

        # Clean response to remove source citations that can't be downloaded
        cleaned_response = clean_response_text(response)
        st.markdown(cleaned_response)

        # Check if user requested a CSV file
        csv_requested = any(keyword in user_input.lower() for keyword in ['csv', 'download', 'export', 'file'])

        # Try to extract tabular data from response and offer CSV download
        if csv_requested:
            df_extracted = extract_table_from_response(cleaned_response)
            if df_extracted is not None and len(df_extracted) > 0:
                # Convert DataFrame to CSV
                csv_buffer = io.StringIO()
                df_extracted.to_csv(csv_buffer, index=False)
                csv_data = csv_buffer.getvalue()

                # Offer download button
                st.success(f"✅ Generated CSV with {len(df_extracted)} rows")
                st.download_button(
                    label="📥 Download CSV File",
                    data=csv_data,
                    file_name="job_listings_export.csv",
                    mime="text/csv",
                    key=f"csv_download_{thread_id}"
                )

        # Handle file downloads from OpenAI (if any)
        if file_ids:
            for file_type, file_id in file_ids:
                if file_type == "file":
                    try:
                        # Retrieve file metadata to get filename
                        file_info = client.files.retrieve(file_id)
                        filename = file_info.filename

                        # Download file content
                        file_content = client.files.content(file_id)
                        file_bytes = file_content.read()

                        # Display download button
                        st.download_button(
                            label=f"📥 Download {filename}",
                            data=file_bytes,
                            file_name=filename,
                            mime="text/csv" if filename.endswith('.csv') else "application/octet-stream"
                        )
                    except Exception as e:
                        st.error(f"Error downloading file: {str(e)}")

        if thread_id and not st.session_state.thread_id:
            st.session_state.thread_id = thread_id

        st.session_state.messages.append({"role": "assistant", "content": cleaned_response})

st.markdown("---")
st.caption("Dataset: 28,257 job listings | Supports CSV, Semantic, and Hybrid methods.")