# Copyright (c) 2023-2024 Geosiris.
# SPDX-License-Identifier: Apache-2.0
"""
Module for exporting mesh data to various file formats.
Supports OBJ, GeoJSON, VTK, and STL formats.
"""

import json
import struct
from enum import Enum
from pathlib import Path
from typing import TYPE_CHECKING, BinaryIO, List, Optional, TextIO, Union

import numpy as np

if TYPE_CHECKING:
    from .mesh import AbstractMesh


class ExportFormat(Enum):
    """Supported mesh export formats."""

    OBJ = "obj"
    GEOJSON = "geojson"
    VTK = "vtk"
    STL = "stl"

    @classmethod
    def from_extension(cls, extension: str) -> "ExportFormat":
        """Get format from file extension."""
        ext = extension.lower().lstrip(".")
        for fmt in cls:
            if fmt.value == ext:
                return fmt
        raise ValueError(f"Unsupported file extension: {extension}")

    @classmethod
    def all_extensions(cls) -> List[str]:
        """Get all supported file extensions."""
        return [fmt.value for fmt in cls]


class ExportOptions:
    """Base class for export options."""

    pass


class STLExportOptions(ExportOptions):
    """Options for STL export."""

    def __init__(self, binary: bool = True, ascii_precision: int = 6):
        """
        Initialize STL export options.

        :param binary: If True, export as binary STL; if False, export as ASCII STL
        :param ascii_precision: Number of decimal places for ASCII format
        """
        self.binary = binary
        self.ascii_precision = ascii_precision


class VTKExportOptions(ExportOptions):
    """Options for VTK export."""

    def __init__(self, binary: bool = False, dataset_name: str = "mesh"):
        """
        Initialize VTK export options.

        :param binary: If True, export as binary VTK; if False, export as ASCII VTK
        :param dataset_name: Name of the dataset in VTK file
        """
        self.binary = binary
        self.dataset_name = dataset_name


class GeoJSONExportOptions(ExportOptions):
    """Options for GeoJSON export."""

    def __init__(self, indent: Optional[int] = 2, properties: Optional[dict] = None):
        """
        Initialize GeoJSON export options.

        :param indent: JSON indentation level (None for compact)
        :param properties: Additional properties to include in features
        """
        self.indent = indent
        self.properties = properties or {}


def export_obj(mesh_list: List["AbstractMesh"], out: BinaryIO, obj_name: Optional[str] = None) -> None:
    """
    Export mesh data to Wavefront OBJ format.

    :param mesh_list: List of AbstractMesh objects to export
    :param out: Binary output stream
    :param obj_name: Optional object name for the OBJ file
    """
    # Lazy import to avoid circular dependency
    from .mesh import PolylineSetMesh

    # Write header
    out.write(b"# Generated by energyml-utils a Geosiris python module\n\n")

    # Write object name if provided
    if obj_name is not None:
        out.write(f"o {obj_name}\n\n".encode("utf-8"))

    point_offset = 0

    for mesh in mesh_list:
        # Write group name using mesh identifier or uuid
        mesh_id = getattr(mesh, "identifier", None) or getattr(mesh, "uuid", "mesh")
        out.write(f"g {mesh_id}\n\n".encode("utf-8"))

        # Write vertices
        for point in mesh.point_list:
            if len(point) > 0:
                out.write(f"v {' '.join(map(str, point))}\n".encode("utf-8"))

        # Write faces or lines depending on mesh type
        indices = mesh.get_indices()
        elt_letter = "l" if isinstance(mesh, PolylineSetMesh) else "f"

        for face_or_line in indices:
            if len(face_or_line) > 1:
                # OBJ indices are 1-based
                indices_str = " ".join(str(idx + point_offset + 1) for idx in face_or_line)
                out.write(f"{elt_letter} {indices_str}\n".encode("utf-8"))

        point_offset += len(mesh.point_list)


