from __future__ import annotations

import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import cv2
import numpy as np
import pytesseract
from PIL import Image

from ocr_pipeline.config import settings
from ocr_pipeline.utils.logging import get_logger

logger = get_logger(__name__)

# Language code mapping: config names -> engine-specific codes
TESSERACT_LANG_MAP = {
    "en": "eng",
    "latin": "lat",
    "de": "deu",
    "fr": "fra",
    "es": "spa",
    "it": "ita",
    "pt": "por",
    "zh": "chi_sim",
    "ja": "jpn",
    "ko": "kor",
}

PADDLEOCR_LANG_MAP = {
    "en": "en",
    "latin": "latin",
    "de": "german",
    "fr": "french",
    "es": "spanish",
    "it": "italian",
    "pt": "portuguese",
    "zh": "ch",
    "ja": "japan",
    "ko": "korean",
}


def map_languages(languages: list[str], engine: str) -> str:
    """Map config language names to engine-specific codes."""
    if engine == "tesseract":
        mapped = [TESSERACT_LANG_MAP.get(lang, lang) for lang in languages]
        return "+".join(mapped)
    else:  # paddleocr
        # PaddleOCR uses primary language only for detection
        primary = languages[0] if languages else "en"
        return PADDLEOCR_LANG_MAP.get(primary, primary)


@dataclass
class OCRResult:
    text: str
    confidence: float
    language: str
    engine: str
    boxes: list[dict[str, Any]]
    raw_result: Any
    processing_time: float


class PaddleOCREngine:
    def __init__(self):
        self._ocr = None
        self._initialized = False

    def _init(self):
        if self._initialized:
            return
        try:
            from paddleocr import PaddleOCR

            paddle_lang = map_languages(settings.ocr.languages, "paddleocr")

            self._ocr = PaddleOCR(
                use_angle_cls=settings.ocr.use_angle_cls,
                lang=paddle_lang,
                use_gpu=settings.ocr.use_gpu,
                show_log=False,
                det_db_thresh=settings.ocr.det_db_thresh,
                det_db_box_thresh=settings.ocr.det_db_box_thresh,
                det_db_unclip_ratio=settings.ocr.det_db_unclip_ratio,
                rec_batch_num=settings.ocr.rec_batch_num,
                cpu_threads=settings.ocr.cpu_threads,
            )
            self._initialized = True
            logger.info("paddleocr_initialized", language=paddle_lang)
        except Exception as e:
            logger.error("paddleocr_init_failed", error=str(e))
            raise

    def process(self, image: np.ndarray) -> OCRResult:
        self._init()
        start = time.time()

        result = self._ocr.ocr(image, cls=True)
        processing_time = time.time() - start

        if not result or not result[0]:
            return OCRResult(
                text="",
                confidence=0.0,
                language=map_languages(settings.ocr.languages, "paddleocr"),
                engine="paddleocr",
                boxes=[],
                raw_result=result,
                processing_time=processing_time,
            )

        texts = []
        confidences = []
        boxes = []

        for line in result[0]:
            box = line[0]
            text = line[1][0]
            conf = line[1][1]

            texts.append(text)
            confidences.append(conf)
            boxes.append(
                {
                    "bbox": box,
                    "text": text,
                    "confidence": conf,
                }
            )

        full_text = "\n".join(texts)
        mean_conf = sum(confidences) / len(confidences) if confidences else 0.0

        return OCRResult(
            text=full_text,
            confidence=mean_conf,
            language=map_languages(settings.ocr.languages, "paddleocr"),
            engine="paddleocr",
            boxes=boxes,
            raw_result=result,
            processing_time=processing_time,
        )


class TesseractEngine:
    def process(self, image: np.ndarray) -> OCRResult:
        start = time.time()

        pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))

        tess_lang = map_languages(settings.ocr.languages, "tesseract")
        config = f"--oem 3 --psm 6 -l {tess_lang}"
        data = pytesseract.image_to_data(
            pil_image, config=config, output_type=pytesseract.Output.DICT
        )

        texts = []
        confidences = []
        boxes = []

        for i in range(len(data["text"])):
            text = data["text"][i].strip()
            conf = data["conf"][i]
            if text and conf > 0:
                texts.append(text)
                confidences.append(conf / 100.0)
                boxes.append(
                    {
                        "bbox": [
                            [data["left"][i], data["top"][i]],
                            [data["left"][i] + data["width"][i], data["top"][i]],
                            [
                                data["left"][i] + data["width"][i],
                                data["top"][i] + data["height"][i],
                            ],
                            [data["left"][i], data["top"][i] + data["height"][i]],
                        ],
                        "text": text,
                        "confidence": conf / 100.0,
                    }
                )

        full_text = " ".join(texts)
        mean_conf = sum(confidences) / len(confidences) if confidences else 0.0

        return OCRResult(
            text=full_text,
            confidence=mean_conf,
            language=map_languages(settings.ocr.languages, "tesseract"),
            engine="tesseract",
            boxes=boxes,
            raw_result=data,
            processing_time=time.time() - start,
        )


class OCREngine:
    def __init__(self):
        self.paddle = PaddleOCREngine()
        self.tesseract = TesseractEngine()
        self.engine_preference = settings.ocr.engine

    def process(self, image: np.ndarray) -> OCRResult:
        """Process a pre-loaded image array."""
        if self.engine_preference == "paddleocr":
            try:
                return self.paddle.process(image)
            except Exception as e:
                logger.warning("paddleocr_failed_fallback", error=str(e))
                return self.tesseract.process(image)

        elif self.engine_preference == "tesseract":
            return self.tesseract.process(image)

        else:
            try:
                return self.paddle.process(image)
            except Exception as e:
                logger.warning("paddleocr_failed_fallback", error=str(e))
                return self.tesseract.process(image)

    def process_file(self, image_path: Path) -> OCRResult:
        """Load and process an image file."""
        logger.info("processing_image", path=str(image_path), engine=self.engine_preference)

        image = cv2.imread(str(image_path))
        if image is None:
            raise ValueError(f"Failed to load image: {image_path}")
        return self.process(image)

    # Alias for backward compatibility
    recognize = process
