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

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

# %% auto #0
__all__ = ['logger', 'Report', 'load_report', 'map_single', 'should_force', 'PipelineResult', 'get_current_step', 'run_pipeline',
           'BatchResult', 'filter_evals', 'is_completed', 'batch_run']

# %% ../nbs/07_pipeline.ipynb #44a0b07e
from fastcore.all import *
from pathlib import Path
from .core import n_tokens, load_prompt
from .readers import load_evals, find_eval, Evaluation, eval_url
from .downloaders import download_eval
from .extract import extract_sections
from .themes import load_enbs, load_ccps, load_gcms, load_srf_outs, load_gcms_lut, fmt_enbs, fmt_ccps, fmt_srf_outs, get_srf_outs
from .mapper import mk_system_blocks, map_themes, sort_by_relevance, get_top_ids, parse_res
from mistocr.core import read_pgs
from mistocr.pipeline import pdf_to_md
from datetime import datetime
import json
import shutil
import logging
import traceback

# %% ../nbs/07_pipeline.ipynb #047209cb
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.WARNING, format='%(name)s - %(levelname)s - %(message)s')
logger.setLevel(logging.DEBUG)

# %% ../nbs/07_pipeline.ipynb #73009447
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()
        self.id = ev.id
        self.pdf_path,self.md_path = None,None
        self.mappings,self.selected_headings = {},[]
        self.curation_status = 'pending'
        self._load_existing()

# %% ../nbs/07_pipeline.ipynb #73009447
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()
        self.id = ev.id
        self.pdf_path,self.md_path = None,None
        self.mappings,self.selected_headings = {},[]
        self.curation_status = 'pending'
        self._load_existing()

# %% ../nbs/07_pipeline.ipynb #6e1a2613
@patch
def _load_existing(self:Report):
    "Load state from saved JSON if it exists"
    p = Path(self.results_path)/f'{self.id}.json'
    if not p.exists(): return
    data = json.loads(p.read_text())
    self.mappings = data.get('mappings', {})
    self.curation_status = data.get('curation_status', 'pending')
    self.selected_headings = data.get('selected_headings', [])


# %% ../nbs/07_pipeline.ipynb #f27d90ae
@patch(cls_method=True)
def from_url(
    cls:Report,
    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
    ):                               # Report initialized from URL
    "Create a Report by finding an evaluation matching the given URL"
    return cls(find_eval(evals, url, by='url'), pdf_url=url, results_path=results_path)

# %% ../nbs/07_pipeline.ipynb #6bccec97
@patch(cls_method=True)
def from_title(
    cls:Report,
    title:str,                      # Title to search for
    evals:list,                     # List of `Evaluation` objects to search
    results_path:str='data/results' # Path to save/load results
    ):                               # Report initialized from title
    "Create a Report by finding an evaluation matching the given title"
    ev = find_eval(evals, title, by='title')
    url = eval_url(ev)
    if not url: raise ValueError(f"{ev.id}: No report URL found")
    return cls(ev, pdf_url=url, results_path=results_path)

# %% ../nbs/07_pipeline.ipynb #8a6108de
@patch(cls_method=True)
def from_id(
    cls:Report,
    id:str,                         # id to search for
    evals:list,                     # List of `Evaluation` objects to search
    results_path:str='data/results' # Path to save/load results
    ):                              # Report initialized from id
    "Create a Report by finding an evaluation matching the given id"
    ev = find_eval(evals, id, by='id')
    url = eval_url(ev)
    if not url: raise ValueError(f"{ev.id}: No report URL found")
    return cls(ev, pdf_url=url, results_path=results_path)

