workflow/report/report_render.py

import datetime
import io
from io import BytesIO

from stqdm import stqdm
import mammoth
import numpy as np
import pandas as pd
import streamlit as st
import streamlit_antd_components as sac
from concurrent.futures import ThreadPoolExecutor, as_completed

from prompt_engineer.sec5_call_llm import *
from workflow.report.report_utils import html_dowmload
from workflow.report.report_html import write_html
from workflow.report.report_word import write_word
from workflow.report.report_markdown import write_markdown
from workflow.report.report_prepare_er import report_prepare


def report_save(agents, auto): 

    report_agent = agents[-1]
    action = report_agent.load_report_format()
    
    if report_agent.load_report_format() == 'HTML':
        not_generate = report_agent.html == None
    if report_agent.load_report_format() == 'Word':
        not_generate = report_agent.word == None
    if report_agent.load_report_format() == 'Markdown':
        not_generate = report_agent.markdown == None

    mode = report_agent.load_gen_mode()
    parallel = (mode == "并行")
    
    if st.button(f"📝 生成 {action} 报告") or (auto and not_generate):
        with st.spinner(f"正在生成 {action} 报告..."):

            report_prepare(agents, parallel=parallel)

            if report_agent.load_report_format() == 'Word':
                write_word(agents)
            elif report_agent.load_report_format() == 'HTML':
                write_html(agents)
            elif report_agent.load_report_format() == 'Markdown':
                write_markdown(agents)


def report_basic_info(agent, auto) -> None:
    outline_length = sac.segmented(
        items=[
            sac.SegmentedItem(label='简要'),
            sac.SegmentedItem(label='标准'),
            sac.SegmentedItem(label='详细'),
        ],
        label='详细程度', index=1, align='center',
        size='sm', radius='sm', use_container_width=True
    )
    agent.save_outline_length(outline_length)

    c1, c2 = st.columns([3, 1])
    with c1:
        report_format = sac.chip(
            items=[
                sac.ChipItem(label='Word', icon=sac.BsIcon(name='file-earmark-word', size=16)),
                sac.ChipItem(label='HTML', icon=sac.BsIcon(name='filetype-html', size=16)),
                sac.ChipItem(label='Markdown', icon=sac.BsIcon(name='file-earmark-code', size=16)),
            ],
            label='选择报告生成格式', index=[0, 2],
            align='start', radius='md', multiple=False,
        )
        agent.save_report_format(report_format)

    with c2:
        mode = sac.segmented(
            items=[
                sac.SegmentedItem(label='并行'),
                sac.SegmentedItem(label='串行'),
            ],
            label='生成模式', align='end', size='sm',
            use_container_width=True, radius='md'
        )
        agent.save_gen_mode(mode)

    user_input = st.text_input("报告生成要求", "默认")
    agent.save_user_input(user_input)

    not_generated = report_agent.load_outline() is None

    # === 并行生成目录 ===
    if st.button("🗒️ 生成目录") or (auto and not_generated):
        with st.spinner("⏳ 正在自动生成目录结构..."):
            summaries = []

            # === 保存当前 Streamlit 状态副本 ===
            session_snapshot = dict(st.session_state)

            def process_summary(idx, sub_agent, session_snapshot):
                """并行执行 summary_html/summary_word（带状态复制）"""
                # 恢复 session_state
                for k, v in session_snapshot.items():
                    st.session_state[k] = v

                # 实际生成逻辑
                if hasattr(sub_agent, "summary_html"):
                    summary = sub_agent.summary_html()
                else:
                    summary = sub_agent.summary_word()

                return idx, summary

            max_workers = min(6, len(agents) - 1)
            results = []

            with ThreadPoolExecutor(max_workers=max_workers) as executor:
                futures = {
                    executor.submit(process_summary, i, sub_agent, session_snapshot): i
                    for i, sub_agent in enumerate(agents[:-1])
                }

                for future in stqdm(as_completed(futures), total=len(futures)):
                    try:
                        idx, summary = future.result()
                        if summary:
                            results.append((idx, summary))
                    except Exception as e:
                        print(f"子模块摘要生成失败: {e}")

            # === 恢复章节原顺序 ===
            results.sort(key=lambda x: x[0])
            summaries = [summary for _, summary in results if summary]

            # === 生成目录 ===
            default_toc = report_agent.generate_toc_from_summary(summaries)
            report_agent.save_outline(default_toc)


def report_outline(agents):

    st.subheader("目录结构预览与编辑")
    load_agent, preproc_agent, visual_agent, coding_agent, report_agent = agents[0], agents[1], agents[2], agents[3], agents[4]

    default_toc = report_agent.load_outline()

    toc_md = st.text_area(
        "您可以在此处编辑目录结构", 
        value=default_toc, 
        height=260
    )

    report_agent.save_outline(toc_md)


def report_execution(report_agent):

    if report_agent.load_report_format() == 'Word':

        full_report = report_agent.load_word()
        if full_report is not None:
            st.download_button(
                label="⬇️ 下载 Word 报告",
                data=full_report,
                file_name="report.docx",
                mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document"
            )

    elif report_agent.load_report_format() == 'HTML':

        full_report = report_agent.load_html()

        if full_report is not None:
            st.download_button(
                label="⬇️ 下载 HTML 报告",
                data=full_report.encode("utf-8"),
                file_name="report.html",
                mime="text/html",
            )

            if st.button("⬇️ 下载 PDF 报告"):
                html_dowmload(full_report)
    elif report_agent.load_report_format() == 'Markdown':

        full_report = report_agent.load_markdown()
        if full_report is not None:

            # 提供下载按钮
            st.download_button(
                label="⬇️ 下载 Markdown 报告",
                data=full_report,
                file_name="report.md",
                mime="text/markdown"
            )


if __name__ == "__main__":

    st.title("报告生成")

    st.markdown("---")

    load_agent   = st.session_state.data_loading_agent
    preproc_agent = st.session_state.data_preprocess_agent
    visual_agent = st.session_state.visualization_agent
    coding_agent = st.session_state.modeling_coding_agent
    planner = st.session_state.planner_agent
    auto = planner.report_auto

    processed_df = preproc_agent.load_processed_df()
    if processed_df is None:
        df = load_agent.load_df()
    else:
        df = processed_df

    if df is None:
        st.warning("⚠️ 请先在数据导入页面加载数据")
        st.stop()

    if isinstance(df, np.ndarray):
        df = pd.DataFrame(df)

    df_shuffled = df.sample(frac=1, random_state=42).reset_index(drop=True)

    report_agent = st.session_state.report_agent
    report_agent.add_df(df_shuffled)

    agents = [load_agent, preproc_agent, visual_agent, coding_agent, report_agent]

    c = st.columns(2)
    with c[0].expander('报告设置', True):
        report_basic_info(report_agent, auto)

    with c[1].expander('报告大纲', True):
        report_outline(agents)
        report_save(agents, auto)
        report_execution(report_agent)