"""
Database-backed YOLO Dataset

Loads images from filesystem and labels from PostgreSQL database.
Generates YOLO format labels dynamically at training time.
"""

from __future__ import annotations

import logging
import random
from dataclasses import dataclass, field
from functools import lru_cache
from pathlib import Path
from typing import Any, Optional

import numpy as np
from PIL import Image

from shared.config import DEFAULT_DPI
from shared.fields import TRAINING_FIELD_CLASSES as FIELD_CLASSES, CLASS_NAMES
from shared.bbox import expand_bbox
from .annotation_generator import YOLOAnnotation

logger = logging.getLogger(__name__)

# Module-level LRU cache for image loading (shared across dataset instances)
@lru_cache(maxsize=256)
def _load_image_cached(image_path: str) -> tuple[np.ndarray, int, int]:
    """
    Load and cache image from disk.

    Args:
        image_path: Path to image file (must be string for hashability)

    Returns:
        Tuple of (image_array, width, height)
    """
    image = Image.open(image_path).convert('RGB')
    width, height = image.size
    image_array = np.array(image)
    return image_array, width, height


def clear_image_cache():
    """Clear the image cache to free memory."""
    _load_image_cached.cache_clear()


@dataclass
class DatasetItem:
    """Single item in the dataset."""
    document_id: str
    image_path: Path
    page_no: int
    labels: list[YOLOAnnotation]
    is_scanned: bool = False  # True if bbox is in pixel coords, False if in PDF points
    csv_split: str | None = None  # CSV-defined split ('train', 'test', etc.)