def export_geojson(
    mesh_list: List["AbstractMesh"], out: TextIO, options: Optional[GeoJSONExportOptions] = None
) -> None:
    """
    Export mesh data to GeoJSON format.

    :param mesh_list: List of AbstractMesh objects to export
    :param out: Text output stream
    :param options: GeoJSON export options
    """
    # Lazy import to avoid circular dependency
    from .mesh import PolylineSetMesh, SurfaceMesh

    if options is None:
        options = GeoJSONExportOptions()

    features = []

    for mesh_idx, mesh in enumerate(mesh_list):
        indices = mesh.get_indices()

        if isinstance(mesh, PolylineSetMesh):
            # Export as LineString features
            for line_idx, line_indices in enumerate(indices):
                if len(line_indices) < 2:
                    continue
                coordinates = [list(mesh.point_list[idx]) for idx in line_indices]
                feature = {
                    "type": "Feature",
                    "geometry": {"type": "LineString", "coordinates": coordinates},
                    "properties": {"mesh_index": mesh_idx, "line_index": line_idx, **options.properties},
                }
                features.append(feature)

        elif isinstance(mesh, SurfaceMesh):
            # Export as Polygon features
            for face_idx, face_indices in enumerate(indices):
                if len(face_indices) < 3:
                    continue
                # GeoJSON Polygon requires closed ring (first point == last point)
                coordinates = [list(mesh.point_list[idx]) for idx in face_indices]
                coordinates.append(coordinates[0])  # Close the ring

                feature = {
                    "type": "Feature",
                    "geometry": {"type": "Polygon", "coordinates": [coordinates]},
                    "properties": {"mesh_index": mesh_idx, "face_index": face_idx, **options.properties},
                }
                features.append(feature)

    geojson = {"type": "FeatureCollection", "features": features}

    json.dump(geojson, out, indent=options.indent)


def export_vtk(mesh_list: List["AbstractMesh"], out: BinaryIO, options: Optional[VTKExportOptions] = None) -> None:
    """
    Export mesh data to VTK legacy format.

    :param mesh_list: List of AbstractMesh objects to export
    :param out: Binary output stream
    :param options: VTK export options
    """
    # Lazy import to avoid circular dependency
    from .mesh import PolylineSetMesh, SurfaceMesh

    if options is None:
        options = VTKExportOptions()

    # Combine all meshes
    all_points = []
    all_polygons = []
    all_lines = []
    vertex_offset = 0

    for mesh in mesh_list:
        all_points.extend(mesh.point_list)
        indices = mesh.get_indices()

        if isinstance(mesh, SurfaceMesh):
            # Adjust face indices
            for face in indices:
                adjusted_face = [idx + vertex_offset for idx in face]
                all_polygons.append(adjusted_face)
        elif isinstance(mesh, PolylineSetMesh):
            # Adjust line indices
            for line in indices:
                adjusted_line = [idx + vertex_offset for idx in line]
                all_lines.append(adjusted_line)

        vertex_offset += len(mesh.point_list)

    # Write VTK header
    out.write(b"# vtk DataFile Version 3.0\n")
    out.write(f"{options.dataset_name}\n".encode("utf-8"))
    out.write(b"ASCII\n")
    out.write(b"DATASET POLYDATA\n")

    # Write points
    out.write(f"POINTS {len(all_points)} float\n".encode("utf-8"))
    for point in all_points:
        out.write(f"{point[0]} {point[1]} {point[2]}\n".encode("utf-8"))

    # Write polygons
    if all_polygons:
        total_poly_size = sum(len(poly) + 1 for poly in all_polygons)
        out.write(f"POLYGONS {len(all_polygons)} {total_poly_size}\n".encode("utf-8"))
        for poly in all_polygons:
            out.write(f"{len(poly)} {' '.join(str(idx) for idx in poly)}\n".encode("utf-8"))

    # Write lines
    if all_lines:
        total_line_size = sum(len(line) + 1 for line in all_lines)
        out.write(f"LINES {len(all_lines)} {total_line_size}\n".encode("utf-8"))
        for line in all_lines:
            out.write(f"{len(line)} {' '.join(str(idx) for idx in line)}\n".encode("utf-8"))


