from pathlib import Path
from typing import Dict, List, Optional, Tuple, Union

from highlighter.core.data_models import DataSample
from highlighter.core.exceptions import require_package

try:
    import cv2
except ModuleNotFoundError as _:
    cv2 = None
import numpy as np
from PIL import Image
from pydantic import BaseModel, field_validator
from shapely import affinity
from shapely import geometry as gm

PathLike = Union[str, Path]
LTRBRect = Tuple[int, int, int, int]

KEY_CROP_ROTATED_RECT = "crop_rotated_rect"

__all__ = [
    "CropArgs",
    "KEY_CROP_ROTATED_RECT",
    "crop_rect_from_poly",
    "crop_rectangle",
    "crop_rotated_rectangle",
]


class CropArgs(BaseModel):
    crop_rotated_rect: Optional[bool] = False
    scale: Optional[float] = None
    pad: Optional[int] = None
    warped_wh: Optional[Union[Tuple[int, int], List[int]]] = None

    @field_validator("warped_wh")
    @classmethod
    def _validate_warped_wh(cls, v):
        """If warped_wh is not None make sure it's a list and not a tuple
        because tuple is not json/safe_yaml serializable.
        """

        if isinstance(v, tuple):
            v = list(v)

        if v is not None:
            assert isinstance(v, list)
            assert len(v) == 2
        return v

    def dict(self, **kwargs):
        d = super().model_dump(**kwargs)
        if not self.crop_rotated_rect:
            d.pop("warped_wh")
        return d

    def __call__(
        self, crop_locations: List[LTRBRect | gm.Polygon], image: Union[DataSample, np.ndarray, Image.Image]
    ):
        if isinstance(image, DataSample):
            return image.crop_content(crop_locations, crop_args=self, as_data_sample=True)
        else:
            if self.crop_rotated_rect:
                crop_fn = crop_rotated_rectangle
            else:
                crop_fn = crop_rectangle

            kwargs = self.model_dump()

            polys = []
            for loc in crop_locations:
                if hasattr(loc, "bounds"):
                    polys.append(loc)
                elif isinstance(loc, (list, tuple)) and len(loc) == 4:
                    x0, y0, x1, y1 = loc
                    polys.append(gm.Polygon([(x0, y0), (x1, y0), (x1, y1), (x0, y1), (x0, y0)]))
                else:
                    raise ValueError(f"Invalid crop location, got {loc}")

            return [crop_fn(image, poly, **kwargs) for poly in polys]


def _get_wh_of_rotated_rect(np_corners: np.ndarray) -> Tuple[int, int]:
    # First 3 corners of the rectangle
    a, b, c = np_corners[0], np_corners[1], np_corners[2]
    ab = int(round(np.linalg.norm(b - a)))
    bc = int(round(np.linalg.norm(c - b)))
    if ab >= bc:
        w, h = ab, bc
    else:
        w, h = bc, ab

    return w, h


def _order_rect_points_bottom_left_ccw(rect: gm.Polygon) -> gm.Polygon:
    # 0. Note, the origin for Shapely is the bottom left with
    #    the positive x and y axis right and up respectivally.
    #    But, our coordinates have the origin at the top left
    #    and positive x,y being rigth and down.
    #    So our top right and clockwise will translate to shapely's
    #    Bottom left and counter-clockwise. See below for an illustration
    #
    #    Shapely coords:
    #
    #    +y
    #    |
    #    | 0,1      2,1
    #    * * * * * *
    #    *         *
    #    * * * * * * --- +x
    #     0,0      2,0

    #    Highlighter coords:
    #
    #     0,0      2,0
    #    * * * * * * --- +x
    #    *         *
    #    * * * * * *
    #    |0,1      2,1
    #    |
    #    +y
    #
    #
    # The following uses the shapely coordinate system, that way, if we
    # look for bottom left and order ccw we should be golden
    #
    # 1. Find the edges that belong to the major axis (longest)
    # 2. Select the bottom edge
    # 3. Find the left most point of the bottom edge
    # 4. Set bottom left point to 0th index
    # 5. remaining points follow ccw

    rect_points = np.array(rect.exterior.coords)[:-1]

    # Find major axis
    a, b, c, d = 0, 1, 2, 3
    x, y = 0, 1
    ab = gm.LineString((rect_points[a], rect_points[b]))
    bc = gm.LineString((rect_points[b], rect_points[c]))
    cd = gm.LineString((rect_points[c], rect_points[d]))
    da = gm.LineString((rect_points[d], rect_points[a]))
    edges = np.array([ab, bc, cd, da])
    longest = edges[np.argsort([edge.length for edge in edges])[-2:]]

    edge_0 = np.array(longest[0].coords)
    edge_1 = np.array(longest[1].coords)

    # Find bottom edge
    if edge_0[..., y].min() < edge_1[..., y].min():
        bottom_edge = edge_0
        top_edge = edge_1
    else:
        bottom_edge = edge_1
        top_edge = edge_0

    # arrange points starting from the left most point of the bottom
    # edge and moving counter clockwise
    if bottom_edge[0][x] > bottom_edge[1][x]:
        bottom_edge = bottom_edge[[1, 0]]
        top_edge = top_edge[[1, 0]]

    return gm.Polygon(bottom_edge.tolist() + top_edge.tolist())


