"""Report Sections Extraction"""

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

# %% auto #0
__all__ = ['CoreSectionsOutput', 'heading_to_key', 'find_path', 'rm_nested', 'headings_to_paths', 'get_text',
           'identify_core_sections', 'get_h1_title', 'extract_sections']

# %% ../nbs/03_extract.ipynb #bb2a3b4d
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
import logging
from .core import load_prompt

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

# %% ../nbs/03_extract.ipynb #42dcbe92
def heading_to_key(hdg:str # Raw markdown heading like "## 1. Intro ..."
                  ) -> str: # Heading text without markdown prefix
    "Strip markdown heading prefix: '## 1. Intro ...' → '1. Intro ...'"
    return hdg.lstrip('#').strip()

# %% ../nbs/03_extract.ipynb #dd600bdc
def find_path(hdgs:dict, # Nested dictionary of headings
              key:str    # Heading key to find
             ) -> list[str]: # Full path to key, or empty list if not found
    "Find full path to a heading key in nested dict"
    for k in hdgs.keys():
        if k == key: return [k]
        res = find_path(hdgs[k], key)
        if res: return [k] + res
    return []

# %% ../nbs/03_extract.ipynb #7b6ffa03
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 #84e7a79e
def headings_to_paths(hdgs:dict,                  # Nested dictionary of headings
                      selected_headings:list[str] # Raw markdown headings like ["## 1. Intro ...", ...]
                     ) -> list[list[str]]:        # Deduplicated paths through heading hierarchy
    "Convert raw headings to deduplicated paths"
    paths = [find_path(hdgs, heading_to_key(h)) for h in selected_headings]
    paths = [p for p in paths if p]  # remove empty (not found)
    return rm_nested(paths)

# %% ../nbs/03_extract.ipynb #518d5333
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 #e7c1a140
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')
    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 #438e0148
def get_h1_title(hdgs:dict) -> str:
    "Get the main title (first h1 heading) from heading dict"
    return f"# {first(hdgs.keys())}" if hdgs else ""

# %% ../nbs/03_extract.ipynb #cb6dd783
@delegates(identify_core_sections)
def extract_sections(
    md:str,                            # Markdown text of full report
    selected_headings:list[str]=None,  # Raw headings like ["## 1. Intro ...", ...]; if None, uses LLM
    **kwargs                           # Additional kwargs passed to `identify_core_sections`
) -> str:                              # Concatenated text of all core sections
    "Extract core sections from report. Uses LLM auto-detection if `selected_headings` is None, otherwise uses provided headings."
    hdgs = create_heading_dict(md)
    
    if selected_headings is not None: paths = headings_to_paths(hdgs, selected_headings)
    else:
        sections = identify_core_sections(hdgs, **kwargs)
        paths = rm_nested(sections['section_paths'])

    if not paths: return ""
    texts = [get_h1_title(hdgs)]
    for p in paths:
        try: texts.append(get_text(p, hdgs))
        except (KeyError, AttributeError) as e: logging.warning(f"Path not found, skipping: {p}")
    return '\n'.join(texts)
