"""OCR evaluation reports"""

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

# %% auto 0
__all__ = ['mistral_api_key', 'src_dir', 'get_doc_subtype', 'clean_pdf_name', 'setup_output_dirs', 'get_pdfs_and_dir',
           'save_page_images', 'process_batch_results', 'create_batch_ocr_job', 'submit_and_monitor_batch_job',
           'download_and_parse_results', 'process_single_evaluation_batch', 'process_all_reports_batch']

# %% ../nbs/03_ocr.ipynb 2
import os
import re
from io import BytesIO
from dotenv import load_dotenv
import base64
from fastcore.all import *
from rich import print
import urllib.parse
from pathlib import Path
from tqdm import tqdm
import json
import logging
import tempfile
from PIL import Image

from mistralai import Mistral
from .readers import load_evals

# %% ../nbs/03_ocr.ipynb 3
load_dotenv()
mistral_api_key = os.getenv("MISTRAL_API_KEY")

# %% ../nbs/03_ocr.ipynb 4
src_dir = Path("../_data/")

# %% ../nbs/03_ocr.ipynb 11
def get_doc_subtype(
    id:str, # ID of the evaluation
    fname:str, # Name of the file
    evals # Evaluations data
    )->str: # Document Subtype
    "Get Document Subtype for a given file in the evaluation dataset"
    eval_data = L(evals).filter(lambda x: x['id']==id)
    if not eval_data: return None
    
    docs = L(eval_data[0]['docs'])
    matches = docs.filter(lambda x: Path(x['File URL']).name==fname)
    return matches[0]['Document Subtype'] if matches else None

# %% ../nbs/03_ocr.ipynb 14
# Note: filename will be cleaned upstream in a next version
def clean_pdf_name(pdf_name: str) -> str:
    """
    Clean PDF name to create folder-friendly string.
    Removes special characters and spaces, replaces with underscores.
    """
    # Remove URL encoding
    pdf_name = urllib.parse.unquote(pdf_name)
    
    # Replace spaces and special characters with underscores
    # Replace any character that is not a word character (\w), whitespace (\s), or hyphen (-) with underscore
    cleaned = re.sub(r'[^\w\s-]', '_', pdf_name)
    
    # Replace any sequence of hyphens or whitespace with a single underscore
    cleaned = re.sub(r'[-\s]+', '_', cleaned)
    
    # Replace multiple consecutive underscores with a single underscore
    cleaned = re.sub(r'_+', '_', cleaned)
    cleaned = cleaned.strip('_')  # Remove leading/trailing underscores
    
    return cleaned.lower()

# %% ../nbs/03_ocr.ipynb 17
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('batch_ocr.log'),
        logging.StreamHandler()  # Also print to console
    ]
)

# %% ../nbs/03_ocr.ipynb 18
def setup_output_dirs(md_library_path="../_data/md_library"):
    "Set up the output directory structure for markdown files"
    md_output_dir = Path(md_library_path)
    mkdir(md_output_dir, parents=True, exist_ok=True, overwrite=False)
    return md_output_dir

# %% ../nbs/03_ocr.ipynb 19
def get_pdfs_and_dir(
    report_path:Path, # Path to the report directory
    md_output_dir:Path # Path to the output directory
    ) -> tuple[list[Path], str]:
    "Get PDFs from report directory and create output directory"
    pdfs = report_path.ls(file_exts='.pdf')
    eval_report_path = report_path.name
    mkdir(md_output_dir / eval_report_path, parents=True, exist_ok=True, overwrite=False)
    return pdfs, eval_report_path

# %% ../nbs/03_ocr.ipynb 22
def save_page_images(page, dest_folder: Path):
    "Save all images from a page to destination folder as PNG"
    images = page.images if hasattr(page, 'images') else page.get('images', [])
    
    for img in images:
        img_data = getattr(img, 'image_base64', img.get('image_base64'))
        img_id = getattr(img, 'id', img.get('id'))
        
        if img_data and img_id:
            img_bytes = base64.b64decode(img_data.split(',')[1])
            pil_img = Image.open(BytesIO(img_bytes))
            output_path = dest_folder / img_id
            pil_img.save(output_path)