class DBYOLODataset:
    """
    YOLO Dataset that reads labels from database.

    This dataset:
    1. Scans temp directory for rendered images
    2. Queries database for bbox data
    3. Generates YOLO labels dynamically
    4. Performs train/val/test split at runtime
    """

    def __init__(
        self,
        images_dir: str | Path,
        db: Any,  # DocumentDB instance
        split: str = 'train',
        train_ratio: float = 0.8,
        val_ratio: float = 0.1,
        seed: int = 42,
        dpi: int = DEFAULT_DPI,  # Must match the DPI used in autolabel_tasks.py for rendering
        min_confidence: float = 0.7,
        bbox_padding_px: int = 20,
        min_bbox_height_px: int = 30,
        limit: int | None = None,
    ):
        """
        Initialize database-backed YOLO dataset.

        Args:
            images_dir: Directory containing temp/{doc_id}/images/*.png
            db: DocumentDB instance for label queries
            split: Which split to use ('train', 'val', 'test')
            train_ratio: Ratio for training set
            val_ratio: Ratio for validation set
            seed: Random seed for reproducible splits
            dpi: DPI used for rendering (for coordinate conversion)
            min_confidence: Minimum match score to include
            bbox_padding_px: Padding around bboxes
            min_bbox_height_px: Minimum bbox height
            limit: Maximum number of documents to use (None for all)
        """
        self.images_dir = Path(images_dir)
        self.db = db
        self.split = split
        self.train_ratio = train_ratio
        self.val_ratio = val_ratio
        self.seed = seed
        self.dpi = dpi
        self.min_confidence = min_confidence
        self.bbox_padding_px = bbox_padding_px
        self.min_bbox_height_px = min_bbox_height_px
        self.limit = limit

        # Load and split dataset
        self.items: list[DatasetItem] = []
        self._all_items: list[DatasetItem] = []  # Cache all items for sharing
        self._doc_ids_ordered: list[str] = []  # Cache ordered doc IDs for consistent splits
        self._load_dataset()

    @classmethod
    def from_shared_data(
        cls,
        source_dataset: 'DBYOLODataset',
        split: str,
    ) -> 'DBYOLODataset':
        """
        Create a new dataset instance sharing data from an existing one.

        This avoids re-loading data from filesystem and database.

        Args:
            source_dataset: Dataset to share data from
            split: Which split to use ('train', 'val', 'test')

        Returns:
            New dataset instance with shared data
        """
        # Create instance without loading (we'll share data)
        instance = object.__new__(cls)

        # Copy configuration
        instance.images_dir = source_dataset.images_dir
        instance.db = source_dataset.db
        instance.split = split
        instance.train_ratio = source_dataset.train_ratio
        instance.val_ratio = source_dataset.val_ratio
        instance.seed = source_dataset.seed
        instance.dpi = source_dataset.dpi
        instance.min_confidence = source_dataset.min_confidence
        instance.bbox_padding_px = source_dataset.bbox_padding_px
        instance.min_bbox_height_px = source_dataset.min_bbox_height_px
        instance.limit = source_dataset.limit

        # Share loaded data
        instance._all_items = source_dataset._all_items
        instance._doc_ids_ordered = source_dataset._doc_ids_ordered

        # Split items for this split
        instance.items = instance._split_dataset_from_cache()
        logger.info("Split '%s': %d items", split, len(instance.items))

        return instance

    def _load_dataset(self):
        """Load dataset items from filesystem and database."""
        # Find all document directories
        temp_dir = self.images_dir / 'temp'
        if not temp_dir.exists():
            logger.warning("Temp directory not found: %s", temp_dir)
            return

        # Collect all document IDs with images
        doc_image_map: dict[str, list[Path]] = {}
        for doc_dir in temp_dir.iterdir():
            if not doc_dir.is_dir():
                continue

            images_path = doc_dir / 'images'
            if not images_path.exists():
                continue

            images = list(images_path.glob('*.png'))
            if images:
                doc_image_map[doc_dir.name] = sorted(images)

        logger.info("Found %d documents with images", len(doc_image_map))

        # Query database for all document labels
        doc_ids = list(doc_image_map.keys())
        doc_labels = self._load_labels_from_db(doc_ids)

        logger.info("Loaded labels for %d documents from database", len(doc_labels))

        # Build dataset items
        all_items: list[DatasetItem] = []
        skipped_no_labels = 0
        skipped_no_db_record = 0
        total_images = 0

        for doc_id, images in doc_image_map.items():
            doc_data = doc_labels.get(doc_id)

            # Skip documents that don't exist in database
            if doc_data is None:
                skipped_no_db_record += len(images)
                total_images += len(images)
                continue

            labels_by_page, is_scanned, csv_split = doc_data

            for image_path in images:
                total_images += 1
                # Extract page number from filename (e.g., "doc_page_000.png")
                page_no = self._extract_page_no(image_path.stem)

                # Get labels for this page
                page_labels = labels_by_page.get(page_no, [])

                if page_labels:  # Only include pages with labels
                    all_items.append(DatasetItem(
                        document_id=doc_id,
                        image_path=image_path,
                        page_no=page_no,
                        labels=page_labels,
                        is_scanned=is_scanned,
                        csv_split=csv_split
                    ))
                else:
                    skipped_no_labels += 1

        logger.info("Total images found: %d", total_images)
        logger.info("Images with labels: %d", len(all_items))
        if skipped_no_db_record > 0:
            logger.info("Skipped %d images (document not in database)", skipped_no_db_record)
        if skipped_no_labels > 0:
            logger.info("Skipped %d images (no labels for page)", skipped_no_labels)

        # Cache all items for sharing with other splits
        self._all_items = all_items

        # Split dataset
        self.items, self._doc_ids_ordered = self._split_dataset(all_items)
        logger.info("Split '%s': %d items", self.split, len(self.items))

    def _load_labels_from_db(self, doc_ids: list[str]) -> dict[str, tuple[dict[int, list[YOLOAnnotation]], bool, str | None]]:
        """
        Load labels from database for given document IDs using batch queries.

        Returns:
            Dict of doc_id -> (page_labels, is_scanned, split)
            where page_labels is {page_no -> list[YOLOAnnotation]}
            is_scanned indicates if bbox is in pixel coords (True) or PDF points (False)
            split is the CSV-defined split ('train', 'test', etc.) or None
        """
        result: dict[str, tuple[dict[int, list[YOLOAnnotation]], bool, str | None]] = {}

        # Query in batches using efficient batch method
        batch_size = 500
        for i in range(0, len(doc_ids), batch_size):
            batch_ids = doc_ids[i:i + batch_size]

            # Use batch query instead of individual queries (N+1 fix)
            docs_batch = self.db.get_documents_batch(batch_ids)

            for doc_id, doc in docs_batch.items():
                if not doc.get('success'):
                    continue

                # Check if scanned PDF (OCR bbox is in pixels, text PDF bbox is in PDF points)
                is_scanned = doc.get('pdf_type') == 'scanned'

                # Get CSV-defined split
                csv_split = doc.get('split')

                page_labels: dict[int, list[YOLOAnnotation]] = {}

                for field_result in doc.get('field_results', []):
                    if not field_result.get('matched'):
                        continue

                    field_name = field_result.get('field_name')

                    # Map supplier_accounts(X) to the actual class name (Bankgiro/Plusgiro)
                    yolo_class_name = field_name
                    if field_name and field_name.startswith('supplier_accounts('):
                        # Extract the account type: "supplier_accounts(Bankgiro)" -> "Bankgiro"
                        yolo_class_name = field_name.split('(')[1].rstrip(')')

                    if yolo_class_name not in FIELD_CLASSES:
                        continue

                    score = field_result.get('score', 0)
                    if score < self.min_confidence:
                        continue

                    bbox = field_result.get('bbox')
                    page_no = field_result.get('page_no', 0)

                    if bbox and len(bbox) == 4:
                        annotation = self._create_annotation(
                            field_name=yolo_class_name,  # Use mapped class name
                            bbox=bbox,
                            score=score
                        )

                        if page_no not in page_labels:
                            page_labels[page_no] = []
                        page_labels[page_no].append(annotation)

                if page_labels:
                    result[doc_id] = (page_labels, is_scanned, csv_split)

        return result

    def _create_annotation(
        self,
        field_name: str,
        bbox: list[float],
        score: float
    ) -> YOLOAnnotation:
        """
        Create a YOLO annotation from bbox.

        Note: bbox is in PDF points (72 DPI), will be normalized later.
        """
        class_id = FIELD_CLASSES[field_name]
        x0, y0, x1, y1 = bbox

        # Store raw PDF coordinates - will be normalized when getting item
        return YOLOAnnotation(
            class_id=class_id,
            x_center=(x0 + x1) / 2,  # center in PDF points
            y_center=(y0 + y1) / 2,
            width=x1 - x0,
            height=y1 - y0,
            confidence=score
        )

    def _extract_page_no(self, stem: str) -> int:
        """Extract page number from image filename."""
        # Format: "{doc_id}_page_{page_no:03d}"
        parts = stem.rsplit('_', 1)
        if len(parts) == 2:
            try:
                return int(parts[1])
            except ValueError:
                pass
        return 0

    def _split_dataset(self, items: list[DatasetItem]) -> tuple[list[DatasetItem], list[str]]:
        """
        Split items into train/val/test based on CSV-defined split field.

        If CSV has 'split' field, use it directly.
        Otherwise, fall back to random splitting based on train_ratio/val_ratio.

        Returns:
            Tuple of (split_items, ordered_doc_ids) where ordered_doc_ids can be
            reused for consistent splits across shared datasets.
        """
        # Group by document ID for proper splitting
        doc_items: dict[str, list[DatasetItem]] = {}
        doc_csv_split: dict[str, str | None] = {}  # Track CSV split per document

        for item in items:
            if item.document_id not in doc_items:
                doc_items[item.document_id] = []
                doc_csv_split[item.document_id] = item.csv_split
            doc_items[item.document_id].append(item)

        # Check if we have CSV-defined splits
        has_csv_splits = any(split is not None for split in doc_csv_split.values())

        doc_ids = list(doc_items.keys())

        if has_csv_splits:
            # Use CSV-defined splits
            logger.info("Using CSV-defined split field for train/val/test assignment")

            # Map split values: 'train' -> train, 'test' -> test, None -> train (fallback)
            # 'val' is taken from train set using val_ratio
            split_doc_ids = []

            if self.split == 'train':
                # Get documents marked as 'train' or no split defined
                train_docs = [doc_id for doc_id in doc_ids
                              if doc_csv_split[doc_id] in (None, 'train', 'Train', 'TRAIN')]

                # Take train_ratio of train docs for actual training, rest for val
                random.seed(self.seed)
                random.shuffle(train_docs)

                n_train = int(len(train_docs) * (self.train_ratio / (self.train_ratio + self.val_ratio)))
                split_doc_ids = train_docs[:n_train]

            elif self.split == 'val':
                # Get documents marked as 'train' and take val portion
                train_docs = [doc_id for doc_id in doc_ids
                              if doc_csv_split[doc_id] in (None, 'train', 'Train', 'TRAIN')]

                random.seed(self.seed)
                random.shuffle(train_docs)

                n_train = int(len(train_docs) * (self.train_ratio / (self.train_ratio + self.val_ratio)))
                split_doc_ids = train_docs[n_train:]

            else:  # test
                # Get documents marked as 'test'
                split_doc_ids = [doc_id for doc_id in doc_ids
                                 if doc_csv_split[doc_id] in ('test', 'Test', 'TEST')]

            # Apply limit if specified
            if self.limit is not None and self.limit < len(split_doc_ids):
                split_doc_ids = split_doc_ids[:self.limit]
                logger.info("Limited to %d documents", self.limit)

        else:
            # Fall back to random splitting (original behavior)
            logger.info("No CSV split field found, using random splitting")

            random.seed(self.seed)
            random.shuffle(doc_ids)

            # Apply limit if specified (before splitting)
            if self.limit is not None and self.limit < len(doc_ids):
                doc_ids = doc_ids[:self.limit]
                logger.info("Limited to %d documents", self.limit)

            # Calculate split indices
            n_total = len(doc_ids)
            n_train = int(n_total * self.train_ratio)
            n_val = int(n_total * self.val_ratio)

            # Split document IDs
            if self.split == 'train':
                split_doc_ids = doc_ids[:n_train]
            elif self.split == 'val':
                split_doc_ids = doc_ids[n_train:n_train + n_val]
            else:  # test
                split_doc_ids = doc_ids[n_train + n_val:]

        # Collect items for this split
        split_items = []
        for doc_id in split_doc_ids:
            split_items.extend(doc_items[doc_id])

        return split_items, doc_ids

    def _split_dataset_from_cache(self) -> list[DatasetItem]:
        """
        Split items using cached data from a shared dataset.

        Uses pre-computed doc_ids order for consistent splits.
        Respects CSV-defined splits if available.
        """
        # Group by document ID and track CSV splits
        doc_items: dict[str, list[DatasetItem]] = {}
        doc_csv_split: dict[str, str | None] = {}

        for item in self._all_items:
            if item.document_id not in doc_items:
                doc_items[item.document_id] = []
                doc_csv_split[item.document_id] = item.csv_split
            doc_items[item.document_id].append(item)

        # Check if we have CSV-defined splits
        has_csv_splits = any(split is not None for split in doc_csv_split.values())

        doc_ids = self._doc_ids_ordered

        if has_csv_splits:
            # Use CSV-defined splits
            if self.split == 'train':
                train_docs = [doc_id for doc_id in doc_ids
                              if doc_id in doc_csv_split and
                              doc_csv_split[doc_id] in (None, 'train', 'Train', 'TRAIN')]

                random.seed(self.seed)
                random.shuffle(train_docs)

                n_train = int(len(train_docs) * (self.train_ratio / (self.train_ratio + self.val_ratio)))
                split_doc_ids = train_docs[:n_train]

            elif self.split == 'val':
                train_docs = [doc_id for doc_id in doc_ids
                              if doc_id in doc_csv_split and
                              doc_csv_split[doc_id] in (None, 'train', 'Train', 'TRAIN')]

                random.seed(self.seed)
                random.shuffle(train_docs)

                n_train = int(len(train_docs) * (self.train_ratio / (self.train_ratio + self.val_ratio)))
                split_doc_ids = train_docs[n_train:]

            else:  # test
                split_doc_ids = [doc_id for doc_id in doc_ids
                                 if doc_id in doc_csv_split and
                                 doc_csv_split[doc_id] in ('test', 'Test', 'TEST')]

        else:
            # Fall back to random splitting
            n_total = len(doc_ids)
            n_train = int(n_total * self.train_ratio)
            n_val = int(n_total * self.val_ratio)

            if self.split == 'train':
                split_doc_ids = doc_ids[:n_train]
            elif self.split == 'val':
                split_doc_ids = doc_ids[n_train:n_train + n_val]
            else:  # test
                split_doc_ids = doc_ids[n_train + n_val:]

        # Collect items for this split
        split_items = []
        for doc_id in split_doc_ids:
            if doc_id in doc_items:
                split_items.extend(doc_items[doc_id])

        return split_items

    def __len__(self) -> int:
        return len(self.items)

    def __getitem__(self, idx: int) -> tuple[np.ndarray, np.ndarray]:
        """
        Get image and labels for index.

        Returns:
            (image, labels) where:
            - image: numpy array (H, W, C)
            - labels: numpy array (N, 5) with [class_id, x_center, y_center, width, height]
        """
        item = self.items[idx]

        # Load image using LRU cache (significant speedup during training)
        image_array, img_width, img_height = _load_image_cached(str(item.image_path))

        # Convert annotations to YOLO format (normalized)
        labels = self._convert_labels(item.labels, img_width, img_height, item.is_scanned)

        return image_array, labels

    def _convert_labels(
        self,
        annotations: list[YOLOAnnotation],
        img_width: int,
        img_height: int,
        is_scanned: bool = False
    ) -> np.ndarray:
        """
        Convert annotations to normalized YOLO format.

        Uses field-specific bbox expansion strategies via expand_bbox.

        Args:
            annotations: List of annotations
            img_width: Actual image width in pixels
            img_height: Actual image height in pixels
            is_scanned: If True, bbox is already in pixels; if False, bbox is in PDF points

        Returns:
            numpy array (N, 5) with [class_id, x_center, y_center, width, height]
        """
        if not annotations:
            return np.zeros((0, 5), dtype=np.float32)

        # Scale factor: PDF points (72 DPI) -> rendered pixels
        # Note: After the OCR coordinate fix, ALL bbox (both text and scanned PDF)
        # are stored in PDF points, so we always apply the same scaling.
        scale = self.dpi / 72.0

        labels = []
        for ann in annotations:
            # Convert center+size to corner coords in PDF points
            half_w = ann.width / 2
            half_h = ann.height / 2
            x0_pdf = ann.x_center - half_w
            y0_pdf = ann.y_center - half_h
            x1_pdf = ann.x_center + half_w
            y1_pdf = ann.y_center + half_h

            # Convert to pixels
            x0_px = x0_pdf * scale
            y0_px = y0_pdf * scale
            x1_px = x1_pdf * scale
            y1_px = y1_pdf * scale

            # Apply uniform bbox expansion
            x0, y0, x1, y1 = expand_bbox(
                bbox=(x0_px, y0_px, x1_px, y1_px),
                image_width=img_width,
                image_height=img_height,
            )

            # Ensure minimum height
            height_px = y1 - y0
            if height_px < self.min_bbox_height_px:
                extra = (self.min_bbox_height_px - height_px) / 2
                y0 = max(0, int(y0 - extra))
                y1 = min(img_height, int(y1 + extra))

            # Convert to YOLO format (normalized center + size)
            x_center = (x0 + x1) / 2 / img_width
            y_center = (y0 + y1) / 2 / img_height
            width = (x1 - x0) / img_width
            height = (y1 - y0) / img_height

            # Clamp to valid range
            x_center = max(0, min(1, x_center))
            y_center = max(0, min(1, y_center))
            width = max(0, min(1, width))
            height = max(0, min(1, height))

            labels.append([ann.class_id, x_center, y_center, width, height])

        return np.array(labels, dtype=np.float32)

    def get_image_path(self, idx: int) -> Path:
        """Get image path for index."""
        return self.items[idx].image_path

    def get_labels_for_yolo(self, idx: int) -> str:
        """
        Get YOLO format labels as string for index.

        Returns:
            String with YOLO format labels (one per line)
        """
        item = self.items[idx]

        # Use cached image loading to avoid duplicate disk reads
        _, img_width, img_height = _load_image_cached(str(item.image_path))

        labels = self._convert_labels(item.labels, img_width, img_height, item.is_scanned)

        lines = []
        for label in labels:
            class_id = int(label[0])
            x_center, y_center, width, height = label[1:5]
            lines.append(f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}")

        return '\n'.join(lines)

    def export_to_yolo_format(
        self,
        output_dir: str | Path,
        split_name: Optional[str] = None
    ) -> int:
        """
        Export dataset to standard YOLO format (images + label files).

        This is useful for training with standard YOLO training scripts.

        Args:
            output_dir: Output directory
            split_name: Name for the split subdirectory (default: self.split)

        Returns:
            Number of items exported
        """
        import shutil

        output_dir = Path(output_dir)
        split_name = split_name or self.split

        images_out = output_dir / split_name / 'images'
        labels_out = output_dir / split_name / 'labels'

        # Clear existing directories before export
        if images_out.exists():
            shutil.rmtree(images_out)
        if labels_out.exists():
            shutil.rmtree(labels_out)

        images_out.mkdir(parents=True, exist_ok=True)
        labels_out.mkdir(parents=True, exist_ok=True)

        count = 0
        for idx in range(len(self)):
            item = self.items[idx]

            # Copy image
            dest_image = images_out / item.image_path.name
            shutil.copy2(item.image_path, dest_image)

            # Write label file
            label_content = self.get_labels_for_yolo(idx)
            label_path = labels_out / f"{item.image_path.stem}.txt"
            with open(label_path, 'w') as f:
                f.write(label_content)

            count += 1

        logger.info("Exported %d items to %s", count, output_dir / split_name)
        return count


