import pytesseract
from pdf2image import convert_from_path
from PIL import Image
from pathlib import Path
from typing import Optional
import tempfile

from config import settings
from logger import get_logger

log = get_logger("ocr")


class OCRExtractor:
    def __init__(self, lang: str = None):
        self.lang = lang or settings.TESSERACT_LANG
        log.debug(f"OCRExtractor initialized: lang={self.lang}")
    
    def extract_from_image(self, image_path: Path | str) -> str:
        log.debug(f"Extracting from image: {image_path}")
        image = Image.open(image_path)
        log.debug(f"Image size: {image.size}")
        
        text = pytesseract.image_to_string(image, lang=self.lang)
        log.info(f"Image OCR: {len(text)} chars extracted")
        
        return text.strip()
    
    def extract_from_pdf(self, pdf_path: Path | str) -> list[dict]:
        log.debug(f"Converting PDF to images: {pdf_path}")
        
        pages = convert_from_path(pdf_path)
        log.debug(f"PDF has {len(pages)} pages")
        
        results = []
        
        for i, page in enumerate(pages, 1):
            log.debug(f"Processing page {i}/{len(pages)}")
            text = pytesseract.image_to_string(page, lang=self.lang)
            if text.strip():
                results.append({
                    "page": i,
                    "text": text.strip()
                })
                log.debug(f"Page {i}: {len(text)} chars")
            else:
                log.debug(f"Page {i}: empty (skipped)")
        
        log.info(f"PDF OCR complete: {len(results)} pages with text")
        return results
    
    def extract(self, file_path: Path | str) -> dict:
        path = Path(file_path)
        log.info(f"Extracting: {path.name}")
        
        if not path.exists():
            log.error(f"File not found: {path}")
            raise FileNotFoundError(f"File not found: {path}")
        
        suffix = path.suffix.lower()
        log.debug(f"File type: {suffix}")
        
        if suffix == ".pdf":
            pages = self.extract_from_pdf(path)
            full_text = "\n\n".join(p["text"] for p in pages)
            result = {
                "source": str(path),
                "type": "pdf",
                "pages": pages,
                "full_text": full_text,
                "total_pages": len(pages)
            }
            log.info(f"PDF extracted: {len(full_text)} chars, {len(pages)} pages")
            return result
        
        elif suffix in {".png", ".jpg", ".jpeg", ".tiff", ".bmp"}:
            text = self.extract_from_image(path)
            result = {
                "source": str(path),
                "type": "image",
                "full_text": text,
                "total_pages": 1
            }
            log.info(f"Image extracted: {len(text)} chars")
            return result
        
        else:
            log.error(f"Unsupported file type: {suffix}")
            raise ValueError(f"Unsupported file type: {suffix}")


def extract_text(file_path: Path | str, lang: str = None) -> dict:
    extractor = OCRExtractor(lang=lang)
    return extractor.extract(file_path)