def crop_rect_from_poly(
    image: Union[np.ndarray, Image.Image],
    polys: Union[gm.Polygon, List[gm.Polygon]],
    crop_args: Optional[Union[CropArgs, Dict]] = None,
) -> Union[List[np.ndarray], List[Image.Image], np.ndarray, Image.Image]:
    single_crop = False
    if isinstance(polys, gm.Polygon):
        polys = [polys]
        single_crop = True

    if crop_args is None:
        crop_args = {}

    if isinstance(crop_args, dict):
        crop_args = CropArgs(**crop_args)
    elif not isinstance(crop_args, CropArgs):
        raise ValueError(f"invalid parameter 'crop_args' expected a CropArgs object, got {crop_args}")

    kwargs = crop_args.dict()

    crop_rotated_rect = kwargs.pop(KEY_CROP_ROTATED_RECT)

    if crop_rotated_rect:
        crop_fn = crop_rotated_rectangle
    else:
        crop_fn = crop_rectangle

    crops = [crop_fn(image, poly, **kwargs) for poly in polys]
    if single_crop:
        return crops[0]
    else:
        return crops


@require_package(cv2, "cv2", "opencv")
def crop_rotated_rectangle(
    image: Union[np.ndarray, Image.Image],
    poly: gm.Polygon,
    scale: Optional[float] = None,
    pad: Optional[int] = None,
    warped_wh: Optional[Tuple[int, int]] = None,
) -> Union[np.ndarray, Image.Image]:
    """Find the minimum_rotated_rectange for the input poly and apply
    cv2.warpPerspective to transform to to rectangular crop of shape warped_wh

    see: https://shapely.readthedocs.io/en/stable/manual.html#object.minimum_rotated_rectangle

    Resulting crops will be of the same type as the input image (np.ndarray | PIL.Image)

    Args:
        image: Image to crop
        poly: Polygon to crop out
        scale: (Optional) factor applied to poly before cropping about the origin
        pad: Optional pad (in pixels) applied to poly before cropping
        warped_wh: The desired output shape of the crop
    """
    assert isinstance(pad, int) or (
        pad is None
    ), f"Croping rotated rec only support a int for 'pad', got: {pad}"

    rot_rect = poly.minimum_rotated_rectangle

    if scale is not None:
        rot_rect = affinity.scale(rot_rect, xfact=scale, yfact=scale, origin=(0, 0, 0))

    if pad is not None:
        rot_rect = rot_rect.buffer(pad, join_style=2)

    # Shapely uses a different coordinate system to us and opencv.
    # We need to make sure the points in the src_points corrispond to
    # the points in dst_points, hence this function is needed. More detail in
    # its doc string.
    rot_rect = _order_rect_points_bottom_left_ccw(rot_rect)
    src_points = np.array(list(zip(*rot_rect.exterior.coords.xy)))[:-1].astype(np.float32)

    if warped_wh is None:
        warped_wh = _get_wh_of_rotated_rect(src_points)
    w, h = warped_wh

    dst_points = np.array([[0, 0], [w, 0], [w, h], [0, h]], dtype=np.float32)

    transform_mat = cv2.getPerspectiveTransform(src_points, dst_points)

    if isinstance(image, Image.Image):
        src = np.array(image)
    elif isinstance(image, np.ndarray):
        src = image
    else:
        raise ValueError(f"Invalid image input type: {type(image)}")

    out = cv2.warpPerspective(src, transform_mat, warped_wh, flags=cv2.INTER_LINEAR)

    return out if isinstance(image, np.ndarray) else Image.fromarray(out)


