pages/3_Reports.py

# /home/user/app/pages/3_Reports.py
import streamlit as st
from datetime import datetime
from typing import List, Dict, Any
from sqlmodel import select

from config.settings import settings
from models import ChatMessage, ChatSession # SQLModel classes for DB query
from models.db import get_session_context
# Import generate_pdf_report AND MockChatMessage from services.pdf_report
from services.pdf_report import generate_pdf_report, MockChatMessage
from services.logger import app_logger
from services.metrics import log_report_generated # Assuming this function exists

# --- 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")
    except st.errors.StreamlitAPIException:
        st.info("Please navigate to the main login page manually.")
    st.stop()

authenticated_user_id = st.session_state.get("authenticated_user_id")
authenticated_username = st.session_state.get("authenticated_username", "User")
app_logger.info(f"User '{authenticated_username}' (ID: {authenticated_user_id}) accessed Reports page.")

# --- Page Title and Info ---
st.title("Consultation Reports")
st.markdown("View and download your past consultation sessions.")
st.info(f"{settings.MAIN_DISCLAIMER_SHORT} PDFs are summaries for review.")

# --- Helper Function to Load User's Chat Session Display Data ---
# @st.cache_data(ttl=300, show_spinner=False) # Consider caching if fetching is slow and data doesn't change too rapidly
def get_user_chat_session_display_data(user_id: int) -> List[Dict[str, Any]]:
    """Fetches essential data for chat sessions (ID, title, start_time, context) as dictionaries."""
    app_logger.debug(f"ReportsPage: Fetching session display data for user_id: {user_id}")
    session_data_list: List[Dict[str, Any]] = []
    try:
        with get_session_context() as db_session:
            statement = select(
                ChatSession.id,
                ChatSession.title,
                ChatSession.start_time,
                ChatSession.patient_context_summary # Fetch this as well
            ).where(ChatSession.user_id == user_id).order_by(ChatSession.start_time.desc())
            
            results = db_session.exec(statement).all() # List of Row objects
            for row_item in results:
                session_data_list.append({
                    "id": row_item.id,
                    "title": row_item.title,
                    "start_time": row_item.start_time, # datetime object
                    "patient_context_summary": row_item.patient_context_summary
                })
        app_logger.debug(f"ReportsPage: Found {len(session_data_list)} session display entries for user {user_id}")
    except Exception as e:
        app_logger.error(f"ReportsPage: Error fetching session display data for user {user_id}: {e}", exc_info=True)
        st.error("Could not load your chat sessions. Please try again later.")
    return session_data_list

# Fetch session display data
chat_session_display_items = get_user_chat_session_display_data(authenticated_user_id)

if not chat_session_display_items:
    st.info("You have no past consultation sessions recorded to display.")
    st.stop()

# --- UI: Select a Session ---
session_options_for_selectbox = []
for s_data_item in chat_session_display_items:
    start_time = s_data_item.get("start_time")
    start_time_str = start_time.strftime('%Y-%m-%d %H:%M') if start_time else "Date N/A"
    title_str = s_data_item.get("title") or f"Consultation on {start_time_str}"
    display_text = f"ID: {s_data_item.get('id')} | Started: {start_time_str} | Title: {title_str}"
    session_options_for_selectbox.append((display_text, s_data_item.get('id'))) # (Display, Value)

selected_option_tuple = st.selectbox(
    "Select a Consultation Session:",
    options=session_options_for_selectbox,
    format_func=lambda x_tuple: x_tuple[0], # Show only the display string
    index=0, # Default to the first (most recent) session
    key="report_session_selector"
)