def export_stl(mesh_list: List["AbstractMesh"], out: BinaryIO, options: Optional[STLExportOptions] = None) -> None:
    """
    Export mesh data to STL format (binary or ASCII).

    Note: STL format only supports triangles. Only triangular faces will be exported.

    :param mesh_list: List of AbstractMesh objects to export
    :param out: Binary output stream
    :param options: STL export options
    """
    # Lazy import to avoid circular dependency
    from .mesh import SurfaceMesh

    if options is None:
        options = STLExportOptions(binary=True)

    # Collect all triangles (only from SurfaceMesh with triangular faces)
    all_triangles = []
    for mesh in mesh_list:
        if isinstance(mesh, SurfaceMesh):
            indices = mesh.get_indices()
            for face in indices:
                # Only export triangular faces
                if len(face) == 3:
                    p0 = np.array(mesh.point_list[face[0]])
                    p1 = np.array(mesh.point_list[face[1]])
                    p2 = np.array(mesh.point_list[face[2]])
                    all_triangles.append((p0, p1, p2))

    if options.binary:
        _export_stl_binary(all_triangles, out)
    else:
        _export_stl_ascii(all_triangles, out, options.ascii_precision)


def _export_stl_binary(triangles: List[tuple], out: BinaryIO) -> None:
    """Export STL in binary format."""
    # Write 80-byte header
    header = b"Binary STL file generated by energyml-utils" + b"\0" * (80 - 44)
    out.write(header)

    # Write number of triangles
    out.write(struct.pack("<I", len(triangles)))

    # Write each triangle
    for p0, p1, p2 in triangles:
        # Calculate normal vector
        v1 = p1 - p0
        v2 = p2 - p0
        normal = np.cross(v1, v2)
        norm = np.linalg.norm(normal)
        if norm > 0:
            normal = normal / norm
        else:
            normal = np.array([0.0, 0.0, 0.0])

        # Write normal
        out.write(struct.pack("<fff", float(normal[0]), float(normal[1]), float(normal[2])))

        # Write vertices
        for point in [p0, p1, p2]:
            out.write(struct.pack("<fff", float(point[0]), float(point[1]), float(point[2])))

        # Write attribute byte count (unused)
        out.write(struct.pack("<H", 0))


def _export_stl_ascii(triangles: List[tuple], out: BinaryIO, precision: int) -> None:
    """Export STL in ASCII format."""
    out.write(b"solid mesh\n")

    for p0, p1, p2 in triangles:
        # Calculate normal vector
        v1 = p1 - p0
        v2 = p2 - p0
        normal = np.cross(v1, v2)
        norm = np.linalg.norm(normal)
        if norm > 0:
            normal = normal / norm
        else:
            normal = np.array([0.0, 0.0, 0.0])

        # Write facet
        line = f"  facet normal {normal[0]:.{precision}e} {normal[1]:.{precision}e} {normal[2]:.{precision}e}\n"
        out.write(line.encode("utf-8"))
        out.write(b"    outer loop\n")

        for point in [p0, p1, p2]:
            line = f"      vertex {point[0]:.{precision}e} {point[1]:.{precision}e} {point[2]:.{precision}e}\n"
            out.write(line.encode("utf-8"))

        out.write(b"    endloop\n")
        out.write(b"  endfacet\n")

    out.write(b"endsolid mesh\n")


