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 # Model definitions
from models.db import get_session_context
from services.pdf_report import generate_pdf_report # Needs to accept more data
from services.logger import app_logger
from services.metrics import log_report_generated

# --- 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.")

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


def get_user_chat_session_display_data(user_id: int) -> List[Dict[str, Any]]:
    app_logger.debug(f"Fetching session display data for user_id: {user_id}")
    session_data_list: List[Dict[str, Any]] = []
    try:
        with get_session_context() as db:
            stmt = select(ChatSession.id, ChatSession.title, ChatSession.start_time, ChatSession.patient_context_summary)\
                   .where(ChatSession.user_id == user_id).order_by(ChatSession.start_time.desc())
            results = db.exec(stmt).all()
            for row in results:
                session_data_list.append({
                    "id": row.id, "title": row.title, "start_time": row.start_time,
                    "patient_context_summary": row.patient_context_summary
                })
        app_logger.debug(f"Found {len(session_data_list)} session display entries for user {user_id}")
    except Exception as e:
        app_logger.error(f"Error fetching session display data for user {user_id}: {e}", exc_info=True)
        st.error("Could not load your chat sessions.")
    return session_data_list

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.")
    st.stop()

session_options = []
for s_data in chat_session_display_items:
    start_time_obj = s_data["start_time"]
    start_time_display = start_time_obj.strftime('%Y-%m-%d %H:%M') if start_time_obj else "N/A"
    title_display = s_data["title"] if s_data["title"] else f"Consult on {start_time_display}"
    display_string = f"ID: {s_data['id']} | Started: {start_time_display} | Title: {title_display}"
    session_options.append((display_string, s_data['id']))

selected_option_tuple = st.selectbox(
    "Select a Consultation Session:", options=session_options, format_func=lambda x: x[0], index=0
)

if selected_option_tuple:
    selected_session_id = selected_option_tuple[1]
    MESSAGES_KEY = f"report_messages_for_session_{selected_session_id}"

    if MESSAGES_KEY not in st.session_state:
        app_logger.info(f"Fetching messages for report, session ID: {selected_session_id}")
        msg_data_list: List[Dict[str, Any]] = []
        try:
            with get_session_context() as db:
                stmt = select(ChatMessage.role, ChatMessage.content, ChatMessage.timestamp, ChatMessage.tool_name)\
                       .where(ChatMessage.session_id == selected_session_id).order_by(ChatMessage.timestamp)
                results = db.exec(stmt).all()
                for row in results:
                    msg_data_list.append({
                        "role": row.role, "content": row.content, "timestamp": row.timestamp,
                        "tool_name": getattr(row, 'tool_name', None)
                    })
            st.session_state[MESSAGES_KEY] = msg_data_list
        except Exception as e:
            app_logger.error(f"Error fetching messages for report (session {selected_session_id}): {e}", exc_info=True)
            st.error("Could not load messages for this report.")
            st.session_state[MESSAGES_KEY] = []
    else:
        app_logger.debug(f"Using cached messages from session_state for report, session ID: {selected_session_id}")
        msg_data_list = st.session_state[MESSAGES_KEY]

    selected_session_data = next((s for s in chat_session_display_items if s['id'] == selected_session_id), None)

    if selected_session_data:
        st.markdown("---")
        st.subheader(f"Report Preview for Session ID: {selected_session_data['id']}")
        start_t = selected_session_data["start_time"]
        st.write(f"**Started:** {start_t.strftime('%Y-%m-%d %H:%M:%S UTC') if start_t else 'N/A'}")
        st.write(f"**Title:** {selected_session_data['title'] or 'Untitled Session'}")
        if selected_session_data.get("patient_context_summary"):
            with st.expander("View Patient Context Provided", expanded=False):
                st.markdown(selected_session_data["patient_context_summary"])
        
        if msg_data_list:
            with st.expander("View Chat Transcript", expanded=False):
                for idx, msg_d in enumerate(msg_data_list):
                    if msg_d["role"] == "system": continue # Don't show system context messages here
                    icon = "🧑‍⚕️" if msg_d["role"] == "assistant" else "👤"
                    if msg_d["role"] == "tool": icon = "🛠️"
                    ts_display = msg_d["timestamp"].strftime('%Y-%m-%d %H:%M:%S') if msg_d["timestamp"] else "N/A"
                    st.markdown(f"**{icon} {msg_d['role'].capitalize()}** ({ts_display})")
                    st.markdown(f"> ```\n{msg_d['content']}\n```")
                    if idx < len(msg_data_list) - 1: st.markdown("---")
            
            st.markdown("---")
            try:
                # Prepare data for PDF report
                # This creates temporary ChatMessage-like objects.
                # Ensure generate_pdf_report can handle these non-ORM objects.
                pdf_report_messages = [
                    ChatMessage(**msg_d) for msg_d in msg_data_list # Spread dict into ChatMessage constructor
                ]
                report_data_for_pdf = {
                    "patient_name": authenticated_username,
                    "session_id": selected_session_data['id'],
                    "session_title": selected_session_data['title'] or 'Untitled Session',
                    "session_start_time": selected_session_data['start_time'],
                    "patient_context_summary": selected_session_data.get('patient_context_summary', "Not provided."),
                    "messages": pdf_report_messages
                }

                pdf_bytes = generate_pdf_report(report_data_for_pdf) # Pass the whole dict
                
                file_ts = datetime.now().strftime('%Y%m%d_%H%M%S')
                pdf_fn = f"Consult_Report_S{selected_session_id}_{file_ts}.pdf"
                st.download_button(
                    label="Download Report as PDF", data=pdf_bytes, file_name=pdf_fn, mime="application/pdf",
                    key=f"dl_pdf_s{selected_session_id}", on_click=log_report_generated,
                    args=(authenticated_user_id, selected_session_id), help="Download consultation summary."
                )
            except Exception as e:
                app_logger.error(f"Error generating PDF for session {selected_session_id}: {e}", exc_info=True)
                st.error(f"Could not generate PDF report. Error: {type(e).__name__}")
        else: st.info("This session has no messages for the report.")
    else: st.error("Selected session data not found.")
else: st.info("Select a session to view details.")