def crop_rectangle(
    image: Union[np.ndarray, Image.Image],
    rect: gm.Polygon,
    scale: Optional[float] = None,
    pad: Optional[int] = None,
) -> Union[np.ndarray, Image.Image]:
    """Crop an XY Axis aligned rectangle crop image

    Resulting crops will be of the same type as the input image (np.ndarray | PIL.Image)

    Args:
        image: Image to crop
        rect: Left, Top, Right, Bottom tuple representing the rectangle to crop
        scale: Optional factor applied to rect before cropping
        pad: Optional pad (in pixels) applied to rect before cropping

    Note: Pillow.Image.crop does not give a 💩 if a crop goes out-of-bounds
    it will just add zeros as needed.
    """
    if hasattr(rect, "bounds"):
        # shapely object
        rect = rect.bounds
    else:
        raise ValueError("rect argument must have a .bounds attribute in order to crop")

    if scale is not None:
        rect = [x * scale for x in rect]

    if pad is None:
        pad = (0, 0)
    elif isinstance(pad, (float, int)):
        pad = (int(pad), int(pad))
    assert isinstance(pad, tuple)

    def pad_rect(rect, w_pad, h_pad):
        l, t, r, b = rect
        l -= w_pad
        r += w_pad
        t -= h_pad
        b += h_pad
        return (l, t, r, b)

    l, t, r, b = [int(i) for i in pad_rect(rect, *pad)]

    if isinstance(image, Image.Image):
        crop = image.crop((l, t, r, b))
    elif isinstance(image, np.ndarray):
        crop = np_crop_with_auto_pad(image, (l, t, r, b))
    else:
        raise ValueError(f"Invalid image input type: {type(image)}")
    return crop


def np_crop_with_auto_pad(
    img: np.ndarray,
    xyxybox: Tuple[int, int, int, int],
) -> np.ndarray:
    """
    Crops an image using NumPy. If the crop region is partially or fully
    outside the image bounds, pads with zeros. This matches the behaviour
    of PIL.Image.crop

    Parameters:
    -----------
    img : np.ndarray
        Input image (H, W) or (H, W, C)
    y, x : int
        Top-left corner of the crop (can be negative)
    h, w : int
        Desired height and width of output crop

    Returns:
    --------
    np.ndarray
        Cropped (and possibly zero-padded) image of shape (h, w) or (h, w, C)
    """
    x0, y0, x1, y1 = xyxybox
    h = y1 - y0
    w = x1 - x0
    if h <= 0 or w <= 0:
        raise ValueError("Crop height and width must be positive")

    # Get image dimensions
    img_h, img_w = img.shape[:2]
    channels = img.shape[2] if img.ndim == 3 else 1

    # Create output array filled with zeros
    dtype = img.dtype
    if img.ndim == 3:
        out = np.zeros((h, w, channels), dtype=dtype)
    else:
        out = np.zeros((h, w), dtype=dtype)

    # Calculate source and destination regions
    # src_y_start, src_y_end → region in original image
    # dst_y_start, dst_y_end → region in output image

    src_y_start = max(y0, 0)
    src_y_end = min(y1, img_h)
    src_x_start = max(x0, 0)
    src_x_end = min(x1, img_w)

    dst_y_start = max(-y0, 0)
    dst_y_end = dst_y_start + (src_y_end - src_y_start)
    dst_x_start = max(-x0, 0)
    dst_x_end = dst_x_start + (src_x_end - src_x_start)

    # If there's any overlap → copy the valid region
    if (src_y_start < src_y_end) and (src_x_start < src_x_end):
        out[dst_y_start:dst_y_end, dst_x_start:dst_x_end] = img[src_y_start:src_y_end, src_x_start:src_x_end]

    return out
