"""
Shifu publish funcs

This module contains functions for publishing shifu.

Author: yfge
Date: 2025-08-07
"""

from flaskr.service.shifu.shifu_draft_funcs import get_latest_shifu_draft
from flaskr.service.common import raise_error
from flaskr.dao import db
from flaskr.service.shifu.models import (
    PublishedShifu,
    PublishedOutlineItem,
    DraftOutlineItem,
    LogPublishedStruct,
)
from flaskr.service.shifu.shifu_outline_funcs import (
    build_outline_tree,
    ShifuOutlineTreeNode,
)
from flaskr.service.shifu.shifu_history_manager import HistoryItem
from flaskr.service.shifu.shifu_struct_manager import get_shifu_outline_tree
from flaskr.util import generate_id
from datetime import datetime
import threading
import queue
from flaskr.service.shifu.shifu_struct_manager import ShifuInfoDto
from flaskr.api.llm import invoke_llm
from flaskr.api.langfuse import (
    create_trace_with_root_span,
    finalize_langfuse_trace,
    get_langfuse_client,
)
from flaskr.service.metering import UsageContext
from flaskr.service.metering.consts import BILL_USAGE_SCENE_DEBUG
from flaskr.util.prompt_loader import load_prompt_template
from flaskr.service.shifu.consts import (
    ASK_MODE_ENABLE,
)
from markdown_flow import (
    MarkdownFlow,
    BlockType,
)
from flaskr.common.i18n_utils import get_markdownflow_output_language
from flaskr.common.shifu_context import (
    get_shifu_context_snapshot,
    apply_shifu_context_snapshot,
)


def _build_frontend_url(base_url: str, path: str) -> str:
    """
    Build a frontend URL based on the provided base URL.
    """
    normalized_base = base_url.rstrip("/") if base_url else ""
    cleaned_path = path if path.startswith("/") else f"/{path}"
    return f"{normalized_base}{cleaned_path}" if normalized_base else cleaned_path


def preview_shifu_draft(
    app, user_id: str, shifu_id: str, variables: dict, base_url: str
):
    """
    Preview shifu draft
    Args:
        app: Flask application instance
        user_id: User ID
        shifu_id: Shifu ID
        variables: Variables
        base_url: Base URL to build preview link
    """
    with app.app_context():
        shifu_draft = get_latest_shifu_draft(shifu_id)
        if not shifu_draft:
            raise_error("server.shifu.shifuNotFound")

        return _build_frontend_url(base_url, f"/c/{shifu_id}?preview=true")