# %% ../nbs/07_pipeline.ipynb #ed425a53
@patch
def _repr_markdown_(self:Report) -> str:
    "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')
    
    pipeline = []
    if self.md_path: pipeline.append('✓ OCR')
    if self.curation_status == 'sections_selected': pipeline.append(f'✓ Curated ({len(self.selected_headings)} headings)')
    if self.mappings:
        mapped = ', '.join(self.mappings.keys())
        pipeline.append(f'✓ Mappings ({mapped})')
    pipeline_str = ' | '.join(pipeline) if pipeline else 'Not processed'
    
    return f"""
## Report: {title}
**Year:** {year} | **Organization:** {org}  
**ID:** `{self.id}`

**Pipeline:** {pipeline_str}  
**Curation:** {self.curation_status}

**Report:** [View Evaluation Report]({self.pdf_url})
"""

# %% ../nbs/07_pipeline.ipynb #6d634bf6
@patch
def save(self:Report,
         path:str=None  # Override default results path
        ) -> Report:    # Reports self for method chaining
    "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,
        report_url=self.pdf_url,
        meta={**self.ev.meta, 'report_url': self.pdf_url},
        docs=self.ev.docs,
        curation_status=self.curation_status,
        selected_headings=self.selected_headings,
        mappings=self.mappings,
        timestamp=datetime.now().isoformat()
    )
    p.write_text(json.dumps(data, indent=2))
    return self

# %% ../nbs/07_pipeline.ipynb #3a58454f
def load_report(id:str,                  # Report ID (hash)
                base_path:str='data'     # Base directory containing pdf/, md/, results/
               ) -> Report:              # The loaded Report
    "Load a saved Report by id"
    results_path = Path(base_path)/'results'
    data = json.loads((results_path/f'{id}.json').read_text())
    
    ev = Evaluation(id=data['id'], meta=data['meta'], docs=data['docs'])
    report = Report(ev, pdf_url=data.get('report_url'), results_path=results_path)
    
    # Infer paths from convention
    pdf_dir = Path(base_path)/'pdf'/id
    md_dir = Path(base_path)/'md'/id
    report.pdf_path = pdf_dir if pdf_dir.exists() else None
    report.md_path = md_dir if md_dir.exists() else None
    
    # Restore new attributes
    report.curation_status = data.get('curation_status', 'pending')
    report.selected_headings = data.get('selected_headings', [])
    report.mappings = data.get('mappings', {})
    
    return report

# %% ../nbs/07_pipeline.ipynb #3cfe4b55
@patch
def download(self:Report,
             dst:str='data/pdf',  # Destination directory for PDFs
             force:bool=False      # Force re-download
            ) -> Report:           # Self for chaining
    "Download evaluation PDF to `dst`/`eval_id`/"
    if self.pdf_path and not force: return self
    self.pdf_path = download_eval(self.ev, dst=dst)
    return self

# %% ../nbs/07_pipeline.ipynb #9fc305a0
@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
              force:bool=False,        # Force re-OCR
              progress:bool=False,     # Show OCR progress messages
              fix_kwargs:dict=None,    # Extra kwargs for fix_hdgs
              desc_kwargs:dict=None    # Extra kwargs for add_img_descs
             ) -> Report:              # Self for chaining
    "Run OCR on PDF and fix heading hierarchy"
    if (Path(dst)/self.id).exists() and not force: return self
    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, progress=progress, fix_kwargs=fix_kwargs, desc_kwargs=desc_kwargs)
    self.md_path = Path(dst)/self.id
    self.save(self.results_path)
    return self

# %% ../nbs/07_pipeline.ipynb #f90d9203
@patch
def get_sections(self:Report) -> str:
    "Extract sections on demand from selected headings"
    if self.md_path is None: raise ValueError("Call ocr() first")
    if self.curation_status != 'sections_selected': raise ValueError("Curation required: use curator app to select headings first")
    if not self.selected_headings: raise ValueError("No headings selected")
    return extract_sections(read_pgs(self.md_path), selected_headings=self.selected_headings)

# %% ../nbs/07_pipeline.ipynb #f7b5bf6c
@patch
def ensure_sys_blocks(self:Report) -> None:
    "Ensure system blocks are available"
    if not hasattr(self, '_sys_blocks'): self._sys_blocks = mk_system_blocks(self.get_sections())

