"""Local ONNX OCR with shared detection and line assembly for CPU/GPU."""
import base64
import io
import json
import time
from pathlib import Path

import numpy as np
import cv2
from PIL import Image, ImageOps
from tokenizers import Tokenizer

from .dbnet_utils import DBPostProcess, DetResizeForTest
from .line_merger import merge_line_boxes


def pixels(image, config):
    size=config['size']
    image=image.convert('RGB').resize((size['width'],size['height']),Image.Resampling(config['resample']))
    x=np.asarray(image,dtype=np.float32)*np.float32(config['rescale_factor'])
    x=(x-np.asarray(config['image_mean'],dtype=np.float32))/np.asarray(config['image_std'],dtype=np.float32)
    return np.ascontiguousarray(x.transpose(2,0,1)[None])


def split_line(image, ratio=20):
    w,h=image.size
    if w/max(h,1)<=ratio:return [image]
    ink=(np.asarray(image.convert('L'))<210).sum(axis=0)
    runs=[];start=None
    for x,value in enumerate(ink):
        if value==0 and start is None:start=x
        elif value!=0 and start is not None:runs.append((start,x));start=None
    if start is not None:runs.append((start,w))
    gaps=[(a,b) for a,b in runs if a>40 and b<w-40 and b-a>=8]
    if gaps:
        a,b=min(gaps,key=lambda r:abs((r[0]+r[1])/2-w/2));cut=(a+b)//2
    else:
        left,right=int(w*.35),int(w*.65)
        cut=left+int(np.argmin(ink[left:right])) if right>left else w//2
    cut=max(1,min(w-1,cut))
    return split_line(image.crop((0,0,cut,h)),ratio)+split_line(image.crop((cut,0,w,h)),ratio)


def preview(image):
    image=image.copy();image.thumbnail((1100,180))
    out=io.BytesIO();image.save(out,format='PNG')
    return 'data:image/png;base64,'+base64.b64encode(out.getvalue()).decode('ascii')


