"""Document restoration prototype. Never silently replace OCR text with model guesses."""
import os, io, re, time, base64, json
from pathlib import Path
import cv2
import fitz
import numpy as np
import requests
from PIL import Image, ImageOps, ImageEnhance

MAX_PAGES=int(os.getenv('MAX_PAGES','15'))
MAX_PIXELS=int(os.getenv('MAX_PIXELS','16000000'))

def rasterize(data, filename):
    pages=[]
    if filename.lower().endswith('.pdf'):
        src=fitz.open(stream=data,filetype='pdf')
        if len(src)>MAX_PAGES: raise ValueError(f'Maximum {MAX_PAGES} pages par document.')
        for page in src:
            pix=page.get_pixmap(matrix=fitz.Matrix(1.7,1.7),alpha=False)
            pages.append(Image.frombytes('RGB',(pix.width,pix.height),pix.samples))
        src.close()
    else:
        img=ImageOps.exif_transpose(Image.open(io.BytesIO(data))).convert('RGB')
        pages=[img]
    for im in pages:
        if im.width*im.height>MAX_PIXELS: raise ValueError('Image trop grande.')
    return pages

def enhance(image):
    """Conservative enhancement; never discard original color information."""
    a=np.array(image.convert('RGB'))
    gray=cv2.cvtColor(a,cv2.COLOR_RGB2GRAY)
    # Rotate only if Hough line evidence is sufficiently strong.
    edges=cv2.Canny(gray,60,180)
    lines=cv2.HoughLinesP(edges,1,np.pi/180,threshold=95,minLineLength=max(70,a.shape[1]//4),maxLineGap=20)
    angles=[]
    if lines is not None:
        for l in lines[:,0]:
            x1,y1,x2,y2=map(int,l)
            angle=np.degrees(np.arctan2(y2-y1,x2-x1))
            if abs(angle)<12: angles.append(angle)
    if len(angles)>=5:
        angle=float(np.median(angles))
        if 0.5<abs(angle)<8:
            h,w=a.shape[:2]; M=cv2.getRotationMatrix2D((w/2,h/2),angle,1)
            a=cv2.warpAffine(a,M,(w,h),flags=cv2.INTER_CUBIC,borderMode=cv2.BORDER_CONSTANT,borderValue=(255,255,255))
    # CLAHE on luminance only, retain colored stamps/signatures.
    lab=cv2.cvtColor(a,cv2.COLOR_RGB2LAB)
    lab[:,:,0]=cv2.createCLAHE(clipLimit=1.5,tileGridSize=(8,8)).apply(lab[:,:,0])
    return Image.fromarray(cv2.cvtColor(lab,cv2.COLOR_LAB2RGB))

def azure_layout(image):
    endpoint=os.getenv('AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT','').rstrip('/')
    key=os.getenv('AZURE_DOCUMENT_INTELLIGENCE_KEY','')
    if not endpoint or not key: raise RuntimeError('Azure Document Intelligence non configuré.')
    ver=os.getenv('AZURE_DOCUMENT_INTELLIGENCE_API_VERSION','2024-11-30')
    buf=io.BytesIO(); image.save(buf,format='PNG')
    url=f'{endpoint}/documentintelligence/documentModels/prebuilt-layout:analyze?api-version={ver}'
    r=requests.post(url,headers={'Ocp-Apim-Subscription-Key':key,'Content-Type':'application/octet-stream'},data=buf.getvalue(),timeout=45)
    r.raise_for_status()
    poll=r.headers.get('Operation-Location')
    if not poll or not poll.startswith(endpoint+'/'): raise RuntimeError('Azure: URL de suivi invalide.')
    result=None
    for _ in range(35):
        p=requests.get(poll,headers={'Ocp-Apim-Subscription-Key':key},timeout=25)
        p.raise_for_status(); result=p.json()
        if result.get('status')=='succeeded': break
        if result.get('status')=='failed': raise RuntimeError('Azure: analyse échouée.')
        time.sleep(1)
    else: raise RuntimeError('Azure: délai de traitement dépassé.')
    ar=result.get('analyzeResult',{}); ap=ar.get('pages',[])
    if not ap: raise RuntimeError('Azure: aucune page analysée.')
    page=ap[0]; pw=float(page.get('width') or image.width); ph=float(page.get('height') or image.height)
    words=[]
    for w in page.get('words',[]):
        poly=w.get('polygon',[])
        if len(poly)<8: continue
        xs=poly[::2]; ys=poly[1::2]
        words.append({'text':w.get('content',''),'box':[min(xs)/pw,min(ys)/ph,max(xs)/pw,max(ys)/ph], 'confidence':w.get('confidence')})
    tables=[]
    for t in ar.get('tables',[]):
        if any((b.get('pageNumber',1)!=1) for b in t.get('boundingRegions',[])): continue
        cells=[]
        for c in t.get('cells',[]):
            polys=[b.get('polygon',[]) for b in c.get('boundingRegions',[]) if b.get('pageNumber',1)==1]
            poly=polys[0] if polys else []
            xs=poly[::2]; ys=poly[1::2]
            cells.append({'row':c.get('rowIndex',0),'col':c.get('columnIndex',0),
                'rowspan':c.get('rowSpan',1),'colspan':c.get('columnSpan',1),
                'text':c.get('content',''),'kind':c.get('kind','content'),
                'box':[min(xs)/pw,min(ys)/ph,max(xs)/pw,max(ys)/ph] if xs and ys else None})
        tables.append({'rows':t.get('rowCount',0),'cols':t.get('columnCount',0),'cells':cells})
    return words, {'provider':'azure','tables':len(tables),'table_data':tables,'words':len(words)}

def local_layout(image):
    try: import pytesseract
    except ImportError: raise RuntimeError('Installer pytesseract et Tesseract OCR ou configurer Azure.')
    try: d=pytesseract.image_to_data(image,output_type=pytesseract.Output.DICT)
    except pytesseract.TesseractNotFoundError: raise RuntimeError('Tesseract OCR absent du serveur. Configurer Azure ou installer Tesseract.')
    words=[]
    for i,t in enumerate(d['text']):
        if not t.strip(): continue
        x,y,w,h=(int(d[k][i]) for k in ('left','top','width','height'))
        conf=float(d['conf'][i]) if str(d['conf'][i])!='-1' else 0
        words.append({'text':t,'box':[x/image.width,y/image.height,(x+w)/image.width,(y+h)/image.height], 'confidence':max(0,conf/100)})
    return words,{'provider':'tesseract','tables':None,'words':len(words)}

def openai_review(image, words):
    """Non-authoritative QA: flag likely OCR errors, NEVER change legal/financial values."""
    if not os.getenv('OPENAI_API_KEY'): return {'status':'not_configured','warnings':[]}
    from openai import OpenAI
    buff=io.BytesIO(); image.thumbnail((1500,1800)); image.save(buff,format='JPEG',quality=82)
    uri='data:image/jpeg;base64,'+base64.b64encode(buff.getvalue()).decode()
    ocr=' '.join(w['text'] for w in words)[:9000]
    response=OpenAI(api_key=os.getenv('OPENAI_API_KEY'),timeout=45).chat.completions.create(
      model=os.getenv('OPENAI_MODEL','gpt-4.1-mini'),temperature=0,
      response_format={'type':'json_object'},messages=[
      {'role':'system','content':'You are a document OCR QA assistant. Return JSON {"warnings":["..."],"document_type":"..."}. Do not infer, fabricate, or rewrite any data. Warn about uncertain dates, amounts, names and illegible zones. Max 8 warnings. Answer in French.'},
      {'role':'user','content':[{'type':'text','text':'Review OCR against image. OCR: '+ocr},{'type':'image_url','image_url':{'url':uri}}]}
      ])
    return {'status':'completed',**json.loads(response.choices[0].message.content)}

def make_pdf(images, layouts, mode):
    doc=fitz.open()
    for im,words in zip(images,layouts):
        w,h=im.size; page=doc.new_page(width=w,height=h)
        if mode in ('clean','searchable'):
            buff=io.BytesIO(); im.save(buff,format='PNG')
            page.insert_image(fitz.Rect(0,0,w,h),stream=buff.getvalue())
        if mode in ('searchable','rebuild'):
            for item in words:
                s=item['text']; x0,y0,x1,y1=item['box']; box=fitz.Rect(x0*w,y0*h,x1*w,y1*h)
                if box.is_empty or box.is_infinite: continue
                # Exact source words, no AI rewriting. Text fit is approximate.
                size=max(5,min(32,box.height*.85))
                estimated=fitz.get_text_length(s,fontname='helv',fontsize=size)
                if estimated>box.width and estimated>0: size=max(3,size*box.width/estimated)
                if mode=='searchable':
                    page.insert_text((box.x0,box.y1-box.height*.12),s,fontsize=size,fontname='helv',render_mode=3)
                else:
                    page.insert_text((box.x0,box.y1-box.height*.12),s,fontsize=size,fontname='helv',color=(0.08,0.13,0.20))
    output=doc.tobytes(garbage=4,deflate=True); doc.close(); return output


def _valid_box(box):
    return box and len(box)==4 and all(isinstance(v,(int,float)) and 0<=v<=1 for v in box) and box[2]>box[0] and box[3]>box[1]

def draw_tables_on_pdf(pdf_bytes, images, diagnostics, mode):
    """Render Azure-detected cell boundaries and text; suppress underlying OCR words in table areas upstream."""
    if mode!='rebuild': return pdf_bytes
    doc=fitz.open(stream=pdf_bytes,filetype='pdf')
    for page,im,meta in zip(doc,images,diagnostics):
        w,h=im.size
        for table in meta.get('table_data',[]):
            for cell in table.get('cells',[]):
                box=cell.get('box')
                if not _valid_box(box): continue
                r=fitz.Rect(box[0]*w,box[1]*h,box[2]*w,box[3]*h)
                page.draw_rect(r,color=(.45,.51,.60),width=.7)
                text=cell.get('text','').strip()
                if not text: continue
                inner=r+ (2,2,-2,-2)
                if inner.width<4 or inner.height<4: continue
                size=min(12,max(5,inner.height/2.3))
                while size>=4:
                    spare=page.insert_textbox(inner,text,fontsize=size,fontname='helv',lineheight=1.05,
                                             color=(.08,.13,.2),overlay=True)
                    if spare>=0: break
                    size-=.5
    data=doc.tobytes(garbage=4,deflate=True);doc.close();return data

def export_docx(images, layouts, diagnostics):
    """Editable semantic document. Word pagination is approximate, not pixel-perfect."""
    from docx import Document
    from docx.shared import Inches,Pt
    doc=Document(); sec=doc.sections[0];sec.top_margin=Inches(.7);sec.bottom_margin=Inches(.7)
    for pi,(im,words,meta) in enumerate(zip(images,layouts,diagnostics)):
        if pi: doc.add_page_break()
        tables=meta.get('table_data',[])
        boxes=[c['box'] for t in tables for c in t['cells'] if _valid_box(c.get('box'))]
        outside=[w for w in words if not any(b[0]<=((w['box'][0]+w['box'][2])/2)<=b[2] and b[1]<=((w['box'][1]+w['box'][3])/2)<=b[3] for b in boxes)]
        # Group OCR words into lines based on their vertical centers.
        lines=[]
        for w in sorted(outside,key=lambda x:(x['box'][1],x['box'][0])):
            cy=(w['box'][1]+w['box'][3])/2
            target=next((line for line in lines if abs(line[0]-cy)<.009),None)
            if target is None: lines.append([cy,[w]])
            else: target[1].append(w)
        blocks=[(y,'line',ws) for y,ws in lines]
        for t in tables:
            positions=[c['box'][1] for c in t['cells'] if _valid_box(c.get('box'))]
            blocks.append((min(positions) if positions else 1,'table',t))
        for _,kind,obj in sorted(blocks,key=lambda b:b[0]):
            if kind=='line':
                paragraph=doc.add_paragraph(' '.join(w['text'] for w in sorted(obj,key=lambda x:x['box'][0])))
                paragraph.paragraph_format.space_after=Pt(2)
            else:
                rows,cols=obj.get('rows',0),obj.get('cols',0)
                if not (0<rows<=250 and 0<cols<=80): continue
                tab=doc.add_table(rows=rows,cols=cols);tab.style='Table Grid'
                for c in obj['cells']:
                    r,k=c['row'],c['col']
                    if r>=rows or k>=cols:continue
                    rr=min(rows-1,r+max(1,c['rowspan'])-1);cc=min(cols-1,k+max(1,c['colspan'])-1)
                    try: cell=tab.cell(r,k).merge(tab.cell(rr,cc)) if (rr!=r or cc!=k) else tab.cell(r,k)
                    except (ValueError,IndexError):cell=tab.cell(r,k)
                    cell.text=c['text']
                    if c.get('kind') in ('columnHeader','rowHeader'):
                        for paragraph in cell.paragraphs:
                            for run in paragraph.runs:run.bold=True
    out=io.BytesIO();doc.save(out);return out.getvalue()

def export_xlsx(diagnostics):
    """One worksheet per detected table, preserving cell merges and textual values."""
    from openpyxl import Workbook
    from openpyxl.styles import Font,PatternFill,Alignment,Border,Side
    from openpyxl.utils import get_column_letter
    wb=Workbook();index=wb.active;index.title='Index'
    index.append(['Page','Tableau','Feuille','Lignes','Colonnes'])
    count=0
    for page_no,meta in enumerate(diagnostics,1):
        for n,t in enumerate(meta.get('table_data',[]),1):
            rows,cols=t.get('rows',0),t.get('cols',0)
            if not (0<rows<=250 and 0<cols<=80):continue
            count+=1;name=f'P{page_no}_Table{n}'[:31];ws=wb.create_sheet(name)
            index.append([page_no,n,name,rows,cols]);ws.freeze_panes='A2'
            for c in t.get('cells',[]):
                r,k=c['row']+1,c['col']+1
                if r>rows or k>cols:continue
                cell=ws.cell(r,k,c.get('text',''))
                cell.alignment=Alignment(wrap_text=True,vertical='center')
                cell.border=Border(bottom=Side(style='hair',color='CBD5E1'),right=Side(style='hair',color='CBD5E1'))
                if c.get('kind') in ('columnHeader','rowHeader'):
                    cell.font=Font(bold=True,color='FFFFFF');cell.fill=PatternFill('solid',fgColor='1E3A5F')
                rr=min(rows,r+max(1,c.get('rowspan',1))-1);cc=min(cols,k+max(1,c.get('colspan',1))-1)
                if rr>r or cc>k:
                    try:ws.merge_cells(start_row=r,start_column=k,end_row=rr,end_column=cc)
                    except ValueError:pass
            for col in range(1,cols+1):ws.column_dimensions[get_column_letter(col)].width=22
            for row in range(1,rows+1):ws.row_dimensions[row].height=29
    index.column_dimensions['A'].width=12;index.column_dimensions['B'].width=14;index.column_dimensions['C'].width=22
    index.column_dimensions['D'].width=12;index.column_dimensions['E'].width=12
    if not count:index.append(['Aucun tableau détecté. Utiliser Azure Layout pour identifier les cellules.'])
    out=io.BytesIO();wb.save(out);return out.getvalue()