# %% ../nbs/07_pipeline.ipynb #5f51ae13
@delegates(map_themes)
def map_single(sys_blocks,                 # System blocks from mk_system_blocks
                theme_type,                # One of: 'enbs', 'ccps', 'gcms', 'outs'
                path=None,                 # Path to theme files
                model='claude-haiku-4-5',  # Model to use for mapping
                gcm_ids=None,              # GCM IDs for output mapping
                **kwargs                   # Additional args passed to map_themes
               ) -> dict:                  # Mapping results
    "Map system blocks (Report) to a single theme type using appropriate prompts and formatting"
    if theme_type == 'enbs': res = map_themes(sys_blocks, fmt_enbs(load_enbs(path)), load_prompt('srf_enablers'), model, **kwargs)
    elif theme_type == 'ccps': res = map_themes(sys_blocks, fmt_ccps(load_ccps(path)), load_prompt('srf_ccps'), model, **kwargs)
    elif theme_type == 'gcms': res = map_themes(sys_blocks, load_gcms(path), load_prompt('gcms'), model, **kwargs)
    elif theme_type == 'outs':
        srf_obj, gcm_lut = load_srf_outs(path), load_gcms_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, **kwargs)
    return parse_res(res)

# %% ../nbs/07_pipeline.ipynb #4c1885d3
@patch
def map_enbs(self:Report,
             force:bool=False,  # Re-run even if already completed
             **kwargs           # Additional args passed to map_single (e.g. path, model)
            ) -> Report:        # Self for chaining
    "Map report sections to Strategic Results Framework enablers"
    if 'enbs' in self.mappings and not force: return self
    self.ensure_sys_blocks()
    self.mappings['enbs'] = map_single(self._sys_blocks, 'enbs', **kwargs)
    self.save(self.results_path)
    return self

# %% ../nbs/07_pipeline.ipynb #67574163
@patch
def map_ccps(self:Report,
             force:bool=False,  # Re-run even if already completed
             **kwargs           # Additional args passed to map_single (e.g. path, model)
            ) -> Report:        # Self for chaining
    "Map report sections to Strategic Results Framework cross-cutting priorities"
    if 'ccps' in self.mappings and not force: return self
    self.ensure_sys_blocks()
    self.mappings['ccps'] = map_single(self._sys_blocks, 'ccps', **kwargs)
    self.save(self.results_path)
    return self

# %% ../nbs/07_pipeline.ipynb #5253ce14
@patch
def map_gcms(self:Report,
             force:bool=False,  # Re-run even if already completed
             **kwargs           # Additional args passed to map_single (e.g. path, model)
            ) -> Report:        # Self for chaining
    "Map report sections to Global Compact for Migration objectives"
    if 'gcms' in self.mappings and not force: return self
    self.ensure_sys_blocks()
    self.mappings['gcms'] = map_single(self._sys_blocks, 'gcms', **kwargs)
    self.save(self.results_path)
    return self

# %% ../nbs/07_pipeline.ipynb #c9110a90
@patch
def map_outs(self:Report,
             gcm_ids=None,      # GCM IDs to filter SRF objectives
             force:bool=False,  # Re-run even if already completed
             **kwargs           # Additional args passed to map_single (e.g. path, model)
            ) -> Report:        # Self for chaining
    "Map report sections to Strategic Results Framework outputs"
    if 'outs' in self.mappings and not force: return self
    self.ensure_sys_blocks()
    if gcm_ids is None:
        top_ids = get_top_ids(self.mappings.get('gcms', []))
        if not top_ids: return self
        gcm_ids = [top_ids[0]]
    self.mappings['outs'] = map_single(self._sys_blocks, 'outs', gcm_ids=gcm_ids, **kwargs)
    self.save(self.results_path)
    return self

# %% ../nbs/07_pipeline.ipynb #fc37da1c
@patch
def map_all(self:Report,
            **kwargs  # Args passed to all mapping methods
           ) -> Report:  # Self for chaining
    "Run all theme mappings in sequence"
    return self.map_enbs(**kwargs).map_ccps(**kwargs).map_gcms(**kwargs).map_outs(**kwargs)

