import cv2
import numpy as np
import os
import urllib.request

class DataScanner:
    """
    Module: Text Engine (ONNX Powered OCR)
    Reads physical documents and digitizes them into Aelion OS.
    """
    def __init__(self):
        self.current_dir = os.path.dirname(os.path.abspath(__file__))
        
        # File Paths
        self.det_path = os.path.join(self.current_dir, "text_detection_DB.onnx")
        self.rec_path = os.path.join(self.current_dir, "text_recognition_CRNN.onnx")

        # --- THE AUTO-DOWNLOADER ---
        self._ensure_models_exist()

        print("[Aelion Vision] Booting Optical Data Scanner...")
        self._init_ai_engines()

    def _ensure_models_exist(self):
        """Downloads the ONNX files exclusively from the Aelion Private Server."""
        files_to_get = {
            self.det_path: "https://ai.aelionyxe.top/models/text_detection_DB_IC15_resnet18_2021sep.onnx",
            self.rec_path: "https://ai.aelionyxe.top/models/text_recognition_CRNN_EN_2021sep.onnx"
        }

        for path, url in files_to_get.items():
            if not os.path.exists(path):
                print(f"[Aelion Vision] Fetching {os.path.basename(path)} from Aelion Server...")
                try:
                    urllib.request.urlretrieve(url, path)
                    print(f"[Aelion Vision] Successfully secured {os.path.basename(path)}.")
                except Exception as e:
                    print(f"[!] Critical Network Error: Failed to reach {url}. {e}")

    def _init_ai_engines(self):
        """Loads the ONNX weights into memory."""
        try:
            # We hardcode the vocabulary into RAM to prevent 404 download errors!
            self.vocabulary = list("0123456789abcdefghijklmnopqrstuvwxyz")

            # 1. Setup the Detector (Finds the boxes)
            self.detector = cv2.dnn.TextDetectionModel_DB(self.det_path)
            self.detector.setBinaryThreshold(0.3).setPolygonThreshold(0.5).setUnclipRatio(2.0).setMaxCandidates(1000)
            self.detector.setInputParams(1.0/255.0, (736, 736), (122.67891434, 116.66876762, 104.00698793))

            # 2. Setup the Recognizer (Reads the words)
            self.recognizer = cv2.dnn.TextRecognitionModel(self.rec_path)
            self.recognizer.setDecodeType("CTC-greedy")
            self.recognizer.setVocabulary(self.vocabulary)
            
            # THE CORE FIX: Set to 1-Channel (Grayscale) mean instead of 3-Channel (BGR)
            self.recognizer.setInputParams(1.0/127.5, (100, 32), (127.5,))
            
            self.is_ready = True
        except Exception as e:
            print(f"[!] Optical Engine failed to start: {e}")
            self.is_ready = False

    def scan_frame(self, frame):
        """
        NEW UPGRADE: Scans a live camera frame directly in RAM!
        Advanced: Converts to grayscale to prevent crashes, and sorts words top-to-bottom.
        """
        if not self.is_ready or frame is None: return []

        # Find text boxes
        boxes, _ = self.detector.detect(frame)
        if len(boxes) == 0: return []

        extracted_text = []
        img_h, img_w = frame.shape[:2]

        # --- ADVANCED UPGRADE: Smart Reading Order ---
        # Extract standard bounding rectangles first
        rectangles = []
        for box in boxes:
            pts = np.array(box).astype(np.int32)
            x, y, w, h = cv2.boundingRect(pts)
            rectangles.append((x, y, w, h))
            
        # Sort the boxes so the AI reads top-to-bottom, left-to-right
        # Grouping by lines roughly every 15 pixels of Y-height
        rectangles = sorted(rectangles, key=lambda r: (r[1] // 15, r[0]))

        # Read the text in each sorted box
        for (x, y, w, h) in rectangles:
            try:
                # Prevent out-of-bounds slicing which causes ghost crashes
                x, y = max(0, x), max(0, y)
                w, h = min(img_w - x, w), min(img_h - y, h)
                
                if w < 5 or h < 5: continue
                
                # Cut the word cleanly out of the main image
                cropped_word = frame[y:y+h, x:x+w]
                if cropped_word.size == 0: continue

                # THE CORE FIX: Convert to Grayscale! The neural net demands 1 channel.
                gray_word = cv2.cvtColor(cropped_word, cv2.COLOR_BGR2GRAY)

                # Feed the Grayscale cut to the AI
                word = self.recognizer.recognize(gray_word)
                
                if word:
                    extracted_text.append(word)
                    
            except Exception as e:
                print(f"[!] Recognition error on word: {e}")
                continue 

        return extracted_text

    def scan_image_to_file(self, image_path, output_filename="scanned_data.txt"):
        """Reads an image from your hard drive and saves the text to a .txt file."""
        if not self.is_ready: return False

        img = cv2.imread(image_path)
        if img is None:
            print(f"[Aelion Vision] Error: Could not load image at {image_path}")
            return False

        print(f"[Aelion Vision] Scanning image layout...")
        words = self.scan_frame(img)

        print(f"[Aelion Vision] Compiling {len(words)} words...")
        with open(output_filename, "a", encoding="utf-8") as file:
            full_text = " ".join(words)
            file.write(full_text + "\n")
            
        print(f"[*] DATA SAVED TO: {output_filename}")
        return True