"""End-to-end pipeline for processing IOM evaluation reports"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/06_pipeline.ipynb.

# %% auto 0
__all__ = ['Report', 'load_report', 'run_pipeline']

# %% ../nbs/06_pipeline.ipynb 3
from fastcore.all import *
from pathlib import Path
from .core import n_tokens, load_prompt
from .readers import load_evals, find_eval, Evaluation
from .downloaders import download_eval
from .extract import extract_sections
from .themes import load_enablers, load_ccp, load_gcms, load_srf_outs, load_gcm_lut, fmt_enablers_ccp, fmt_srf_outs, get_srf_outs
from .mapper import mk_system_blocks, map_themes, sort_by_centrality, get_top_ids, parse_json_response
from mistocr.core import read_pgs
from mistocr.pipeline import pdf_to_md
import json

# %% ../nbs/06_pipeline.ipynb 5
class Report:
    "An evaluation report with full pipeline support"
    def __init__(self,
                 ev:Evaluation,                   # The evaluation metadata object
                 pdf_url:str=None,                # Optional direct URL to PDF
                 results_path:str='data/results'  # Path to save/load results
                ):
        store_attr('ev,pdf_url,results_path')
        self.id = ev.id
        self.pdf_path,self.md_path,self.sections,self.mappings = None,None,None,{}
    
    @classmethod
    def from_url(cls,
                 url:str,                         # URL of the evaluation PDF
                 evals:list,                      # List of `Evaluation` objects to search
                 results_path:str='data/results'  # Path to save/load results
                ): return cls(find_eval(evals, url, by='url'), pdf_url=url, results_path=results_path)
    
    @classmethod
    def from_title(cls,
                   title:str,                      # Title to search for
                   evals:list,                     # List of `Evaluation` objects to search
                   results_path:str='data/results' # Path to save/load results
                  ): return cls(find_eval(evals, title, by='title'), results_path=results_path)

# %% ../nbs/06_pipeline.ipynb 6
@patch
def _repr_markdown_(self:Report):
    "Display report metadata and processing status in Jupyter notebooks"
    title = self.ev.meta.get('Title', 'Untitled')
    year = self.ev.meta.get('Year', 'n/a')
    org = self.ev.meta.get('Evaluation Commissioner', 'Unknown')
    
    status = []
    if self.pdf_path: status.append(f'✓ PDF downloaded')
    if self.md_path: status.append(f'✓ Markdown converted')
    if self.sections: status.append(f'✓ Sections extracted (~{n_tokens(self.sections)} tokens)')
    if self.mappings:
        mapped = ', '.join(self.mappings.keys())
        status.append(f'✓ Mappings: {mapped}')
    status_str = ' | '.join(status) if status else 'Not processed'
    
    return f"""
## Report: {title}
**Year:** {year} | **Organization:** {org}  
**ID:** `{self.id}`

**Processing Status:**  
{status_str}