class OCR:
    def __init__(self, root=None, backend='cpu', threads=4):
        import onnxruntime as ort
        ort.disable_telemetry_events()
        self.root=Path(root or Path(__file__).resolve().parents[1])
        self.config=json.loads((self.root/'config.json').read_text(encoding='utf-8'))
        self.preprocess=json.loads((self.root/'preprocessor_config.json').read_text(encoding='utf-8'))
        self.tokenizer=Tokenizer.from_file(str(self.root/'tokenizer.json'))
        self.tokenizer.no_padding();self.tokenizer.no_truncation()
        requested_backend=backend
        if backend not in ('cpu','gpu','auto'):raise ValueError('backend must be cpu, gpu or auto')
        if backend in ('gpu','auto') and hasattr(ort,'preload_dlls'):
            ort.preload_dlls()
        if backend=='auto':backend='gpu' if 'CUDAExecutionProvider' in ort.get_available_providers() else 'cpu'
        if backend=='gpu' and 'CUDAExecutionProvider' not in ort.get_available_providers():
            raise RuntimeError('CUDA недоступна. Установите GPU-зависимости или выберите CPU.')
        options=ort.SessionOptions();options.intra_op_num_threads=threads;options.inter_op_num_threads=1
        if backend=='gpu':
            variant=self.config['variants']['fp16']
            providers=['CUDAExecutionProvider','CPUExecutionProvider']
            try:
                self.encoder=ort.InferenceSession(str(self.root/variant['encoder']),sess_options=options,providers=providers)
                self.decoder=ort.InferenceSession(str(self.root/variant['decoder']),sess_options=options,providers=providers)
                if any('CUDAExecutionProvider' not in s.get_providers() for s in (self.encoder,self.decoder)):
                    raise RuntimeError('Модель не загрузилась на CUDA. Выберите CPU; скрытая смена GPU на CPU запрещена.')
                self.encoder.disable_fallback()
                self.decoder.disable_fallback()
            except Exception:
                if requested_backend!='auto':raise
                backend='cpu'
        if backend=='cpu':
            variant=self.config['variants']['int8']
            providers=['CPUExecutionProvider']
            self.encoder=ort.InferenceSession(str(self.root/variant['encoder']),sess_options=options,providers=providers)
            self.decoder=ort.InferenceSession(str(self.root/variant['decoder']),sess_options=options,providers=providers)
        self.backend=backend
        self.detector=ort.InferenceSession(str(self.root/'models/dbnet.onnx'),sess_options=options,providers=['CPUExecutionProvider'])
        self.resize=DetResizeForTest(limit_side_len=736,limit_type='min')
        self.post=DBPostProcess(thresh=.2,box_thresh=.2,max_candidates=1000,unclip_ratio=1.3,use_dilation=True,score_mode='fast')

    def recognize_line(self,image):
        hidden=self.encoder.run(None,{'pixel_values':pixels(image,self.preprocess)})[0]
        ids=np.array([[self.config['decoder_start_token_id']]],dtype=np.int64)
        truncated=True
        for _ in range(self.config['max_new_tokens']):
            logits=self.decoder.run(None,{'input_ids':ids,'encoder_hidden_states':hidden})[0]
            if not np.isfinite(logits).all():raise RuntimeError('Non-finite decoder output')
            token=int(np.argmax(logits[0,-1]))
            ids=np.concatenate((ids,np.array([[token]],dtype=np.int64)),axis=1)
            if token==self.config['eos_token_id']:truncated=False;break
        return self.tokenizer.decode(ids[0].tolist(),skip_special_tokens=True).strip(),truncated

    def detect(self,image):
        # The original detector upsizes the short side. Bound unusually narrow
        # images so a one-line crop cannot allocate a gigantic detector tensor.
        scale=max(1.,736/min(image.size))
        if max(image.size)*scale>4096:
            scale=4096/max(image.size)
            width=max(32,round(image.width*scale/32)*32)
            height=max(32,round(image.height*scale/32)*32)
            data=dict(image=cv2.resize(np.asarray(image),(width,height)),
                      shape=np.array([image.height,image.width,height/image.height,width/image.width]))
        else:data=self.resize({'image':np.asarray(image)})
        x=np.asarray(data['image'],dtype=np.float32)*np.float32(1/255)
        x=(x-np.array([.485,.456,.406],dtype=np.float32))/np.array([.229,.224,.225],dtype=np.float32)
        x=np.ascontiguousarray(x.transpose(2,0,1)[None])
        output=self.detector.run(None,{self.detector.get_inputs()[0].name:x})[0]
        boxes=self.post(output,np.asarray([data['shape']]))[0]['points']
        valid=[]
        for box in boxes:
            box=np.asarray(box,dtype=np.float32)
            box[:,0]=np.clip(box[:,0],0,image.width-1);box[:,1]=np.clip(box[:,1],0,image.height-1)
            if np.linalg.norm(box[0]-box[1])>3 and np.linalg.norm(box[0]-box[3])>3:valid.append(box)
        return merge_line_boxes(valid,x_margin_ratio=3.5,y_overlap_thresh=.45)

    def recognize(self,image,boxes=None,include_images=True):
        t=time.perf_counter()
        image=ImageOps.exif_transpose(image).convert('RGB')
        if np.asarray(image.convert('L')).mean()<125:image=ImageOps.invert(image)
        boxes=self.detect(image) if boxes is None else boxes
        lines=[]
        for box in boxes:
            a,b,c,d=[int(v) for v in box]
            if c-a<4 or d-b<4:continue
            crop=image.crop((max(0,a-4),max(0,b-2),min(image.width,c+4),min(image.height,d+2)))
            parts=[self.recognize_line(ch) for ch in split_line(crop)]
            text=' '.join(s for s,_ in parts if s).strip()
            line=dict(bbox=[a,b,c,d],text=text,truncated=any(end for _,end in parts),chunks=len(parts))
            if include_images:line['image']=preview(crop)
            lines.append(line)
        return dict(backend=self.backend,lines=lines,full_text='\n'.join(row['text'] for row in lines),
                    seconds=time.perf_counter()-t,no_text_detected=not lines)