"""PDF -> source-traceable, reflowable EPUB. No translation calls or heavy models."""
import hashlib
import html
import json
import platform
import re
import sys
import uuid
from pathlib import Path
from ebooklib import epub
from .layout import paragraphs, repeated_furniture, join_text

SCHEMA = 1


def digest(path):
    h = hashlib.sha256()
    with Path(path).open('rb') as f:
        for chunk in iter(lambda:f.read(1024*1024), b''): h.update(chunk)
    return h.hexdigest()


def save_json(path, data):
    path = Path(path); path.parent.mkdir(parents=True, exist_ok=True)
    temp = path.with_suffix('.tmp'); temp.write_text(json.dumps(data, ensure_ascii=False, indent=2)+'\n'); temp.replace(path)


def page_selection(spec, count):
    if not spec: return list(range(count))
    selected = set()
    for item in spec.split(','):
        if not re.fullmatch(r'\d+(?:-\d+)?', item): raise ValueError('Pages must look like 1-20,25 (1-based inclusive)')
        pair = [int(x) for x in item.split('-')]; lo, hi = pair[0], pair[-1]
        if not 1 <= lo <= hi <= count: raise ValueError('Page range is outside the source PDF')
        selected.update(range(lo-1, hi))
    return sorted(selected)


def native_lines(page):
    text = page.get_textpage()
    try:
        lines=[]; width,height=page.get_size()
        for i in range(text.count_rects()):
            left,bottom,right,top=text.get_rect(i)
            value=text.get_text_bounded(left,bottom,right,top).strip()
            if value:
                lines.append({'text':value,'bbox':[left/width,1-top/height,right/width,1-bottom/height], 'confidence':None})
        return lines
    finally: text.close()


def image_boxes(page):
    import pypdfium2.raw as raw
    w,h=page.get_size(); boxes=[]
    for obj in page.get_objects(filter=[raw.FPDF_PAGEOBJ_IMAGE]):
        left,bottom,right,top=obj.get_bounds()
        boxes.append([max(0,left/w), max(0,1-top/h), min(1,right/w), min(1,1-bottom/h)])
    return boxes


def validate_lines(lines):
    if not isinstance(lines,list): raise ValueError('OCR lines must be an array')
    for line in lines:
        box=line.get('bbox',[])
        if not isinstance(line.get('text'),str) or len(box)!=4 or not all(isinstance(v,(int,float)) and 0<=v<=1 for v in box) or box[0]>=box[2] or box[1]>=box[3]:
            raise ValueError('Invalid OCR text or normalized top-left bounding box')


def apply_corrections(pages, correction, source_sha):
    """Source-bound exact edits; never overwrite the raw OCR checkpoint."""
    if correction.get('source_sha256')!=source_sha or not isinstance(correction.get('edits'),list):
        raise ValueError('Corrections must be bound to this PDF SHA and contain an edits array')
    lines={line['id']:line for page in pages for line in page['lines']};applied=[];seen=set()
    for edit in correction['edits']:
        key=edit.get('line_id');line=lines.get(key)
        if key in seen or line is None or line['text']!=edit.get('expected') or not isinstance(edit.get('replacement'),str) or not edit['replacement'].strip():
            raise ValueError('Stale, duplicate or invalid correction: '+str(key))
        line['text']=edit['replacement'];seen.add(key);applied.append(edit)
    return applied