def create_datasets(
    images_dir: str | Path,
    db: Any,
    train_ratio: float = 0.8,
    val_ratio: float = 0.1,
    seed: int = 42,
    limit: int | None = None,
    **kwargs
) -> dict[str, DBYOLODataset]:
    """
    Create train/val/test datasets.

    This function loads data once and shares it across all splits for efficiency.

    Args:
        images_dir: Directory containing temp/{doc_id}/images/
        db: DocumentDB instance
        train_ratio: Training set ratio
        val_ratio: Validation set ratio
        seed: Random seed
        limit: Maximum number of documents to use (None for all)
        **kwargs: Additional arguments for DBYOLODataset

    Returns:
        Dict with 'train', 'val', 'test' datasets
    """
    # Create first dataset which loads all data
    logger.info("Loading dataset (this may take a few minutes for large datasets)...")
    first_dataset = DBYOLODataset(
        images_dir=images_dir,
        db=db,
        split='train',
        train_ratio=train_ratio,
        val_ratio=val_ratio,
        seed=seed,
        limit=limit,
        **kwargs
    )

    # Create other splits by sharing loaded data
    datasets = {'train': first_dataset}
    for split in ['val', 'test']:
        datasets[split] = DBYOLODataset.from_shared_data(
            first_dataset,
            split=split,
        )

    return datasets
