pages/3_Reports.py

# /home/user/app/pages/3_Reports.py
import streamlit as st
from datetime import datetime
from typing import List
from sqlmodel import select # For SQLModel queries

from config.settings import settings
from models import ChatMessage, ChatSession # User model not directly needed here if ID is sufficient
from models.db import get_session_context # SQLModel session context
from services.pdf_report import generate_pdf_report
from services.logger import app_logger
from services.metrics import log_report_generated # Assuming this function exists and is set up

# --- Page Configuration (Typically set in main app.py) ---
# st.set_page_config(page_title=f"Reports - {settings.APP_TITLE}", layout="wide")

# --- Authentication Check ---
if not st.session_state.get("authenticated_user_id"):
    st.warning("Please log in to access reports.")
    try:
        st.switch_page("app.py") # Redirect to the main login page
    except st.errors.StreamlitAPIException as e:
        # Handle cases where switch_page might not work (e.g., not in MPA mode during dev)
        app_logger.warning(f"Reports: Could not switch page (maybe not in MPA mode or other issue): {e}")
        st.info("Please navigate to the main login page manually.")
    st.stop() # Halt script execution for unauthenticated users

# --- Get Authenticated User Info ---
authenticated_user_id = st.session_state.get("authenticated_user_id")
authenticated_username = st.session_state.get("authenticated_username", "User") # Default to "User"
app_logger.info(f"User '{authenticated_username}' (ID: {authenticated_user_id}) accessed Reports page.")

# --- Page Title and Description ---
st.title("Consultation Reports")
st.markdown("View and download your past consultation sessions.")

# --- Helper Function to Load User's Chat Sessions ---
# @st.cache_data(ttl=300) # Cache for 5 minutes. Invalidate if session titles change or new sessions added frequently.
# For caching to work correctly, arguments to the cached function must be hashable. user_id (int) is.
def get_user_chat_sessions(user_id: int) -> List[ChatSession]:
    """Fetches all chat sessions for a given user ID, ordered by start time descending."""
    app_logger.debug(f"Fetching chat sessions for user_id: {user_id}")
    sessions: List[ChatSession] = []
    try:
        with get_session_context() as db_session: # db_session is a SQLModel Session
            statement = select(ChatSession).where(ChatSession.user_id == user_id).order_by(ChatSession.start_time.desc())
            results = db_session.exec(statement)
            sessions = results.all()
        app_logger.debug(f"Found {len(sessions)} sessions for user_id: {user_id}")
    except Exception as e:
        app_logger.error(f"Error fetching chat sessions for user_id {user_id}: {e}", exc_info=True)
        st.error("Could not load your chat sessions. Please try again later.")
    return sessions

# Fetch the chat sessions for the current authenticated user
chat_sessions = get_user_chat_sessions(authenticated_user_id)

# --- Display Logic ---
if not chat_sessions:
    st.info("You have no past consultation sessions recorded.")
    st.stop() # No further processing if no sessions

# Prepare options for the selectbox: (Display String, Session ID)
session_options = []
for s in chat_sessions:
    start_time_display = s.start_time.strftime('%Y-%m-%d %H:%M') if s.start_time else "Date N/A"
    # Use a more descriptive title if the session title is None
    title_display = s.title if s.title else f"Consultation on {start_time_display}"
    display_string = f"ID: {s.id} | Started: {start_time_display} | Title: {title_display}"
    session_options.append((display_string, s.id)) # Store (display_text, actual_id)

# Let the user select a session
selected_option_tuple = st.selectbox(
    "Select a Consultation Session:",
    options=session_options,
    format_func=lambda x: x[0], # Show only the display string in the selectbox
    index=0 # Default to the first (most recent) session
)

if selected_option_tuple:
    selected_session_id = selected_option_tuple[1] # Extract the session ID
    app_logger.info(f"User '{authenticated_username}' selected session ID: {selected_session_id} for report viewing.")

    # Find the full ChatSession object from the already fetched list
    selected_session_object = next((s for s in chat_sessions if s.id == selected_session_id), None)

    if selected_session_object:
        st.markdown("---") # Visual separator
        st.subheader(f"Details for Session ID: {selected_session_object.id}")

        start_time_detail = selected_session_object.start_time.strftime('%Y-%m-%d %H:%M:%S UTC') if selected_session_object.start_time else "Not recorded"
        st.write(f"**Started:** {start_time_detail}")
        st.write(f"**Title:** {selected_session_object.title or 'Untitled Session'}")

        # Fetch messages for the selected session
        messages_for_report: List[ChatMessage] = []
        try:
            with get_session_context() as db_session:
                statement = select(ChatMessage).where(ChatMessage.session_id == selected_session_id).order_by(ChatMessage.timestamp)
                results = db_session.exec(statement)
                messages_for_report = results.all()
        except Exception as e:
            app_logger.error(f"Error fetching messages for session {selected_session_id}: {e}", exc_info=True)
            st.error("Could not load messages for this session.")

        if messages_for_report:
            with st.expander("View Chat Transcript", expanded=False):
                for msg_idx, msg in enumerate(messages_for_report):
                    icon = "🧑‍⚕️" if msg.role == "assistant" else "👤"
                    if msg.role == "tool": icon = "🛠️"
                    timestamp_display = msg.timestamp.strftime('%Y-%m-%d %H:%M:%S') if msg.timestamp else "Time N/A"

                    # Displaying role and timestamp
                    st.markdown(f"**{icon} {msg.role.capitalize()}** ({timestamp_display})")
                    # Displaying content in a blockquote or code block for better readability
                    st.markdown(f"> ```\n{msg.content}\n```") # Markdown code block
                    if msg_idx < len(messages_for_report) - 1:
                        st.markdown("---") # Separator between messages

            # PDF Download Button
            st.markdown("---") # Visual separator before download button
            try:
                # Use the authenticated username as the patient name for the report
                pdf_bytes = generate_pdf_report(messages_for_report, patient_name=authenticated_username)

                current_time_str = datetime.now().strftime('%Y%m%d_%H%M%S')
                pdf_file_name = f"Consultation_Report_Session{selected_session_id}_{current_time_str}.pdf"

                st.download_button(
                    label="Download Report as PDF",
                    data=pdf_bytes,
                    file_name=pdf_file_name,
                    mime="application/pdf",
                    key=f"download_pdf_btn_{selected_session_id}", # Unique key for the button
                    on_click=log_report_generated,
                    # Pass necessary arguments to the on_click callback if it needs them
                    args=(authenticated_user_id, selected_session_id), # Example arguments
                    help="Click to download the consultation transcript as a PDF file."
                )
            except Exception as e:
                app_logger.error(f"Error generating PDF for session {selected_session_id} (User: '{authenticated_username}'): {e}", exc_info=True)
                st.error(f"Could not generate PDF report for this session. Error: {type(e).__name__}")
        else:
            st.info("This session has no messages recorded.")
    else:
        # This case should ideally not happen if selected_session_id comes from chat_sessions
        app_logger.error(f"Selected session ID {selected_session_id} not found in fetched sessions for user '{authenticated_username}'.")
        st.error("Selected session could not be found. Please try again.")
else:
    # This case handles if session_options is empty, though the earlier check for `if not chat_sessions` should catch it.
    st.info("Please select a session from the dropdown menu to view details.")