# --- Display Details and PDF Download for Selected Session ---
if selected_option_tuple:
    selected_session_id_val = selected_option_tuple[1] # The actual session ID
    app_logger.info(f"ReportsPage: User '{authenticated_username}' selected session ID: {selected_session_id_val} for viewing.")

    # Find the full data for the selected session from our fetched list of dicts
    selected_session_dict_data = next(
        (s_item for s_item in chat_session_display_items if s_item['id'] == selected_session_id_val), None
    )

    if selected_session_dict_data:
        st.markdown("---")
        st.subheader(f"Report Preview for Session ID: {selected_session_dict_data['id']}")

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

        patient_context_from_session = selected_session_dict_data.get("patient_context_summary")
        if patient_context_from_session:
            with st.expander("View Patient Context Provided for this Session", expanded=False):
                st.markdown(patient_context_from_session)
        
        # --- Fetch and Display Messages for the Selected Session ---
        # Store fetched messages in session_state to avoid re-fetching on every interaction within the page
        MESSAGES_DATA_SESS_KEY = f"report_page_messages_for_session_{selected_session_id_val}"

        if MESSAGES_DATA_SESS_KEY not in st.session_state:
            app_logger.info(f"ReportsPage: Fetching DB messages for session ID: {selected_session_id_val}")
            db_messages_as_dicts: List[Dict[str, Any]] = []
            try:
                with get_session_context() as db:
                    # Fetch specific columns needed for display and PDF
                    stmt = select(
                        ChatMessage.role, ChatMessage.content, ChatMessage.timestamp, ChatMessage.tool_name
                    ).where(ChatMessage.session_id == selected_session_id_val).order_by(ChatMessage.timestamp)
                    
                    message_results = db.exec(stmt).all() # List of Row objects
                    for msg_row in message_results:
                        db_messages_as_dicts.append({
                            "role": msg_row.role, "content": msg_row.content,
                            "timestamp": msg_row.timestamp, # datetime object
                            "tool_name": getattr(msg_row, 'tool_name', None)
                        })
                st.session_state[MESSAGES_DATA_SESS_KEY] = db_messages_as_dicts
                app_logger.info(f"ReportsPage: Fetched and stored {len(db_messages_as_dicts)} messages for session {selected_session_id_val} in st.session_state.")
            except Exception as e_msg_fetch:
                app_logger.error(f"ReportsPage: Error fetching DB messages for session {selected_session_id_val}: {e_msg_fetch}", exc_info=True)
                st.error("Could not load messages for this session's transcript.")
                st.session_state[MESSAGES_DATA_SESS_KEY] = [] # Store empty list on error
        else:
            app_logger.debug(f"ReportsPage: Using cached messages from st.session_state for session ID: {selected_session_id_val}")
        
        messages_to_display_and_pdf = st.session_state[MESSAGES_DATA_SESS_KEY]

        if messages_to_display_and_pdf:
            with st.expander("View Chat Transcript", expanded=False):
                for msg_idx, msg_dict_item in enumerate(messages_to_display_and_pdf):
                    msg_role = str(msg_dict_item.get("role", "N/A"))
                    msg_content = str(msg_dict_item.get("content", ""))
                    msg_timestamp_obj = msg_dict_item.get("timestamp")
                    
                    # Skip verbose system messages about context in transcript view
                    if msg_role == "system" and "Initial Patient Context Set:" in msg_content:
                        continue

                    icon_char = "🧑‍⚕️" if msg_role == "assistant" else "👤"
                    if msg_role == "tool": icon_char = "🛠️"
                    if msg_role == "system": icon_char = "⚙️"
                    
                    timestamp_str_display = msg_timestamp_obj.strftime('%Y-%m-%d %H:%M:%S') if msg_timestamp_obj else "N/A"
                    
                    st.markdown(f"**{icon_char} {msg_role.capitalize()}** ({timestamp_str_display})")
                    st.markdown(f"> ```\n{msg_content}\n```") # Using markdown code block for content
                    if msg_idx < len(messages_to_display_and_pdf) - 1:
                        st.markdown("---")
            
            # --- PDF Download Button ---
            st.markdown("---")
            try:
                # Prepare list of MockChatMessage instances for the PDF generator
                pdf_report_mock_messages: List[MockChatMessage] = []
                for msg_d_item in messages_to_display_and_pdf:
                    pdf_report_mock_messages.append(
                        MockChatMessage( # Use the imported MockChatMessage class
                            role=msg_d_item.get("role"),
                            content=msg_d_item.get("content"),
                            timestamp=msg_d_item.get("timestamp"),
                            tool_name=msg_d_item.get("tool_name")
                        )
                    )

                # Prepare the consolidated data dictionary for generate_pdf_report
                report_data_for_pdf_generation = {
                    "patient_name": authenticated_username, # This is the clinician's username
                    "session_id": selected_session_dict_data['id'],
                    "session_title": selected_session_dict_data.get('title') or 'Untitled Session',
                    "session_start_time": selected_session_dict_data.get('start_time'), # datetime object
                    "patient_context_summary": selected_session_dict_data.get('patient_context_summary', "Not provided."),
                    "messages": pdf_report_mock_messages # List of MockChatMessage instances
                }
                
                app_logger.debug(f"ReportsPage: Data for PDF generation (session {selected_session_id_val}): Keys={list(report_data_for_pdf_generation.keys())}, NumMessages={len(pdf_report_mock_messages)}")

                pdf_bytes_output = generate_pdf_report(report_data_for_pdf_generation)
                
                current_time_filename_str = datetime.now().strftime('%Y%m%d_%H%M%S')
                pdf_report_filename = f"Consultation_Report_S{selected_session_id_val}_{current_time_filename_str}.pdf"

                st.download_button(
                    label="Download Report as PDF",
                    data=pdf_bytes_output,
                    file_name=pdf_report_filename,
                    mime="application/pdf",
                    key=f"download_pdf_button_session_{selected_session_id_val}", # Unique key
                    on_click=log_report_generated,
                    args=(authenticated_user_id, selected_session_id_val), # Example args for metrics
                    help="Click to download a PDF summary of this consultation session."
                )
            except Exception as e_pdf_gen:
                app_logger.error(f"ReportsPage: Error generating PDF for session {selected_session_id_val} (User: '{authenticated_username}'): {e_pdf_gen}", exc_info=True)
                st.error(f"Could not generate PDF report at this time. Error: {type(e_pdf_gen).__name__}")
        else:
            st.info("This session has no messages recorded to include in the report.")
    else:
        app_logger.error(f"ReportsPage: Selected session ID {selected_session_id_val} not found in fetched display data for user '{authenticated_username}'. This is unexpected.")
        st.error("An error occurred: The selected session data could not be retrieved. Please try again or select another session.")
else:
    st.info("Please select a session from the dropdown menu above to view details and generate a report.")