pages/3_Reports.py

# /home/user/app/pages/3_Reports.py
import streamlit as st
from datetime import datetime
from typing import List # For type hinting

from config.settings import settings
from models import ChatMessage, ChatSession, User # User might not be needed if ID is used
from models.db import get_session_context # Or from models import get_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

# Page config typically in 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")
    except st.errors.StreamlitAPIException as e:
        if "st.switch_page can only be called when running in MPA mode" in str(e):
            app_logger.warning("Reports: Running in single-page mode or st.switch_page issue. Stopping script.")
            st.info("Please navigate to the main login page.")
        else:
            app_logger.error(f"Reports: Error during st.switch_page: {e}")
            st.error("Redirection error. Please go to the login page manually.")
    st.stop()

# Get authenticated user's ID and username
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.")


st.title("Consultation Reports")
st.markdown("View and download your past consultation sessions.")

# --- Load User's Chat Sessions ---
# @st.cache_data(ttl=300) # Cache for 5 minutes; ensure function args are hashable
def get_user_chat_sessions(user_id: int) -> List[ChatSession]: # Return type hint
    app_logger.debug(f"Fetching chat sessions for user_id: {user_id}")
    with get_session_context() as db:
        # If using SQLModel: from sqlmodel import select
        # sessions = db.exec(select(ChatSession).where(ChatSession.user_id == user_id).order_by(ChatSession.start_time.desc())).all()
        sessions = db.query(ChatSession).filter(ChatSession.user_id == user_id).order_by(ChatSession.start_time.desc()).all()
        app_logger.debug(f"Found {len(sessions)} sessions for user_id: {user_id}")
        return sessions

# Fetch sessions (caching is good, but ensure it invalidates appropriately if sessions are added/updated frequently)
# For simplicity without complex cache invalidation, direct fetch on each load:
chat_sessions = get_user_chat_sessions(authenticated_user_id)

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

# --- Display Sessions and Download Option ---
# Create a list of tuples for 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 "N/A"
    title_display = s.title or f"Session on {start_time_display}"
    display_string = f"ID: {s.id} | Started: {start_time_display} | Title: {title_display}"
    session_options.append((display_string, s.id))


selected_option_tuple = st.selectbox(
    "Select a Consultation Session:",
    options=session_options,
    format_func=lambda x: x[0] # Display the descriptive string from the tuple
)

if selected_option_tuple:
    selected_session_id = selected_option_tuple[1] # Get the session_id from the tuple
    app_logger.info(f"User {authenticated_username} selected session ID: {selected_session_id} for report.")

    # Find the selected session object (already fetched in chat_sessions)
    selected_session = next((s for s in chat_sessions if s.id == selected_session_id), None)

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

        # Fetch messages for the selected session
        with get_session_context() as db:
            # If using SQLModel: from sqlmodel import select
            # messages = db.exec(select(ChatMessage).where(ChatMessage.session_id == selected_session.id).order_by(ChatMessage.timestamp)).all()
            messages: List[ChatMessage] = db.query(ChatMessage).filter(ChatMessage.session_id == selected_session.id).order_by(ChatMessage.timestamp).all()

        if messages:
            with st.expander("View Chat Transcript", expanded=False):
                for msg_idx, msg in enumerate(messages):
                    icon = "🧑‍⚕️" if msg.role == "assistant" else "👤"
                    if msg.role == "tool": icon = "🛠️"
                    timestamp_display = msg.timestamp.strftime('%H:%M:%S') if msg.timestamp else "N/A"
                    st.markdown(f"**{icon} {msg.role.capitalize()} ({timestamp_display}):**")
                    # Use st.code or st.markdown for content display. Using a quote block.
                    st.markdown(f"> ```\n{msg.content}\n```")
                    if msg_idx < len(messages) - 1:
                        st.markdown("---") # Separator between messages

            # PDF Download Button
            st.markdown("---") # Separator before download button
            try:
                # Use authenticated_username for patient_name, or a more specific field if available
                pdf_bytes = generate_pdf_report(messages, patient_name=authenticated_username)

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

                st.download_button(
                    label="Download Report as PDF",
                    data=pdf_bytes,
                    file_name=pdf_file_name,
                    mime="application/pdf",
                    key=f"download_pdf_{selected_session.id}", # Unique key for the button
                    on_click=log_report_generated, # Pass any necessary args if log_report_generated needs them
                    args=(authenticated_user_id, selected_session_id) # Example: if metrics logger needs user_id and session_id
                )
            except Exception as e:
                app_logger.error(f"Error generating PDF for session {selected_session.id} for user {authenticated_username}: {e}", exc_info=True)
                st.error(f"Could not generate PDF report at this time. Error: {type(e).__name__}")
        else:
            st.info("This session has no messages recorded.")
else:
    st.info("Select a session from the dropdown to view details and download the report.")