# /home/user/app/pages/3_Reports.py
import streamlit as st
from datetime import datetime
from typing import List
from sqlmodel import select # <--- IMPORT SELECT FOR SQLMODEL QUERIES

from config.settings import settings
from models import ChatMessage, ChatSession, User # User not directly used if ID is sufficient
from models.db 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

# --- Auth Check (same as before) ---
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.")
    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.")

# --- Load User's Chat Sessions ---
# @st.cache_data(ttl=300) # Consider re-enabling with SQLModel if args are hashable
def get_user_chat_sessions(user_id: int) -> List[ChatSession]:
    app_logger.debug(f"Fetching chat sessions for user_id: {user_id}")
    with get_session_context() as db: # db is a SQLModel Session
        # --- SQLMODEL QUERY ---
        statement = select(ChatSession).where(ChatSession.user_id == user_id).order_by(ChatSession.start_time.desc())
        results = db.exec(statement)
        sessions = results.all()
        # --------------------
        app_logger.debug(f"Found {len(sessions)} sessions for user_id: {user_id}")
        return sessions

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 (largely same, but ensure messages query is updated) ---
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]
)

if selected_option_tuple:
    selected_session_id = selected_option_tuple[1]
    app_logger.info(f"User {authenticated_username} selected session ID: {selected_session_id} for report.")
    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}")
        # ... (display session details - same as before)

        with get_session_context() as db: # db is a SQLModel Session
            # --- SQLMODEL QUERY ---
            statement = select(ChatMessage).where(ChatMessage.session_id == selected_session.id).order_by(ChatMessage.timestamp)
            results = db.exec(statement)
            messages: List[ChatMessage] = results.all()
            # --------------------

        if messages:
            # ... (display transcript and PDF download - same as before)
            # The generate_pdf_report(messages, ...) part doesn't change.
            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}):**")
                    st.markdown(f"> ```\n{msg.content}\n```") # Using markdown code block for content
                    if msg_idx < len(messages) - 1:
                        st.markdown("---")

            st.markdown("---")
            try:
                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}",
                    on_click=log_report_generated,
                    args=(authenticated_user_id, selected_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: {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.")

# (Ensure the rest of 3_Reports.py from the previous good version is here if anything was omitted for brevity)
# Example for start_time display in selected_session details:
if selected_session:
    # ... other details
    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'}")
    # ... rest of message display and PDF generation ...