import os, glob
import fitz 
# import easyocr
import pymupdf
import json

from typing import Dict, List, Optional, Any 
from src.utils import initialize_a_logger


class DocumentIO:
    def __init__(self):
        self.logger = initialize_a_logger("./logs/DocumentIO.log")
        self.input_filename = ""
        self.output_filename = ""
        self.cur_page_number = None
        self.cur_page_content = None
        self.content = {}


    async def upload_document(self, input_filename:str, ext="json"):
        self.input_filename = './' + input_filename
        self.logger.info(f"input filename {self.input_filename}")
        
        self.output_filename = self.input_filename.replace("pdf",ext)
        self.logger.info(f"output filename { self.output_filename}")
        result = await self.load_json_file(self.output_filename)
        return result

    def create_document_with_easyocr(self, file_path):
        doc = fitz.open(file_path)
        reader = easyocr.Reader(['en'])  # Initialize once for efficiency
        results = {}
        
        for page_num in range(doc.page_count):
            page = doc[page_num]
                        
            # Convert page to image
            pix = page.get_pixmap(dpi=500)  # Higher DPI for better OCR
            img_data = pix.tobytes("png")
        
            ocr_results = reader.readtext(img_data)
                
            # Extract text and confidence scores
            page_text = ""
            word_details = []
            
            for (_ , text, confidence) in ocr_results:
                if confidence > 0.2:  # Filter low-confidence results
                    page_text += text + " "
                    word_details.append({
                        'text': text
                    })
            
            results[f'page_{page_num + 1}'] = {
                'text': page_text.strip(),
            }
            
        doc.close()
        return results


    async def load_json_file(self, output_filename: str) -> Dict[str, Any]:
        if not os.path.exists(output_filename):
            results: Dict[str, Any] = self.create_document_with_easyocr(self.input_filename)
            self.content = results
            
            with open(output_filename, "w", encoding="utf-8") as f:
                json.dump(results, f, indent=4, ensure_ascii=False)
        else:
            with open(output_filename, "r", encoding="utf-8") as f:
                try:
                    results = json.load(f)
                    self.content = results
                except json.JSONDecodeError:
                    if self.logger:
                        self.logger.debug("Error: load_json_file failed!")
                    self.content = {}
        return self.content

    def load_page(self, page_id: int) -> Dict[str, Any]:
        page_key = f"page_{page_id}"
        if not self.content:
            if self.logger:
                self.logger.debug("Content not loaded yet. Run load_json_file first.")
            return {}

        try:
            if len(self.content) == 0:
                self.content = self.load_json_file(self.output_filename)
                self.logger.info("Load json file in load page")

            page_data = self.content.get(page_key, {})
        except Exception as e:
            if self.logger:
                self.logger.debug(f"Error accessing page {page_id}: {e}")
            return {}

        return page_data
    
    def extract_header(self, pdf_path:str, page_num:int=2, header_height_ratio=0.1):
        """
        Extracts text from the top portion of a page (header area).
        """
        doc = fitz.open(pdf_path)
        page = doc[page_num]
        blocks = page.get_text("blocks")  # list of (x0, y0, x1, y1, text, block_no, block_type)

        page_height = page.rect.height
        header_cutoff = page_height * header_height_ratio

        header_text = []
        for b in blocks:
            x0, y0, x1, y1, text, *_ = b
            if y1 <= header_cutoff:  # only take text in the top area
                header_text.append(text.strip())

        return "\n".join(header_text)


    def extract_footer(self, pdf_path, page_num=2, footer_height_ratio=0.1):
        """
        Extracts text from the bottom portion of a page (footer area).
        """
        doc = fitz.open(pdf_path)
        page = doc[page_num]
        blocks = page.get_text("blocks")  # (x0, y0, x1, y1, text, block_no, block_type)

        page_height = page.rect.height
        footer_cutoff = page_height * (1 - footer_height_ratio)

        footer_text = []
        for b in blocks:
            x0, y0, x1, y1, text, *_ = b
            if y0 >= footer_cutoff:  # only take text in the bottom area
                footer_text.append(text.strip())

        return "\n".join(footer_text)

    def create_document_with_pymupdf(self, filepath):
        doc_map = {}
        doc = pymupdf.open(filepath)
        
        for i,page in enumerate(doc): # iterate the document pages
            header = self.extract_header(filepath, page_num=i)
            footer = self.extract_footer(filepath, page_num=i)
            text = page.get_text()
            text = text.replace(header, "")
            text = text.replace(footer, "")
            doc_map[f"page {i}"]= text
        return doc_map