def export_mesh(
    mesh_list: List["AbstractMesh"],
    output_path: Union[str, Path],
    format: Optional[ExportFormat] = None,
    options: Optional[ExportOptions] = None,
) -> None:
    """
    Export mesh data to a file in the specified format.

    :param mesh_list: List of Mesh objects to export
    :param output_path: Output file path
    :param format: Export format (auto-detected from extension if None)
    :param options: Format-specific export options
    """
    path = Path(output_path)

    # Auto-detect format from extension if not specified
    if format is None:
        format = ExportFormat.from_extension(path.suffix)

    # Determine if file should be opened in binary or text mode
    binary_formats = {ExportFormat.OBJ, ExportFormat.STL, ExportFormat.VTK}
    text_formats = {ExportFormat.GEOJSON}

    if format in binary_formats:
        with path.open("wb") as f:
            if format == ExportFormat.OBJ:
                export_obj(mesh_list, f)
            elif format == ExportFormat.STL:
                export_stl(mesh_list, f, options)
            elif format == ExportFormat.VTK:
                export_vtk(mesh_list, f, options)
    elif format in text_formats:
        with path.open("w", encoding="utf-8") as f:
            if format == ExportFormat.GEOJSON:
                export_geojson(mesh_list, f, options)
    else:
        raise ValueError(f"Unsupported format: {format}")


# UI Helper Functions


def supported_formats() -> List[str]:
    """
    Get list of supported export formats.

    :return: List of format names (e.g., ['obj', 'geojson', 'vtk', 'stl'])
    """
    return ExportFormat.all_extensions()


def format_description(format: Union[str, ExportFormat]) -> str:
    """
    Get human-readable description of a format.

    :param format: Format name or ExportFormat enum
    :return: Description string
    """
    if isinstance(format, str):
        format = ExportFormat.from_extension(format)

    descriptions = {
        ExportFormat.OBJ: "Wavefront OBJ - 3D geometry format (triangles and lines)",
        ExportFormat.GEOJSON: "GeoJSON - Geographic data format (lines and polygons)",
        ExportFormat.VTK: "VTK Legacy - Visualization Toolkit format",
        ExportFormat.STL: "STL - Stereolithography format (triangles only)",
    }
    return descriptions.get(format, "Unknown format")


def format_filter_string(format: Union[str, ExportFormat]) -> str:
    """
    Get file filter string for UI dialogs (Qt, tkinter, etc.).

    :param format: Format name or ExportFormat enum
    :return: Filter string (e.g., "OBJ Files (*.obj)")
    """
    if isinstance(format, str):
        format = ExportFormat.from_extension(format)

    filters = {
        ExportFormat.OBJ: "OBJ Files (*.obj)",
        ExportFormat.GEOJSON: "GeoJSON Files (*.geojson)",
        ExportFormat.VTK: "VTK Files (*.vtk)",
        ExportFormat.STL: "STL Files (*.stl)",
    }
    return filters.get(format, "All Files (*.*)")


def all_formats_filter_string() -> str:
    """
    Get file filter string for all supported formats.
    Useful for Qt QFileDialog or similar UI components.

    :return: Filter string with all formats
    """
    filters = [format_filter_string(fmt) for fmt in ExportFormat]
    return ";;".join(filters)


def get_format_options_class(format: Union[str, ExportFormat]) -> Optional[type]:
    """
    Get the options class for a specific format.

    :param format: Format name or ExportFormat enum
    :return: Options class or None if no options available
    """
    if isinstance(format, str):
        format = ExportFormat.from_extension(format)

    options_map = {
        ExportFormat.STL: STLExportOptions,
        ExportFormat.VTK: VTKExportOptions,
        ExportFormat.GEOJSON: GeoJSONExportOptions,
    }
    return options_map.get(format)


def supports_lines(format: Union[str, ExportFormat]) -> bool:
    """
    Check if format supports line primitives.

    :param format: Format name or ExportFormat enum
    :return: True if format supports lines
    """
    if isinstance(format, str):
        format = ExportFormat.from_extension(format)

    return format in {ExportFormat.OBJ, ExportFormat.GEOJSON, ExportFormat.VTK}


def supports_triangles(format: Union[str, ExportFormat]) -> bool:
    """
    Check if format supports triangle primitives.

    :param format: Format name or ExportFormat enum
    :return: True if format supports triangles
    """
    # All formats support triangles
    return True