# %% ../nbs/07_pipeline.ipynb #b63be3f5
def should_force(force,     # Bool to force all steps, or set of step names to force
                 step       # Step name to check
                ) -> bool:  # Whether to force the step
    "Check if step should be forced - handles bool or set of step names"
    if isinstance(force, bool): return force
    return step in force

# %% ../nbs/07_pipeline.ipynb #8cf30346
class PipelineResult:
    def __init__(self, report, status, error=None, step=None):
        store_attr()

# %% ../nbs/07_pipeline.ipynb #51cc0271
def get_current_step(report):
    "Infer which step the report is at based on its state"
    if not report.md_path: return 'ocr'
    if report.curation_status != 'sections_selected': return 'curation'
    return 'mapping'

# %% ../nbs/07_pipeline.ipynb #0dd9587d
async def run_pipeline(evals:list,                  # List of `Evaluation` objects to search
                       url:str=None,                # URL of the evaluation PDF
                       id:str=None,                 # Evaluation ID
                       title:str=None,              # Evaluation title
                       base_path:str='data',        # Base directory (contains pdf/, md/, results/)
                       add_img_desc:bool=True,      # Whether to add image descriptions during OCR
                       fix_kwargs:dict=None,        # Extra kwargs for fix_hdgs in OCR
                       desc_kwargs:dict=None,       # Extra kwargs for add_img_descs in OCR
                       force:bool|set=False,        # Force re-run: True for all, or set of step names
                       delete_pdf:bool=True,        # Delete PDF after OCR
                       verbose:bool=False,          # Print progress messages
                       **kwargs                     # Additional arguments passed to mapping functions
                      ) -> PipelineResult:          # Pipeline result with status and report
    "Run pipeline as far as possible: download → ocr → [curate] → map_all"
    if sum(x is not None for x in (url, id, title)) != 1:
        raise ValueError("Provide exactly one of: url, id, title")
    base = Path(base_path)
    pdf_dst, md_dst, results_path = base/'pdf', base/'md', base/'results'
    log = logger.info if verbose else lambda x: None
    
    if url: report = Report.from_url(url, evals, results_path=results_path)
    elif id: report = Report.from_id(id, evals, results_path=results_path)
    else: report = Report.from_title(title, evals, results_path=results_path)
    
    try: report = load_report(report.id, base_path=base_path)
    except FileNotFoundError: pass
    
    try:
        if not report.md_path or should_force(force, 'ocr'):
            if not report.pdf_path or should_force(force, 'download'):
                log("Downloading PDF...")
                report.download(dst=pdf_dst, force=should_force(force, 'download'))
            log("Running OCR...")
            await report.ocr(
                dst=md_dst, 
                add_img_desc=add_img_desc, 
                force=should_force(force, 'ocr'), 
                progress=verbose, 
                fix_kwargs=fix_kwargs, 
                desc_kwargs=desc_kwargs)

        if delete_pdf and report.pdf_path:
            shutil.rmtree(report.pdf_path)
            report.pdf_path = None
                
        if report.curation_status != 'sections_selected':
            report.save()
            log("Awaiting curation. Use curator app to select headings, then re-run.")
            return PipelineResult(report, 'awaiting_curation', step='curation')
        
        log("Mapping enablers...")
        report.map_enbs(force=should_force(force, 'enbs'), **kwargs)
        log("Mapping CCPs...")
        report.map_ccps(force=should_force(force, 'ccps'), **kwargs)
        log("Mapping GCM objectives...")
        report.map_gcms(force=should_force(force, 'gcms'), **kwargs)
        log("Mapping outputs...")
        report.map_outs(force=should_force(force, 'outs'), **kwargs)
        log("Pipeline complete!")
        return PipelineResult(report, 'completed')
    
    except Exception as e:
        step = get_current_step(report)
        logger.error(f"{report.id[:8]}... failed at {step}: {e}\n{traceback.format_exc()}")
        return PipelineResult(report, 'failed', error=str(e), step=step)