# %% ../nbs/03_ocr.ipynb 24
def process_batch_results(results, md_output_dir):
    "Process batch OCR results and save to appropriate folders"
    for result in results:
        try:
            # Parse custom_id to get eval_id and pdf_name
            eval_id, pdf_name = result['custom_id'].split('_', 1)
            
            # Get OCR response
            ocr_response = result['response']['body']
            
            # Create folder structure
            pdf_clean_name = clean_pdf_name(pdf_name)
            pdf_dir = md_output_dir / eval_id / pdf_clean_name
            pdf_dir.mkdir(parents=True, exist_ok=True)
            
            # Save each page markdown
            for page in ocr_response['pages']:
                page_num = page['index'] + 1
                page_path = pdf_dir / f"page_{page_num}.md"
                page_path.write_text(page['markdown'])
            
            # Save images if they exist
            img_dir = pdf_dir / 'img'
            for page in ocr_response['pages']:
                if page.get('images'):
                    img_dir.mkdir(parents=True, exist_ok=True)
                    save_page_images(page, img_dir)
            
            logging.info(f"Saved {len(ocr_response['pages'])} pages for {pdf_clean_name}")
            
        except Exception as e:
            logging.error(f"Error processing result {result.get('custom_id', 'unknown')}: {e}")

# %% ../nbs/03_ocr.ipynb 25
def create_batch_ocr_job(
    pdf_paths: List[Path],
    eval_report_path: str,
    model: str = "mistral-ocr-latest",
    include_images: bool = True,
    api_key: str = mistral_api_key
):
    "Create batch entries for PDFs from one evaluation report"
    cli = Mistral(api_key=api_key)
    
    batch_entries = []
    for pdf_path in pdf_paths:
        uploaded_pdf = cli.files.upload(
            file={
                "file_name": pdf_path.stem,
                "content": pdf_path.read_bytes(),
            },
            purpose="ocr"
        )
        
        signed_url = cli.files.get_signed_url(file_id=uploaded_pdf.id)
        entry = {
            "custom_id": f"{eval_report_path}_{pdf_path.name}",
            "body": {
                "document": {
                    "type": "document_url",
                    "document_url": signed_url.url,
                },
                "include_image_base64": include_images
            }
        }
        batch_entries.append(entry)
        logging.info(f"Added {pdf_path.name} to batch for {eval_report_path}")
        
    return batch_entries, cli


# %% ../nbs/03_ocr.ipynb 26
def submit_and_monitor_batch_job(batch_entries, eval_report_path, cli):
    "Submit batch job and monitor until completion"
    with tempfile.NamedTemporaryFile(mode='w', suffix='.jsonl', delete=True) as temp_file:
        # Write batch entries to temp file
        for entry in batch_entries:
            temp_file.write(json.dumps(entry) + '\n')
        temp_file.flush()
        
        # Upload and create job
        batch_data = cli.files.upload(
            file={"file_name": f"batch_{eval_report_path}.jsonl", 
                  "content": open(temp_file.name, "rb")},
            purpose="batch"
        )
        
        created_job = cli.batch.jobs.create(
            input_files=[batch_data.id],
            model="mistral-ocr-latest",
            endpoint="/v1/ocr"
        )
        
        logging.info(f"Batch job created for {eval_report_path}: {created_job.id}")
        
        # Monitor completion
        while True:
            job = cli.batch.jobs.get(job_id=created_job.id)
            logging.info(f"Job status: {job.status} - {job.succeeded_requests}/{job.total_requests} completed")
            
            if job.status not in ["QUEUED", "RUNNING"]:
                break
            time.sleep(10)
        
        return job

# %% ../nbs/03_ocr.ipynb 27
def download_and_parse_results(job, cli):
    "Download and parse batch job results"
    response = cli.files.download(file_id=job.output_file)
    content = response.read().decode('utf-8')
    
    results = []
    for line in content.strip().split('\n'):
        if line:
            results.append(json.loads(line))
    
    logging.info(f"Downloaded and parsed {len(results)} OCR results")
    return results

# %% ../nbs/03_ocr.ipynb 28
def process_single_evaluation_batch(report: Path, md_output_dir: Path):
    "Process one evaluation report using batch OCR"
    logging.info(f"Processing evaluation: {report.name}")
    pdfs, eval_report_path = get_pdfs_and_dir(report, md_output_dir)
    
    if not pdfs:
        logging.warning(f"No PDFs found in {eval_report_path}")
        return
    
    batch_entries, cli = create_batch_ocr_job(pdfs, eval_report_path)
    
    job = submit_and_monitor_batch_job(batch_entries, eval_report_path, cli)
    
    if job and job.status == "SUCCESS":
        results = download_and_parse_results(job, cli)
        process_batch_results(results, md_output_dir)
        logging.info(f"Completed processing evaluation: {eval_report_path}")
    else:
        logging.error(f"Job failed for {eval_report_path}")

# %% ../nbs/03_ocr.ipynb 29
def process_all_reports_batch(
    reports: list[Path],
    md_library_path="../_data/md_library"
):
    "Process evaluation reports using batch OCR"
    logging.info(f"Starting batch OCR processing for {len(reports)} reports")
    md_output_dir = setup_output_dirs(md_library_path)
    
    for report in tqdm(reports, desc="Processing reports"):
        process_single_evaluation_batch(report, md_output_dir)
    
    logging.info("Batch OCR processing completed for all reports")