**Documents:** {len(self.ev.docs)} available
"""

# %% ../nbs/06_pipeline.ipynb 13
@patch
def save(self:Report,
         path:str=None  # Override default results path
        ) -> Report:
    "Save report state to JSON"
    p = Path(path or self.results_path)/f'{self.id}.json'
    p.parent.mkdir(parents=True, exist_ok=True)
    data = dict(id=self.id, ev_meta=self.ev.meta, ev_docs=self.ev.docs, sections=self.sections, mappings=self.mappings,
                pdf_path=str(self.pdf_path) if self.pdf_path else None, md_path=str(self.md_path) if self.md_path else None,
                results_path=self.results_path)
    p.write_text(json.dumps(data, indent=2))
    return self

# %% ../nbs/06_pipeline.ipynb 14
def load_report(id:str,                  # Report ID (hash)
                path:str='data/results'  # Results directory
               ) -> Report:
    "Load a saved Report by id"
    data = json.loads((Path(path)/f'{id}.json').read_text())
    ev = Evaluation(id=data['id'], meta=data['ev_meta'], docs=data['ev_docs'])
    report = Report(ev, results_path=data.get('results_path', path))
    report.sections,report.mappings = data['sections'],data['mappings']
    if data.get('pdf_path'): report.pdf_path = Path(data['pdf_path'])
    if data.get('md_path'): report.md_path = Path(data['md_path'])
    return report

# %% ../nbs/06_pipeline.ipynb 19
@patch
def download(self:Report,
             dst:str='data/pdfs'  # Destination directory for PDFs
            ) -> Report:
    "Download evaluation PDF to `dst`/`eval_id`/"
    self.pdf_path = download_eval(self.ev, dst=dst)
    self.save(self.results_path)
    return self

# %% ../nbs/06_pipeline.ipynb 23
@patch
async def ocr(self:Report,
              dst:str='data/md',       # Destination directory for markdown files
              add_img_desc:bool=True,  # Whether to add image descriptions
              **kwargs                 # Additional args passed to pdf_to_md
             ) -> Report:
    "Run OCR on PDF and fix heading hierarchy"
    if self.pdf_path is None: raise ValueError("Call download() first")
    if self.pdf_url: pdf_file = self.pdf_path/Path(self.pdf_url).name
    else: pdf_file = first(self.pdf_path.glob('*.pdf'))
    await pdf_to_md(pdf_file, Path(dst)/self.id, add_img_desc=add_img_desc, **kwargs)
    self.md_path = Path(dst)/self.id
    self.save(self.results_path)
    return self

# %% ../nbs/06_pipeline.ipynb 28
@patch
def extract(self:Report, **kwargs):
    "Extract core sections from markdown"
    if self.md_path is None: raise ValueError("Call ocr() first")
    md = read_pgs(self.md_path)
    self.sections = extract_sections(md, **kwargs)
    self.save(self.results_path)
    return self

# %% ../nbs/06_pipeline.ipynb 32
@patch
def _ensure_sys_blocks(self:Report):
    if self.sections is None: raise ValueError("Call extract() first")
    if not hasattr(self, '_sys_blocks'): self._sys_blocks = mk_system_blocks(self.sections)

# %% ../nbs/06_pipeline.ipynb 33
def _map_single(sys_blocks,                # System blocks from mk_system_blocks
                theme_type,                 # One of: 'enablers', 'ccps', 'gcm', 'outputs'
                path='files/themes',        # Path to theme files
                model='claude-haiku-4-5',   # Model to use for mapping
                gcm_ids=None                # GCM IDs for output mapping
               ):
    "Map system blocks (Report) to a single theme type using appropriate prompts and formatting"
    if theme_type == 'enablers': res = map_themes(sys_blocks, fmt_enablers_ccp(load_enablers(path)), load_prompt('srf_enablers'), model)
    elif theme_type == 'ccps': res = map_themes(sys_blocks, fmt_enablers_ccp(load_ccp(path)), load_prompt('srf_ccps'), model)
    elif theme_type == 'gcm': res = map_themes(sys_blocks, load_gcms(path), load_prompt('gcm'), model)
    elif theme_type == 'outputs':
        srf_obj, gcm_lut = load_srf_outs(path), load_gcm_lut(path)
        output_ids = get_srf_outs(gcm_lut, gcm_ids)
        res = map_themes(sys_blocks, fmt_srf_outs(srf_obj, output_ids), load_prompt('srf_outputs'), model)
    return parse_json_response(res)

# %% ../nbs/06_pipeline.ipynb 35
@patch
def map_enablers(self:Report, **kwargs):
    self._ensure_sys_blocks()
    self.mappings['enablers'] = _map_single(self._sys_blocks, 'enablers', **kwargs)
    self.save(self.results_path)
    return self

# %% ../nbs/06_pipeline.ipynb 40
@patch
def map_ccps(self:Report, **kwargs):
    self._ensure_sys_blocks()
    self.mappings['ccps'] = _map_single(self._sys_blocks, 'ccps', **kwargs)
    self.save(self.results_path)
    return self

# %% ../nbs/06_pipeline.ipynb 44
@patch
def map_gcm(self:Report, **kwargs):
    self._ensure_sys_blocks()
    self.mappings['gcm'] = _map_single(self._sys_blocks, 'gcm', **kwargs)
    self.save(self.results_path)
    return self

# %% ../nbs/06_pipeline.ipynb 48
@patch
def map_outputs(self:Report, gcm_ids=None, **kwargs):
    self._ensure_sys_blocks()
    if gcm_ids is None: gcm_ids = [get_top_ids(self.mappings['gcm'])[0]] if self.mappings['gcm'] else []
    self.mappings['outputs'] = _map_single(self._sys_blocks, 'outputs', gcm_ids=gcm_ids, **kwargs)
    self.save(self.results_path)
    return self

# %% ../nbs/06_pipeline.ipynb 52
@patch
def map_all(self:Report, **kwargs): return self.map_enablers(**kwargs).map_ccps(**kwargs).map_gcm(**kwargs).map_outputs(**kwargs)

# %% ../nbs/06_pipeline.ipynb 54
async def run_pipeline(url:str,                       # URL of the evaluation PDF
                       evals:list,                    # List of `Evaluation` objects to search
                       pdf_dst:str='data/pdfs',       # Destination directory for PDFs
                       md_dst:str='data/md',          # Destination directory for markdown files
                       results_path:str='data/results', # Path to save/load results
                       ocr_kwargs:dict=None,          # Additional arguments passed to ocr (e.g. add_img_desc, model)
                       **kwargs                       # Additional arguments passed to mapping functions
                      ) -> Report:                    # Fully processed report with all mappings
    "Run complete pipeline: download → ocr → extract → map_themes"
    print(f"Creating report from URL...")
    report = Report.from_url(url, evals, results_path=results_path)
    print(f"Step 1/7: Downloading PDF...")
    report.download(dst=pdf_dst)
    print(f"Step 2/7: Running OCR...")
    await report.ocr(dst=md_dst, **(ocr_kwargs or {}))
    print(f"Step 3/7: Extracting sections...")
    report.extract()
    print(f"Step 4/7: Mapping enablers...")
    report.map_enablers(**kwargs)
    print(f"Step 5/7: Mapping CCPs...")
    report.map_ccps(**kwargs)
    print(f"Step 6/7: Mapping GCM objectives...")
    report.map_gcm(**kwargs)
    print(f"Step 7/7: Mapping outputs...")
    report.map_outputs(**kwargs)
    print(f"Pipeline complete!")
    return report