def publish_shifu_draft(
    app,
    user_id: str,
    shifu_id: str,
    base_url: str,
    sync_summary: bool = False,
):
    """
    Publish shifu draft
    will copy all draft data to published data
    and save history to database
    and run summary generation in background by default
    and return published shifu url
    Args:
        app: Flask application instance
        user_id: User ID
        shifu_id: Shifu ID
        base_url: Base URL to build published link
        sync_summary: If True, generate summary/ask prompts synchronously in the
            current process (useful for one-off console commands). Default False
            keeps existing async background behavior.
    Returns:
        str: Shifu published URL
    """
    with app.app_context():
        now_time = datetime.now()
        shifu_draft = get_latest_shifu_draft(shifu_id)
        if not shifu_draft:
            raise_error("server.shifu.shifuNotFound")
        PublishedShifu.query.filter_by(shifu_bid=shifu_id).update({"deleted": 1})
        PublishedOutlineItem.query.filter_by(shifu_bid=shifu_id).update({"deleted": 1})
        shifu_published = PublishedShifu()
        shifu_published.shifu_bid = shifu_id
        shifu_published.title = shifu_draft.title
        shifu_published.description = shifu_draft.description
        shifu_published.avatar_res_bid = shifu_draft.avatar_res_bid
        shifu_published.keywords = shifu_draft.keywords
        shifu_published.llm = shifu_draft.llm
        shifu_published.llm_temperature = shifu_draft.llm_temperature
        shifu_published.price = shifu_draft.price
        shifu_published.created_user_bid = shifu_draft.created_user_bid
        shifu_published.updated_user_bid = user_id
        shifu_published.updated_at = now_time
        shifu_published.llm_system_prompt = shifu_draft.llm_system_prompt
        shifu_published.ask_enabled_status = shifu_draft.ask_enabled_status
        shifu_published.ask_llm = shifu_draft.ask_llm
        shifu_published.ask_llm_temperature = shifu_draft.ask_llm_temperature
        shifu_published.ask_llm_system_prompt = shifu_draft.ask_llm_system_prompt
        shifu_published.ask_provider_config = (
            getattr(shifu_draft, "ask_provider_config", "{}") or "{}"
        )
        # TTS Configuration
        shifu_published.tts_enabled = shifu_draft.tts_enabled
        shifu_published.tts_provider = getattr(shifu_draft, "tts_provider", "") or ""
        shifu_published.tts_model = getattr(shifu_draft, "tts_model", "") or ""
        shifu_published.tts_voice_id = shifu_draft.tts_voice_id
        shifu_published.tts_speed = shifu_draft.tts_speed
        shifu_published.tts_pitch = shifu_draft.tts_pitch
        shifu_published.tts_emotion = shifu_draft.tts_emotion
        # Learner language setting
        shifu_published.use_learner_language = getattr(
            shifu_draft, "use_learner_language", 0
        )
        db.session.add(shifu_published)
        db.session.flush()
        outline_tree = build_outline_tree(app, shifu_id)

        def publish_outline_item(node: ShifuOutlineTreeNode, history_item: HistoryItem):
            outline_item = PublishedOutlineItem()
            draft_outline_item: DraftOutlineItem = node.outline
            outline_item.shifu_bid = shifu_id
            outline_item.outline_item_bid = draft_outline_item.outline_item_bid
            outline_item.title = draft_outline_item.title
            outline_item.position = draft_outline_item.position
            outline_item.type = draft_outline_item.type
            outline_item.hidden = draft_outline_item.hidden
            outline_item.parent_bid = draft_outline_item.parent_bid
            outline_item.llm = draft_outline_item.llm
            outline_item.llm_temperature = draft_outline_item.llm_temperature
            outline_item.llm_system_prompt = draft_outline_item.llm_system_prompt
            outline_item.ask_enabled_status = draft_outline_item.ask_enabled_status
            outline_item.ask_llm = draft_outline_item.ask_llm
            outline_item.ask_llm_temperature = draft_outline_item.ask_llm_temperature
            outline_item.ask_llm_system_prompt = (
                draft_outline_item.ask_llm_system_prompt
            )
            outline_item.created_user_bid = user_id
            outline_item.created_at = now_time
            outline_item.updated_user_bid = user_id
            outline_item.updated_at = now_time
            outline_item.prerequisite_item_bids = (
                draft_outline_item.prerequisite_item_bids
            )
            outline_item.content = draft_outline_item.content
            db.session.add(outline_item)
            db.session.flush()
            markdown_flow = MarkdownFlow(
                draft_outline_item.content
            ).set_output_language(get_markdownflow_output_language())
            blocks = markdown_flow.get_all_blocks()
            outline_item_history_item = HistoryItem(
                bid=node.outline_id,
                id=outline_item.id,
                type="outline",
                children=[],
                child_count=len(blocks),
            )
            history_item.children.append(outline_item_history_item)
            if node.children and len(node.children) > 0:
                for child in node.children:
                    publish_outline_item(child, outline_item_history_item)

        history_item = HistoryItem(
            bid=shifu_id, id=shifu_published.id, type="shifu", children=[]
        )
        for node in outline_tree:
            publish_outline_item(node, history_item)

        shifu_log_published_struct = LogPublishedStruct()
        shifu_log_published_struct.struct_bid = generate_id(app)
        shifu_log_published_struct.shifu_bid = shifu_id
        shifu_log_published_struct.struct = history_item.to_json()
        shifu_log_published_struct.created_user_bid = user_id
        shifu_log_published_struct.created_at = now_time
        db.session.add(shifu_log_published_struct)
        db.session.commit()
        parent_shifu_context = get_shifu_context_snapshot()
        if sync_summary:
            _run_summary_with_error_handling(app, shifu_id, parent_shifu_context)
        else:
            thread = threading.Thread(
                target=_run_summary_with_error_handling,
                args=(app, shifu_id, parent_shifu_context),
            )
            thread.daemon = True  # Ensure thread doesn't prevent app shutdown
            thread.start()
        return _build_frontend_url(base_url, f"/c/{shifu_id}")


def _run_summary_with_error_handling(app, shifu_id, shifu_context_snapshot=None):
    """
    Run shifu summary generation with error handling
    Args:
        app: Flask application instance
        shifu_id: Shifu ID
    """
    try:
        apply_shifu_context_snapshot(shifu_context_snapshot)
        get_shifu_summary(app, shifu_id)
    except Exception as e:
        app.logger.error(f"Failed to generate shifu summary for {shifu_id}: {str(e)}")


