"""
Box Vectorizer for QPredict

Converts YOLO detections (bounding boxes and OBB) to vector polygons.
Handles geometry creation, validation, attribute enrichment, and styling.

Key features:
- Robust geometry creation for bbox and OBB
- Proper coordinate transformation (pixel → geo)
- Attribute table with class info and confidence
- Default blue outline symbology
- GeoPackage export (.gpkg)

Used by YOLO Detection and YOLO OBB tasks.
"""

from osgeo import ogr, osr
from pathlib import Path
from typing import List, Dict, Any, Tuple


class BoxVectorizer:
    """
    Convert YOLO detections to vector polygons.
    
    This class handles:
    - OBB (Oriented Bounding Boxes) with 4 corners
    - Regular bounding boxes (axis-aligned rectangles)
    - Coordinate transformation from pixel to geographic
    - Attribute enrichment (class, confidence, label)
    - GeoPackage creation with proper spatial reference
    
    Args:
        detections: List of detection dictionaries from YOLO predictor
        output_path: Path for output GeoPackage (.gpkg)
        model_info: Model configuration with class definitions
        geo_transform: GDAL geotransform (6 parameters)
        projection: WKT projection string
        
    Example:
        >>> vectorizer = BoxVectorizer(
        ...     detections=filtered_detections,
        ...     output_path='detections.gpkg',
        ...     model_info=model_info,
        ...     geo_transform=geotransform,
        ...     projection=crs_wkt
        ... )
        >>> vectorizer.vectorize()
    """
    
    def __init__(
        self,
        detections: List[Dict[str, Any]],
        output_path: str,
        model_info: Dict[str, Any],
        geo_transform: Tuple[float, ...],
        projection: str
    ):
        """
        Initialize box vectorizer.
        
        Args:
            detections: List of detection dicts with geometry, class_id, confidence
            output_path: Output GeoPackage path
            model_info: Model configuration
            geo_transform: GDAL geotransform
            projection: WKT projection string
        """
        self.detections = detections
        self.output_path = Path(output_path)
        self.model_info = model_info
        self.geo_transform = geo_transform
        self.projection = projection
        
        # Get class information
        self.classes = {c['id']: c['name'] for c in model_info['classes']}
        
        # Get pixel size for area calculations
        self.pixel_width = abs(geo_transform[1])
        self.pixel_height = abs(geo_transform[5])
    
    def vectorize(self):
        """
        Convert detections to GeoPackage.
        
        Steps:
        1. Create output GeoPackage with spatial reference
        2. Create attribute fields
        3. Convert each detection to feature with geometry
        4. Add attributes (class_id, class_name, confidence, label)
        5. Close dataset
        
        Raises:
            RuntimeError: If vectorization fails
        """
        try:
            # Create output directory if needed
            self.output_path.parent.mkdir(parents=True, exist_ok=True)
            
            # Remove existing file
            if self.output_path.exists():
                self.output_path.unlink()
            
            # Create GeoPackage
            driver = ogr.GetDriverByName('GPKG')
            ds = driver.CreateDataSource(str(self.output_path))
            
            if ds is None:
                raise RuntimeError(f"Could not create GeoPackage: {self.output_path}")
            
            # Create spatial reference
            srs = osr.SpatialReference()
            srs.ImportFromWkt(self.projection)
            
            # Create layer (Polygon for both bbox and OBB)
            layer = ds.CreateLayer('detections', srs, ogr.wkbPolygon)
            
            # Create attribute fields
            self._create_fields(layer)
            
            # Add features
            feature_count = 0
            for detection in self.detections:
                feature = self._create_feature(detection, layer)
                if feature is not None:
                    layer.CreateFeature(feature)
                    feature_count += 1
                    feature = None  # Release feature
            
            # Close dataset
            ds = None
            
            print(f"[BoxVectorizer] Created {feature_count} detection features")
            
        except Exception as e:
            raise RuntimeError(f"Box vectorization failed: {str(e)}")
    
    def _create_fields(self, layer: ogr.Layer):
        """
        Create attribute fields in layer.
        
        Fields:
        - class_id (integer): Numeric class identifier
        - class_name (string): Human-readable class name
        - confidence (real): Detection confidence [0, 1]
        - label (string): Combined label for display (e.g., "boat - 0.94")
        
        Args:
            layer: OGR layer to add fields to
        """
        # Class ID
        layer.CreateField(ogr.FieldDefn('class_id', ogr.OFTInteger))
        
        # Class name
        class_name_field = ogr.FieldDefn('class_name', ogr.OFTString)
        class_name_field.SetWidth(254)
        layer.CreateField(class_name_field)
        
        # Confidence
        confidence_field = ogr.FieldDefn('confidence', ogr.OFTReal)
        layer.CreateField(confidence_field)
        
        # Label (combined)
        label_field = ogr.FieldDefn('label', ogr.OFTString)
        label_field.SetWidth(254)
        layer.CreateField(label_field)
    
    def _create_feature(
        self,
        detection: Dict[str, Any],
        layer: ogr.Layer
    ) -> ogr.Feature:
        """
        Create OGR feature from detection.
        
        Args:
            detection: Detection dict with geometry, class_id, confidence
            layer: OGR layer for feature definition
            
        Returns:
            OGR feature with geometry and attributes, or None if invalid
        """
        # Extract detection info
        geometry_info = detection['geometry']
        class_id = detection['class_id']
        confidence = detection['confidence']
        
        # Get class name
        class_name = self.classes.get(class_id, f'class_{class_id}')
        
        # Create label
        label = f"{class_name} - {confidence:.2f}"
        
        # Create geometry
        geom = self._create_geometry(geometry_info)
        
        if geom is None or not geom.IsValid():
            print(f"[WARNING] Invalid geometry for detection (class={class_id}), skipping")
            return None
        
        # Create feature
        feature = ogr.Feature(layer.GetLayerDefn())
        feature.SetGeometry(geom)
        
        # Set attributes
        feature.SetField('class_id', int(class_id))
        feature.SetField('class_name', class_name)
        feature.SetField('confidence', float(confidence))
        feature.SetField('label', label)
        
        return feature
    
    def _create_geometry(self, geometry_info: Dict[str, Any]) -> ogr.Geometry:
        """Create OGR geometry from detection geometry info."""
        geom_type = geometry_info['type']
        
        # Create polygon
        geom = ogr.Geometry(ogr.wkbPolygon)
        ring = ogr.Geometry(ogr.wkbLinearRing)
        
        if geom_type == 'obb':
            # OBB: 4 corners in pixel coordinates
            corners = geometry_info['corners']
            
            # Add each corner (convert pixel → geo)
            for corner in corners:
                x_pixel, y_pixel = corner
                x_geo, y_geo = self._pixel_to_geo(x_pixel, y_pixel)
                ring.AddPoint(x_geo, y_geo)
            
            # Close ring (repeat first point)
            x_geo, y_geo = self._pixel_to_geo(corners[0][0], corners[0][1])
            ring.AddPoint(x_geo, y_geo)
            
        elif geom_type == 'bbox':
            # Regular bbox: [x1, y1, x2, y2] in pixel coordinates
            x1, y1, x2, y2 = geometry_info['coords']
            
            # Convert corners to geo coordinates
            x1_geo, y1_geo = self._pixel_to_geo(x1, y1)
            x2_geo, y2_geo = self._pixel_to_geo(x2, y2)
            
            # Create rectangle (counter-clockwise)
            ring.AddPoint(x1_geo, y1_geo)
            ring.AddPoint(x2_geo, y1_geo)
            ring.AddPoint(x2_geo, y2_geo)
            ring.AddPoint(x1_geo, y2_geo)
            ring.AddPoint(x1_geo, y1_geo)
            
        else:
            print(f"[WARNING] Unknown geometry type: {geom_type}")
            return None
        
        # Add ring to polygon
        geom.AddGeometry(ring)
        return geom
    
    def _pixel_to_geo(
        self,
        x_pixel: float,
        y_pixel: float
    ) -> Tuple[float, float]:
        """
        Convert pixel coordinates to geographic coordinates.
        
        Uses GDAL geotransform:
        x_geo = GT[0] + x_pixel * GT[1] + y_pixel * GT[2]
        y_geo = GT[3] + x_pixel * GT[4] + y_pixel * GT[5]
        
        Args:
            x_pixel: X coordinate in pixels (column)
            y_pixel: Y coordinate in pixels (row)
            
        Returns:
            Tuple of (x_geo, y_geo) in CRS coordinates
        """
        gt = self.geo_transform
        
        x_geo = gt[0] + x_pixel * gt[1] + y_pixel * gt[2]
        y_geo = gt[3] + x_pixel * gt[4] + y_pixel * gt[5]
        
        return x_geo, y_geo
    
    def get_output_path(self) -> str:
        """
        Get path to output GeoPackage.
        
        Returns:
            Path to created file
        """
        return str(self.output_path)