def convert(source, output, work, language='zh-CN', title=None, author=None,
            engine='auto', pages=None, dpi=180, cover_page=None, ocr_timeout=120, image_pages=None, corrections=None, ocr_url=None, ocr_model='firered-ocr', ocr_prompt=None, geometry_work=None):
    import pypdfium2 as pdfium
    source,output,work=Path(source),Path(output),Path(work)
    if not source.is_file(): raise ValueError('PDF source does not exist')
    if source.resolve()==output.resolve(): raise ValueError('Output must not overwrite source')
    if output.suffix.lower()!='.epub': raise ValueError('Output must have .epub suffix')
    if not 96<=dpi<=400: raise ValueError('DPI must be between 96 and 400')
    work.mkdir(parents=True,exist_ok=True); source_sha=digest(source)
    source_identity=work/'source.json'
    if source_identity.exists() and json.loads(source_identity.read_text())['sha256']!=source_sha:
        raise ValueError('Work directory belongs to another source; use a new directory')
    save_json(source_identity,{'path':str(source.resolve()),'sha256':source_sha})
    document=pdfium.PdfDocument(source); ocr=None
    try:
        selected=page_selection(pages,len(document)); partial=len(selected)!=len(document)
        preserved_images=set(i+1 for i in page_selection(image_pages,len(document))) if image_pages else set()
        if cover_page is not None and not 1<=cover_page<=len(document): raise ValueError('Cover page outside source PDF')
        geometry_dir=Path(geometry_work) if geometry_work else None
        if geometry_dir and json.loads((geometry_dir/'source.json').read_text())['sha256']!=source_sha:
            raise ValueError('Reference geometry belongs to another PDF')
        backend=('vision' if platform.system()=='Darwin' else 'tesseract') if engine=='auto' else engine
        outlines=[]
        for bookmark in document.get_toc():
            dest=bookmark.get_dest()
            if dest and dest.get_index() in selected:
                outlines.append({'title':bookmark.get_title(),'page':dest.get_index()+1,'level':bookmark.level})
        results=[]; warnings=[]
        for index in selected:
            page=document[index]
            try:
                boxes=image_boxes(page)
                scan=any((b[2]-b[0])*(b[3]-b[1])>.65 for b in boxes)
                native=native_lines(page) if engine=='native' or (engine=='auto' and not scan) else []
                text_size=sum(len(l['text']) for l in native)
                use_ocr=engine not in {'native','auto'} or (engine=='auto' and (scan or (text_size==0 and bool(boxes))))
                is_cover=index+1==cover_page
                is_image=is_cover or index+1 in preserved_images
                if use_ocr and not is_image and ocr is None:
                    if backend=='firered':
                        from .firered import FireRedOCR
                        ocr=FireRedOCR(work,language,ocr_timeout,base_url=ocr_url,model=ocr_model,prompt_path=ocr_prompt)
                    else:
                        from .ocr import VisionOCR,TesseractOCR
                        ocr=(VisionOCR if backend=='vision' else TesseractOCR)(work,language,ocr_timeout)
                used='cover-image' if is_cover else 'page-image' if is_image else backend if use_ocr else 'native'
                stamp={'source':source_sha,'page':index+1,'engine':used,'ocr_identity':ocr.identity if use_ocr and not is_image else None,'language':language,'dpi':dpi,'pdfium':str(pdfium.PYPDFIUM_INFO),'schema':SCHEMA}
                cache=work/'pages'/f'{index+1:04d}.json'
                cached=json.loads(cache.read_text()) if cache.exists() else None
                render=work/'renders'/f'{index+1:04d}.jpg';render.parent.mkdir(exist_ok=True)
                if cached and cached.get('stamp')==stamp and render.exists():
                    result=cached
                else:
                    bitmap=page.render(scale=dpi/72)
                    try: bitmap.to_pil().convert('RGB').save(render,quality=88)
                    finally: bitmap.close()
                    recognition=[] if is_image else ocr.recognize(render) if use_ocr else native
                    extra=recognition if isinstance(recognition,dict) else {}
                    lines=extra.get('lines',[]) if isinstance(recognition,dict) else recognition
                    validate_lines(lines)
                    for n,line in enumerate(lines): line['id']=f'p{index+1:04d}-l{n+1:04d}'
                    result={'number':index+1,'width':page.get_width(),'height':page.get_height(),'engine':used,'scan':scan,'lines':lines,'image_boxes':boxes,'render':str(render.resolve()),'stamp':stamp,**extra}
                    save_json(cache,result)
                results.append(result)
                print(f'PDF page {index+1}/{len(document)}: {used}, {len(result["lines"])} lines',file=sys.stderr,flush=True)
            finally: page.close()
        applied=apply_corrections(results,json.loads(Path(corrections).read_text()),source_sha) if corrections else []
        repeated=repeated_furniture(results)
        for page in results:
            if 'markdown' in page:
                from .markdown_layout import markdown_blocks
                geometry=json.loads((geometry_dir/'pages'/f'{page["number"]:04d}.json').read_text()) if geometry_dir else None
                if geometry and (geometry['number']!=page['number'] or geometry['stamp']['source']!=source_sha):
                    raise ValueError('Reference page geometry provenance mismatch')
                blocks=markdown_blocks(page,language,geometry);removed=[]
            else:
                blocks,removed=paragraphs(page,language,repeated)
            page['blocks']=blocks;page['removed_furniture']=removed
            count=sum(len(b['text']) for b in blocks) if 'markdown' in page else sum(len(l['text']) for l in page['lines'])
            page['image_fallback']=page['engine'] in {'cover-image','page-image'} or count<30
            if page['image_fallback'] and page['engine'] not in {'cover-image','page-image'}:
                warnings.append({'page':page['number'],'kind':'page_image_fallback','detail':'Blank/illustrated/short-text page retained as an image; inspect source, not claimed as reflowed text.'})
            removed_ids={l['id'] for l in removed}
            low=[l['id'] for l in page['lines'] if l['id'] not in removed_ids and l.get('confidence') is not None and l['confidence']<.6]
            if low: warnings.append({'page':page['number'],'kind':'low_ocr_confidence','line_ids':low})
            if engine=='native' and page['scan']: warnings.append({'page':page['number'],'kind':'existing_scan_text_layer','detail':'Existing OCR text layer reused without verification.'})
        book={'schema':SCHEMA,'source_sha256':source_sha,'source':str(source.resolve()),'source_pages':len(document),'selected_pages':[i+1 for i in selected],'partial':partial,'language':language,'title':title or source.stem,'author':author or 'Unknown','outlines':outlines,'image_pages':sorted(preserved_images),'cover_page':cover_page,'geometry_work':str(geometry_dir.resolve()) if geometry_dir else None,'pages':results,'warnings':warnings,'applied_corrections':applied,'scope':'horizontal single-column; no automatic footnote/table reconstruction'}
        save_json(work/'book.json',book)
        export_epub(book,output,work,cover_page)
        from src.qa.epub_audit import audit_epub
        audit=audit_epub(output)
        report={'source_sha256':source_sha,'output_sha256':digest(output),'output':str(output.resolve()),'partial':partial,'processed_pages':len(results),'source_pages':len(document),'image_only_pages':[p['number'] for p in results if p['image_fallback']],'removed_furniture_lines':sum(len(p['removed_furniture']) for p in results),'engines':{name:sum(p['engine']==name for p in results) for name in set(p['engine'] for p in results)},'warnings':warnings,'applied_corrections':applied,'epub_audit':audit,'visual_review':'pending','ocr_accuracy_review':'pending','translation':'not requested; English EPUB can be passed to main.py separately'}
        save_json(work/'report.json',report)
        if audit['issues']: raise ValueError('EPUB structural audit failed; inspect report.json')
        return report
    finally:
        if ocr: ocr.close()
        document.close()