def get_shifu_summary(app, shifu_id: str):
    """
    Obtain the shifu summary information
    Args:
        app: Flask application instance
        shifu_id: Shifu ID
    """
    with app.app_context():
        shifu: PublishedShifu = (
            PublishedShifu.query.filter(PublishedShifu.shifu_bid == shifu_id)
            .order_by(PublishedShifu.id.desc())
            .first()
        )
        if not shifu:
            app.logger.error(f"get_shifu_summary shifu_id: {shifu_id} not found")
            return

        # Get the prompt word template
        summary_prompt_template = load_prompt_template("summary")
        ask_prompt_template = load_prompt_template("ask")

        # Get course data
        outline_tree, outline_ids, outline_item_map = _get_shifu_data(app, shifu_id)

        # Generate summaries
        outline_summary_map = _generate_summaries(
            app, outline_tree, outline_item_map, summary_prompt_template, shifu
        )

        # Generate ask_prompt
        _generate_ask_prompts(
            app,
            outline_tree,
            outline_ids,
            outline_summary_map,
            outline_item_map,
            ask_prompt_template,
        )
        shifu.ask_enabled_status = ASK_MODE_ENABLE
        db.session.commit()
        return


def _generate_ask_prompts(
    app,
    shifu_info: ShifuInfoDto,
    outline_ids: list[str],
    outline_summary_map: dict[str, dict],
    outline_item_map: dict[str, PublishedOutlineItem],
    ask_prompt_template: str,
):
    """
    Generate ask_prompt for each section
    Args:
        app: Flask application instance
        shifu_info: Shifu info
        outline_ids: Section ID list
        outline_summary_map: Summary mapping
        outline_item_map: Outline item mapping
        ask_prompt_template: Ask template
    Returns:
        None
    """
    for chapter in shifu_info.outline_items:
        for section in chapter.children:
            # Split outline_summary_map into learned and unlearned parts based on current section ID
            current_section_id = section.bid
            # Find the index of current section in outline_ids
            current_index = outline_ids.index(current_section_id)
            # Split content into learned and unlearned parts
            learned_summaries = []
            unlearned_summaries = []
            for i, section_id in enumerate(outline_ids):
                if section_id in outline_summary_map:
                    if i <= current_index:
                        # Current section and all previous sections (learned)
                        learned_summaries.append(outline_summary_map[section_id])
                    else:
                        # All sections after current section (unlearned)
                        unlearned_summaries.append(outline_summary_map[section_id])

            # Build text for learned content
            learned_text = _build_summary_text(learned_summaries, is_learned=True)

            # Build text for unlearned content
            unlearned_text = _build_summary_text(unlearned_summaries, is_learned=False)

            ask_prompt = _make_ask_prompt(
                app, ask_prompt_template, learned_text, unlearned_text
            )
            outline_item = outline_item_map.get(section.bid)
            if outline_item:
                outline_item.ask_llm_system_prompt = ask_prompt


def _generate_summaries(
    app,
    outline_tree: ShifuInfoDto,
    outline_item_map: dict[str, PublishedOutlineItem],
    summary_prompt_template,
    shifu: PublishedShifu,
) -> dict[str, dict]:
    """
    Generate summaries for all sections
    Args:
        app: Flask application instance
        outline_tree: Outline tree
        outline_item_map: Outline item mapping
        summary_prompt_template: Summary template
        shifu: Course information
    Returns:
        Summary mapping
    """
    outline_summary_map = {}

    # Get model configuration
    model_name = shifu.ask_llm or shifu.llm
    temperature = shifu.ask_llm_temperature or shifu.llm_temperature or 0.3
    if not model_name:
        model_name = app.config.get("DEFAULT_LLM_MODEL", "")

    for chapter in outline_tree.outline_items:
        for section in chapter.children:
            outline_item = outline_item_map.get(section.bid)
            now_lesson_script_prompts = ""
            if outline_item and bool(outline_item.content):
                app.logger.info(
                    f"outline_item: {outline_item.outline_item_bid} has mdflow content,make summary from mdflow"
                )
                mdflow = MarkdownFlow(outline_item.content).set_output_language(
                    get_markdownflow_output_language()
                )
                blocks = mdflow.get_all_blocks()
                for block in blocks:
                    if block.block_type == BlockType.CONTENT:
                        now_lesson_script_prompts += "\n" + block.content
            else:
                now_lesson_script_prompts = outline_item.content

            final_prompt = summary_prompt_template.format(
                all_script_content=now_lesson_script_prompts
            )

            summary = _get_summary(
                app,
                prompt=final_prompt,
                model_name=model_name,
                temperature=temperature,
            )

            # Update section information
            outline_item = outline_item_map.get(section.bid)
            if outline_item:
                outline_item.summary = summary
                outline_item.ask_enabled_status = ASK_MODE_ENABLE

                # Store summary information
                outline_summary_map[section.bid] = {
                    "chapter_id": chapter.bid,
                    "chapter_name": chapter.title,
                    "section_id": section.bid,
                    "section_name": section.title,
                    "content": summary,
                }

    return outline_summary_map


