"""Report Sections Extraction"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/03_extract.ipynb.

# %% auto 0
__all__ = ['CoreSectionsOutput', 'get_text', 'rm_nested', 'identify_core_sections', 'extract_sections']

# %% ../nbs/03_extract.ipynb 3
from fastcore.all import *
from operator import getitem
from pydantic import BaseModel
from lisette.core import completion, mk_msg
from toolslm.md_hier import create_heading_dict
import json
from .core import load_prompt

# %% ../nbs/03_extract.ipynb 4
class CoreSectionsOutput(BaseModel):
    "Identify the core sections of the report"
    section_paths: list[list[str]]
    reasoning: str

# %% ../nbs/03_extract.ipynb 9
def get_text(ks:list[str], # List of exact key strings forming path through nested dict
             hdgs:dict # Nested dictionary of headings created by `create_heading_dict`
            ) -> str: # Extracted markdown text for the section
    "Navigate through nested heading levels and return the text content"
    return L(ks).reduce(getitem, hdgs).text

# %% ../nbs/03_extract.ipynb 12
def rm_nested(paths:list[list[str]] # List of section paths, where each path is a list of keys
             ) -> list[list[str]]: # Filtered list with nested paths removed
    "Remove paths that are children of other paths in the list"
    paths = sorted(paths, key=len)
    keep = []
    for p in paths:
        if not any(p[:len(k)] == k for k in keep): keep.append(p)
    return keep

# %% ../nbs/03_extract.ipynb 15
def identify_core_sections(
    hdgs:dict, # Nested dictionary of report headings from `create_heading_dict`
    sp:str=None, # System prompt for section identification
    response_format:type[BaseModel]=CoreSectionsOutput, # Pydantic model for structured output
    model:str='claude-sonnet-4-5' # LLM model to use for identification
) -> dict: # Dictionary with 'section_paths' and 'reasoning' keys
    "Use LLM to identify core sections (exec summary, intro, conclusions, recommendations) from ToC"
    if sp is None: sp = load_prompt('select_sections', 'files/prompts')
    res = completion(model=model, messages=[mk_msg(f"Here is the table of contents as a nested dictionary:\n\n{hdgs}")], 
                     system=[{"type": "text", "text": sp}], response_format=response_format)
    return json.loads(res.choices[0].message.content)

# %% ../nbs/03_extract.ipynb 18
@delegates(identify_core_sections)
def extract_sections(
    md:str, # Markdown text of full report
    **kwargs # Additional keyword arguments passed to `identify_core_sections`
) -> str: # Concatenated text of all core sections
    "Extract and concatenate core sections (exec summary, intro, conclusions, recommendations) from report markdown"
    hdgs = create_heading_dict(md)
    sections = identify_core_sections(hdgs, **kwargs)
    paths = rm_nested(sections['section_paths'])
    return '\n'.join([get_text(p, hdgs) for p in paths])