STYLE='''body {margin: 5%; line-height:1.65;} p {margin:.5em 0;} html[lang^="zh"] p {text-indent:2em;} h1,h2 {text-align:center;page-break-before:always;} img {max-width:100%;height:auto;} figure {margin:1em 0;text-align:center;} figcaption {font-size:.8em;} .source-page {font-size:.7em;color:#777;text-indent:0;} .attribution {text-align:right !important;text-indent:0 !important;margin:.4em 0;} .image-page {text-align:center;text-indent:0;} .image-page img {max-height:90vh;width:auto;object-fit:contain;}'''


def export_epub(data, output, work, cover_page):
    from PIL import Image
    import io
    book=epub.EpubBook(); identity={k:data[k] for k in ['source_sha256','language','title','author','selected_pages']}
    identity['content_sha256']=hashlib.sha256(json.dumps([[p['number'],p['image_fallback'],[(b['kind'],b['text']) for b in p['blocks']]] for p in data['pages']],ensure_ascii=False).encode()).hexdigest()
    book.set_identifier('urn:uuid:'+str(uuid.uuid5(uuid.NAMESPACE_URL,json.dumps(identity,sort_keys=True))))
    label='【部分样例】' if data['partial'] else ''
    book.set_title(label+data['title']+'【PDF重排初稿】');book.set_language(data['language']);book.add_author(data['author'])
    book.add_metadata('DC','source','urn:sha256:'+data['source_sha256'])
    book.add_metadata('DC','description','PDF结构重建初稿；OCR未逐字校对。'+('仅含选定页，非完整书。' if data['partial'] else ''))
    css=epub.EpubItem(uid='style',file_name='styles/book.css',media_type='text/css',content=STYLE.encode());book.add_item(css)
    chapters=[]; toc_records=[]; body=[]; current=None; previous=None; previous_page=None
    def finish():
        if current is not None:
            current.content=('<html xmlns:epub="http://www.idpf.org/2007/ops"><head></head><body>'+''.join(body)+'</body></html>').encode()
            current.add_link(href='../styles/book.css',rel='stylesheet',type='text/css')
            book.add_item(current);chapters.append(current)
    for page in data['pages']:
        number=page['number']; outlines=[o for o in data['outlines'] if o['page']==number]
        labels=[o['title'] for o in outlines]
        if current is None or labels or page['image_fallback'] or previous_page is None or number!=previous_page+1:
            finish();body=[];previous=None
            current=epub.EpubHtml(title=' / '.join(labels) or f'PDF {number}',file_name=f'chapters/p{number:04d}.xhtml',lang=data['language'])
        for outline in outlines:
            toc_records.append((epub.Link(current.file_name+f'#pdf-page-{number}',outline['title'],f'nav-{len(toc_records)+1}'),outline['level']))
        marker=f'<span id="pdf-page-{number}" epub:type="pagebreak" role="doc-pagebreak" aria-label="PDF {number}"></span>'
        if not page['image_fallback']:
            for heading in labels: body.append('<h1>'+html.escape(heading)+'</h1>')
        if page['image_fallback']:
            body.append(marker)
            image=Path(page['render']).read_bytes();name=f'images/page-{number:04d}.jpg'
            book.add_item(epub.EpubItem(uid=f'image-{number}',file_name=name,media_type='image/jpeg',content=image))
            body.append(f'<div class="image-page"><img src="../{name}" alt="PDF第{number}页原貌"/></div>')
            if number==cover_page: book.set_cover('images/cover.jpg',image,create_page=False)
            previous=None
        else:
            entries=[dict(b) for b in page['blocks']]
            for n,box in enumerate(page['image_boxes']):
                area=(box[2]-box[0])*(box[3]-box[1])
                if .02<area<.65: entries.append({'kind':'image','bbox':box,'image_index':n})
            if all(b.get('bbox') for b in entries):entries.sort(key=lambda b:(b['bbox'][1],b['bbox'][0]))
            marker_pending=True
            for index,block in enumerate(entries):
                if block['kind']=='image':
                    with Image.open(page['render']) as picture:
                        box=block['bbox'];region=picture.crop(tuple(int(v*(picture.width if i%2==0 else picture.height)) for i,v in enumerate(box)))
                        stream=io.BytesIO();region.save(stream,format='PNG')
                    n=block['image_index'];name=f'images/p{number:04d}-{n}.png'
                    book.add_item(epub.EpubItem(uid=f'figure-{number}-{n}',file_name=name,media_type='image/png',content=stream.getvalue()))
                    if marker_pending:body.append(marker);marker_pending=False
                    body.append(f'<figure><img src="../{name}" alt="PDF第{number}页图片"/></figure>');previous=None;continue
                plain=re.sub(r'\s','',block['text'])
                if any(plain==re.sub(r'\s','',label) or (index==0 and plain in {'序','序言','前言','自序','引言','后记','目录'} and re.sub(r'\s','',label).startswith(plain)) for label in labels):continue
                # Merge only obvious prose continuation across consecutive source pages.
                continuation=marker_pending and previous and not labels and previous_page==number-1 and previous['kind']=='paragraph' and block['kind']=='paragraph' and previous.get('role')!='attribution' and block.get('role')!='attribution' and not block['indented'] and not re.search(r'[。！？.!?:：;；”’"]$',previous['text']) and body[-1].endswith('</p>')
                if continuation:
                    extra=join_text(previous['text'],block['text'],data['language'])[len(previous['text']):]
                    body[-1]=body[-1][:-4]+marker+f'<span id="{block["id"]}"></span>'+html.escape(extra)+'</p>'
                    marker_pending=False;previous={**block,'text':join_text(previous['text'],block['text'],data['language'])};continue
                if marker_pending:body.append(marker);marker_pending=False
                tag='h2' if block['kind']=='heading' else 'p'
                attributes=' class="attribution"' if block.get('alignment')=='right' else ''
                body.append(f'<{tag} id="{block["id"]}"{attributes}>{html.escape(block["text"])}</{tag}>')
                if block['kind']=='heading' and not data['outlines']:
                    toc_records.append((epub.Link(current.file_name+'#'+block['id'],block['text'],f'nav-{len(toc_records)+1}'),0))
                previous=block
            if marker_pending:body.append(marker)
        previous_page=number
        if page['image_fallback']:
            finish();current=None;body=[];previous=None
    finish()
    if not toc_records:toc_records=[(epub.Link(c.file_name,c.title,f'nav-{i+1}'),0) for i,c in enumerate(chapters)]
    roots=[];stack=[]
    for link,level in toc_records:
        node={'link':link,'children':[]}
        while stack and stack[-1][0]>=level:stack.pop()
        (stack[-1][1]['children'] if stack else roots).append(node);stack.append((level,node))
    def flatten(nodes):return [(node['link'],flatten(node['children'])) if node['children'] else node['link'] for node in nodes]
    book.toc=flatten(roots);book.add_item(epub.EpubNcx());nav=epub.EpubNav();nav.title='目录' if data['language'].startswith('zh') else 'Contents';book.add_item(nav);book.spine=['nav']+chapters
    output=Path(output);output.parent.mkdir(parents=True,exist_ok=True);temp=output.with_suffix('.tmp.epub')
    try:
        epub.write_epub(str(temp),book,{'epub3_pages':False})
        from src.qa.epub_audit import audit_epub
        if audit_epub(temp)['issues']:raise ValueError('Generated EPUB failed structural audit; existing output retained')
        temp.replace(output)
    finally:temp.unlink(missing_ok=True)