# %% ../nbs/07_pipeline.ipynb #b002bb4c
class BatchResult:
    "Result of running pipeline on a batch of evaluations"
    def __init__(self,
                 completed=None,         # List of successfully completed report IDs
                 awaiting_curation=None, # List of report IDs awaiting curation
                 failed=None,            # List of (report_id, error) tuples
                 skipped=None            # List of skipped report IDs (no report URL found)
                ):
        store_attr()
        self.completed = completed or []
        self.awaiting_curation = awaiting_curation or []
        self.failed = failed or []
        self.skipped = skipped or []
    
    def save(self, path):
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        data = dict(completed=self.completed, awaiting_curation=self.awaiting_curation, 
                    failed=self.failed, skipped=self.skipped)
        Path(path).write_text(json.dumps(data, indent=2))

# %% ../nbs/07_pipeline.ipynb #f72b5d42
def filter_evals(evals,      # List of Evaluation objects
                 year=None,  # Filter by year
                 ids=None    # Filter by list of IDs
                ):           # Filtered list of evaluations
    "Filter evaluations by year or list of IDs"
    if ids: return [e for e in evals if e.id in ids]
    if year: return [e for e in evals if e.meta.get('Year') == year]
    return evals


# %% ../nbs/07_pipeline.ipynb #7c691f19
def is_completed(report):
    "Check if report is completed"
    required = {'enbs', 'ccps', 'gcms', 'outs'}
    return required.issubset(report.mappings.keys()) if report.mappings else False

# %% ../nbs/07_pipeline.ipynb #9a36dbcc
async def batch_run(
    evals,                        # List of Evaluation objects to process
    base_path:str='../../data',   # Base directory (contains pdf/, md/, results/)
    year:str=None,                # Filter evaluations by year (e.g. '2023')
    ids:list=None,                # Filter evaluations by specific IDs
    force:bool|set=False,         # Force re-run: True for all, or set of step names
    stop_after:int=None,          # Stop after processing N reports (not yet implemented)
    delete_pdf:bool=False,        # Delete PDFs after OCR to save space
    add_img_desc:bool=True,       # Whether to add image descriptions during OCR
    fix_kwargs:dict=None,         # Extra kwargs for fix_hdgs in OCR
    desc_kwargs:dict=None,        # Extra kwargs for add_img_descs in OCR
    **kwargs                      # Additional arguments passed to run_pipeline function
) -> BatchResult:                 # Contains completed, awaiting_curation, failed, skipped lists
    "Run pipeline on a batch of reports"
    filtered = filter_evals(evals, year=year, ids=ids)
    result = BatchResult()
    
    for ev in filtered:
        try: report = load_report(ev.id, base_path=base_path)
        except FileNotFoundError: report = None
        
        if report and is_completed(report) and not force:
            logger.info(f"{ev.id[:8]}... skipped (already completed)")
            result.skipped.append(ev.id)
            continue
            
        url = eval_url(ev)
        if not url:
            logger.info(f"{ev.id[:8]}... failed (no report URL)")
            result.failed.append({'id': ev.id, 'step': 'url_lookup', 'error': 'Report URL not found'})
            continue
        
        pr = await run_pipeline(evals, url=url, base_path=base_path, force=force, delete_pdf=delete_pdf,
                                add_img_desc=add_img_desc, fix_kwargs=fix_kwargs, desc_kwargs=desc_kwargs, **kwargs)
        
        logger.info(f"{ev.id[:8]}... {pr.status}" + (f" ({pr.step}: {pr.error})" if pr.error else ""))
        
        if pr.status == 'completed': result.completed.append(ev.id)
        elif pr.status == 'awaiting_curation': result.awaiting_curation.append(ev.id)
        elif pr.status == 'failed': result.failed.append({'id': ev.id, 'step': pr.step, 'error': pr.error})

    timestamp = datetime.now().strftime('%Y-%m-%d_%H-%M')
    save_path = Path(base_path)/'batch_runs'/f'{timestamp}.json'
    result.save(save_path)
    return result

