# Copyright 2025 Thulium Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""HTR Pipeline - Core handwriting text recognition orchestration.

This module provides the main HTRPipeline class that orchestrates the
complete workflow from raw image input to structured text output,
integrating preprocessing, segmentation, recognition, and postprocessing.

Classes:
    HTRPipeline: Main pipeline for end-to-end handwriting recognition.

Example:
    >>> from thulium.pipeline.htr_pipeline import HTRPipeline
    >>> pipeline = HTRPipeline(device="cuda", language="en")
    >>> result = pipeline.process("document.jpg")
    >>> print(result.full_text)
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Any
from typing import Dict
from typing import List
from typing import Optional
from typing import Tuple
from typing import Union

import torch
import yaml
from PIL import Image

from thulium.api.types import Line
from thulium.api.types import PageResult
from thulium.data.language_profiles import get_language_profile
from thulium.data.language_profiles import LanguageProfile
from thulium.data.loaders import load_image
from thulium.models.segmentation.line_segmentation import LineSegmenter
from thulium.models.wrappers.htr_model import HTRModel

logger = logging.getLogger(__name__)


class HTRPipeline:
    """Orchestrates the handwritten text recognition process.

    This class binds together the segmentation, recognition, and decoding
    components into a cohesive pipeline. It handles device management,
    configuration loading, and the execution flow for both single images
    and batches.

    Attributes:
        device: Torch device (CPU or CUDA).
        segmenter: Line segmentation model.
        recognizer: HTR recognition model.
    """

    def __init__(
        self,
        config_path: Optional[str] = None,
        device: str = "auto",
        language: str = "en",
    ) -> None:
        """Initialize the pipeline.

        Args:
            config_path: Path to YAML configuration file.
            device: 'auto', 'cpu', or 'cuda'.
            language: Default language code.
        """
        self.device = self._resolve_device(device)
        self.config = self._load_config(config_path) if config_path else {}
        self.language_profile = get_language_profile(language)

        logger.info("Initializing HTRPipeline on %s for %s", self.device, language)

        # Initialize Models
        # In a real scenario, we would load weights here.
        self.segmenter = LineSegmenter(in_channels=3).to(self.device)
        self.recognizer = HTRModel(
            num_classes=self.language_profile.get_vocab_size()
        ).to(self.device)

        self.segmenter.eval()
        self.recognizer.eval()

    def _resolve_device(self, device_str: str) -> torch.device:
        """Resolve device string to torch.device."""
        if device_str == "auto":
            return torch.device("cuda" if torch.cuda.is_available() else "cpu")
        return torch.device(device_str)

    def _load_config(self, path: str) -> Dict[str, Any]:
        """Load configuration from YAML."""
        try:
            with open(path, "r", encoding="utf-8") as f:
                return yaml.safe_load(f)
        except Exception as e:
            logger.warning("Failed to load config %s: %s. Using API defaults.", path, e)
            return {}

    def process(self, image_source: Union[str, Path, Image.Image]) -> PageResult:
        """Process a single document image.

        Args:
            image_source: Path to image or PIL Image object.

        Returns:
            PageResult containing recognized text and layout info.
        """
        # 1. Load Image
        if isinstance(image_source, (str, Path)):
            image = load_image(image_source)
        else:
            image = image_source

        # 2. Preprocess & Segment
        # Convert to tensor
        import torchvision.transforms.functional as TF
        img_tensor = TF.to_tensor(image).unsqueeze(0).to(self.device)

        # Run segmenter
        with torch.no_grad():
            seg_map = self.segmenter(img_tensor)
            # Post-process segmentation map to get bounding boxes
            # This is a complex step usually involving connected components
            # For this 'from scratch' impl, we'll simulate output or use simple thresholding
            line_bboxes = self._extract_bboxes_from_mask(seg_map, image.size)

        # 3. Recognize Lines
        recognized_lines = []
        for bbox in line_bboxes:
            # Crop line
            x, y, w, h = bbox
            line_img = image.crop((x, y, x + w, y + h))
            
            # Recognize
            text, conf = self._recognize_line(line_img)
            
            recognized_lines.append(
                Line(text=text, confidence=conf, bbox=bbox)
            )

        # 4. Construct Result
        full_text = "\n".join([l.text for l in recognized_lines])
        
        return PageResult(
            full_text=full_text,
            lines=recognized_lines,
            language=self.language_profile.code,
            metadata={"device": str(self.device)}
        )

    def _extract_bboxes_from_mask(self, seg_map: torch.Tensor, original_size: Tuple[int, int]) -> List[Tuple[int, int, int, int]]:
        """Extract bounding boxes from segmentation probability map.
        
        Args:
            seg_map: (1, 1, H, W) tensor.
            original_size: (W, H) tuple.
            
        Returns:
             List of (x, y, w, h) tuples.
        """
        # Real implementation would use cv2.findContours or kornia
        # Here we provide a heuristic fallback if no segments found (e.g. init weights)
        # to ensure pipeline runs end-to-end.
        
        # In a real trained model, seg_map > 0.5 would give binary mask.
        prob = torch.sigmoid(seg_map)
        if prob.max() < 0.1:
             # Fallback: treat whole image as one line if model is untrained
             return [(0, 0, original_size[0], original_size[1])]
        
        # Placeholder for connected components logic
        # Current logic: One big box
        return [(0, 0, original_size[0], original_size[1])]

    def _recognize_line(self, line_img: Image.Image) -> Tuple[str, float]:
        """Run recognition on a single line image."""
        # Preprocess
        import torchvision.transforms.functional as TF
        # Resize height to fixed size (e.g. 32 or 64)
        target_height = 64
        w, h = line_img.size
        new_w = int(w * (target_height / h))
        line_img = line_img.resize((new_w, target_height), Image.Resampling.BILINEAR)
        
        tensor = TF.to_tensor(line_img).unsqueeze(0).to(self.device)
        
        with torch.no_grad():
            log_probs = self.recognizer(tensor) # (T, B, C)
            
            # Decode
            # This calls the greedy decoder inside the recognizer manually or via helper
            # recognizer.forward returns log_probs. We need decoder.decode_greedy
            preds = self.recognizer.decoder.decode_greedy(log_probs.transpose(0, 1))
            
            # TODO: Convert indices to text using Vocabulary
            # For now, return string representation of indices
            pred_indices = preds[0]
            # text = "".join([self.vocab[i] for i in pred_indices])
            text = f"Recognized: {pred_indices}" 
            confidence = 0.9 # Placeholder
            
        return text, confidence