def _get_shifu_data(
    app, shifu_id: str
) -> tuple[
    ShifuInfoDto,
    list[str],
    dict[str, PublishedOutlineItem],
]:
    """
    Get shifu related data
    Args:
        app: Flask application instance
        shifu_id: shifu ID
    Returns:
        (outline_tree, outline_ids, outline_item_map)
    """

    outline_ids = []

    shifu_outline_tree = get_shifu_outline_tree(app, shifu_id, is_preview=False)

    q = queue.Queue()
    for item in shifu_outline_tree.outline_items:
        q.put(item)
    while not q.empty():
        item = q.get()
        outline_ids.append(item.bid)
        if item.children:
            for child in item.children:
                q.put(child)

    # Get all section data
    outline_items = (
        PublishedOutlineItem.query.filter(
            PublishedOutlineItem.outline_item_bid.in_(outline_ids),
            PublishedOutlineItem.deleted == 0,
        )
        .order_by(PublishedOutlineItem.id.desc())
        .all()
    )
    outline_item_map = {
        outline_item.outline_item_bid: outline_item for outline_item in outline_items
    }

    return shifu_outline_tree, outline_ids, outline_item_map


def _make_ask_prompt(
    app, ask_prompt: str, learned_text: str, unlearned_text: str
) -> str:
    """
    Make ask prompt
    Args:
        app: Flask application instance
        ask_prompt: Ask prompt
        learned_text: Learned text
        unlearned_text: Unlearned text
    Returns:
        Ask prompt
    """
    result = ask_prompt.format(
        learned=("\n" + learned_text) if learned_text else "",
        unlearned=("\n" + unlearned_text) if unlearned_text else "",
        shifu_system_message="{shifu_system_message}",
    )
    return result


def _get_summary(app, prompt, model_name, user_id=None, temperature=0.8):
    """
    Call the AI model to generate summary
    Args:
        app: Flask application instance
        prompt: Prompt to be summarized
        model_name: Model name to use
        user_id: Optional, user ID
        temperature: Optional, sampling temperature
    Returns:
        Summary text
    """
    # Create langfuse trace/span
    trace, span = create_trace_with_root_span(
        client=get_langfuse_client(),
        trace_payload={
            "user_id": user_id or "shifu-summary",
            "input": prompt,
            "name": "shifu_summary",
        },
        root_span_payload={
            "name": "shifu_summary",
            "input": prompt,
        },
    )
    summary = ""
    try:
        response = invoke_llm(
            app,
            user_id or "shifu-summary",
            span,
            model_name,
            prompt,
            temperature=temperature,
            generation_name="shifu_summary",
            usage_context=UsageContext(
                user_bid=user_id or "shifu-summary",
                shifu_bid="",
                usage_scene=BILL_USAGE_SCENE_DEBUG,
                billable=0,
            ),
            usage_scene=BILL_USAGE_SCENE_DEBUG,
            billable=0,
        )
        for chunk in response:
            summary += getattr(chunk, "result", "")
        return summary
    finally:
        finalize_langfuse_trace(
            trace=trace,
            root_span=span,
            trace_payload={"output": summary},
            root_span_payload={"output": summary},
        )


def _build_summary_text(summaries: list[dict], is_learned: bool) -> str:
    """
    Build a summary text based on whether it's learned or unlearned
    Args:
        summaries: List of summary dictionaries
        is_learned: Boolean indicating whether the summary is for learned or unlearned
    Returns:
        Built summary text
    """
    if not summaries:
        return ""

    result_lines = []
    chapter_titles_added = set()

    for summary in summaries:
        chapter_id = summary["chapter_id"]
        chapter_name = summary["chapter_name"]
        section_name = summary["section_name"]
        content = summary["content"]

        # Check if chapter title needs to be added
        if chapter_id not in chapter_titles_added:
            # First time encountering this chapter, add chapter title
            result_lines.append(f"### {chapter_name}")
            chapter_titles_added.add(chapter_id)

        # Add section title and content
        result_lines.append(f"#### {section_name}")
        result_lines.append(content)
        result_lines.append("")  # Add empty line separator

    return "\n".join(result_lines)
