"""Training Export Endpoints."""

import logging
from datetime import datetime

from fastapi import APIRouter, HTTPException

from backend.web.core.auth import AdminTokenDep, DocumentRepoDep, AnnotationRepoDep
from backend.web.schemas.admin import (
    ExportRequest,
    ExportResponse,
)
from backend.web.schemas.common import ErrorResponse
from shared.bbox import expand_bbox

logger = logging.getLogger(__name__)


def register_export_routes(router: APIRouter) -> None:
    """Register export endpoints on the router."""

    @router.post(
        "/export",
        response_model=ExportResponse,
        responses={
            400: {"model": ErrorResponse, "description": "Invalid request"},
            401: {"model": ErrorResponse, "description": "Invalid token"},
        },
        summary="Export annotations",
        description="Export annotations in YOLO format for training.",
    )
    async def export_annotations(
        request: ExportRequest,
        admin_token: AdminTokenDep,
        docs: DocumentRepoDep,
        annotations: AnnotationRepoDep,
    ) -> ExportResponse:
        """Export annotations for training."""
        from backend.web.services.storage_helpers import get_storage_helper

        # Get storage helper for reading images and exports directory
        storage = get_storage_helper()

        if request.format not in ("yolo", "coco", "voc"):
            raise HTTPException(
                status_code=400,
                detail=f"Unsupported export format: {request.format}",
            )

        documents = docs.get_labeled_for_export(admin_token)

        if not documents:
            raise HTTPException(
                status_code=400,
                detail="No labeled documents available for export",
            )

        # Get exports directory from StorageHelper
        exports_base = storage.get_exports_base_path()
        if exports_base is None:
            raise HTTPException(
                status_code=500,
                detail="Storage not configured for local access",
            )
        export_dir = exports_base / f"export_{datetime.utcnow().strftime('%Y%m%d_%H%M%S')}"
        export_dir.mkdir(parents=True, exist_ok=True)

        (export_dir / "images" / "train").mkdir(parents=True, exist_ok=True)
        (export_dir / "images" / "val").mkdir(parents=True, exist_ok=True)
        (export_dir / "labels" / "train").mkdir(parents=True, exist_ok=True)
        (export_dir / "labels" / "val").mkdir(parents=True, exist_ok=True)

        total_docs = len(documents)
        train_count = int(total_docs * request.split_ratio)
        train_docs = documents[:train_count]
        val_docs = documents[train_count:]

        total_images = 0
        total_annotations = 0

        for split, docs in [("train", train_docs), ("val", val_docs)]:
            for doc in docs:
                doc_annotations = annotations.get_for_document(str(doc.document_id))

                if not doc_annotations:
                    continue

                for page_num in range(1, doc.page_count + 1):
                    page_annotations = [a for a in doc_annotations if a.page_number == page_num]

                    if not page_annotations and not request.include_images:
                        continue

                    # Get image from storage
                    doc_id = str(doc.document_id)
                    if not storage.admin_image_exists(doc_id, page_num):
                        continue

                    # Download image and save to export directory
                    image_name = f"{doc.document_id}_page{page_num}.png"
                    dst_image = export_dir / "images" / split / image_name
                    image_content = storage.get_admin_image(doc_id, page_num)
                    dst_image.write_bytes(image_content)
                    total_images += 1

                    # Get image dimensions for bbox expansion
                    img_dims = storage.get_admin_image_dimensions(doc_id, page_num)
                    if img_dims is None:
                        # Fall back to standard A4 at 300 DPI if dimensions unavailable
                        img_width, img_height = 2480, 3508
                    else:
                        img_width, img_height = img_dims

                    label_name = f"{doc.document_id}_page{page_num}.txt"
                    label_path = export_dir / "labels" / split / label_name

                    with open(label_path, "w") as f:
                        for ann in page_annotations:
                            # Convert normalized coords to pixel coords
                            half_w = (ann.width * img_width) / 2
                            half_h = (ann.height * img_height) / 2
                            x0 = ann.x_center * img_width - half_w
                            y0 = ann.y_center * img_height - half_h
                            x1 = ann.x_center * img_width + half_w
                            y1 = ann.y_center * img_height + half_h

                            # Apply uniform bbox expansion
                            ex0, ey0, ex1, ey1 = expand_bbox(
                                bbox=(x0, y0, x1, y1),
                                image_width=img_width,
                                image_height=img_height,
                            )

                            # Convert back to normalized YOLO format
                            new_x_center = (ex0 + ex1) / 2 / img_width
                            new_y_center = (ey0 + ey1) / 2 / img_height
                            new_width = (ex1 - ex0) / img_width
                            new_height = (ey1 - ey0) / img_height

                            # Clamp to valid range
                            new_x_center = max(0, min(1, new_x_center))
                            new_y_center = max(0, min(1, new_y_center))
                            new_width = max(0, min(1, new_width))
                            new_height = max(0, min(1, new_height))

                            line = f"{ann.class_id} {new_x_center:.6f} {new_y_center:.6f} {new_width:.6f} {new_height:.6f}\n"
                            f.write(line)
                            total_annotations += 1

        from shared.fields import FIELD_CLASSES

        yaml_content = f"""# Auto-generated YOLO dataset config
path: {export_dir.absolute()}
train: images/train
val: images/val

nc: {len(FIELD_CLASSES)}
names: {list(FIELD_CLASSES.values())}
"""
        (export_dir / "data.yaml").write_text(yaml_content)

        return ExportResponse(
            status="completed",
            export_path=str(export_dir),
            total_images=total_images,
            total_annotations=total_annotations,
            train_count=len(train_docs),
            val_count=len(val_docs),
            message=f"Exported {total_images} images with {total_annotations} annotations",
        )
