from __future__ import annotations

import importlib.util
from pathlib import Path
from typing import Any
from zipfile import ZIP_DEFLATED, ZipFile
from xml.sax.saxutils import escape

import numpy as np

from cad_converter import CADConverter
from config_loader import ConfigLoader
from evaluation import Evaluator
from io_handler import IOHandler
from preprocessing.curvature import CurvatureEstimator
from preprocessing.filters import PreprocessingRunner
from preprocessing.noise import GaussianNoiseAdder
from preprocessing.plane import PlaneEstimator
from registration import PointCloudRegistration
from visualization import Visualizer


class MeasurementPipeline:
    """Main workflow entry class for geometry measurement tasks."""

    def __init__(self, config_path: str | None = None) -> None:
        self.config = ConfigLoader().load(config_path)
        self.io = IOHandler()
        self.cad_converter = CADConverter()
        self.noise_adder = GaussianNoiseAdder()
        self.filters = PreprocessingRunner()
        self.registration = PointCloudRegistration()
        self.evaluator = Evaluator()
        self.curvature = CurvatureEstimator()
        self.plane = PlaneEstimator()
        self.visualizer = Visualizer()

    def run_cad2mesh(self, cad_path: str | None = None, mesh_path: str | None = None, precision: float | None = None) -> Path:
        paths = self.config["paths"]
        cad_config = self.config["cad"]
        input_step = self.io.require_cad(cad_path or paths["cad"])
        output_stl = self.io.ensure_parent_dir(mesh_path or paths["output_mesh"] or paths["mesh"])
        return self.cad_converter.convert_step_to_stl(
            input_step=input_step,
            output_stl=output_stl,
            precision=float(precision if precision is not None else cad_config["precision"]),
        )

    def run_denoise(
        self,
        pointcloud_path: str | None = None,
        noise_std: float | None = None,
        no_visualize: bool = False,
    ) -> dict[str, Any]:
        path = pointcloud_path or self.config["paths"]["pointcloud"]
        clean = self.io.load_point_cloud(path)
        noise_config = self.config["noise"]
        adder = GaussianNoiseAdder(
            noise_std=float(noise_std if noise_std is not None else noise_config["std"]),
            seed=noise_config.get("seed"),
        )
        noisy = adder.apply(clean)
        noise_level = adder.noise_level(clean, noisy)

        filter_results = {}
        if self.config["filters"].get("enabled", True):
            filter_results = self.filters.run(noisy, self.config["filters"])

        print(f"Clean points: {len(clean.points)}")
        print(f"Noisy points: {len(noisy.points)}")
        print(f"Mean noise offset: {noise_level:.6f}")
        for name, result in filter_results.items():
            cd, _, _ = self.evaluator.chamfer_distance(clean, result.pointcloud)
            result.metrics["chamfer_distance"] = cd
            print(f"{name}: keep={result.metrics['keep_ratio']:.2f}% CD={cd:.6f}")

        cd_noisy, _, _ = self.evaluator.chamfer_distance(clean, noisy)
        print(f"Noisy without filtering: CD={cd_noisy:.6f}")

        if self._should_visualize(no_visualize):
            first_result = next(iter(filter_results.values()), None)
            self.visualizer.show_denoise(
                clean,
                noisy,
                first_result.pointcloud if first_result else None,
                first_result.name if first_result else "Filtered",
            )

        return {
            "clean": clean,
            "noisy": noisy,
            "noise_level": noise_level,
            "filter_results": filter_results,
            "noisy_chamfer_distance": cd_noisy,
        }

    def run_registration(
        self,
        pointcloud_path: str | None = None,
        mesh_path: str | None = None,
        voxel_size: float | None = None,
        roll: float | None = None,
        pitch: float | None = None,
        yaw: float | None = None,
        trans_range: float | None = None,
        no_visualize: bool = False,
    ) -> dict[str, Any]:
        paths = self.config["paths"]
        reg_config = self.config["registration"]
        source = self.io.load_point_cloud(pointcloud_path or paths["noisy_pointcloud"] or paths["pointcloud"])
        mesh_file = str(self.io.require_file(mesh_path or paths["mesh"], "Mesh"))
        return self._run_registration_with_pointcloud(
            source=source,
            mesh_file=mesh_file,
            voxel_size=voxel_size,
            roll=roll,
            pitch=pitch,
            yaw=yaw,
            trans_range=trans_range,
            no_visualize=no_visualize,
            reg_config=reg_config,
        )

    def _run_registration_with_pointcloud(
        self,
        source: Any,
        mesh_file: str,
        voxel_size: float | None,
        roll: float | None,
        pitch: float | None,
        yaw: float | None,
        trans_range: float | None,
        no_visualize: bool,
        reg_config: dict[str, Any],
    ) -> dict[str, Any]:

        result = self.registration.run_simulated_registration(
            source=source,
            voxel_size=float(voxel_size if voxel_size is not None else reg_config["voxel_size"]),
            roll=float(roll if roll is not None else reg_config["roll"]),
            pitch=float(pitch if pitch is not None else reg_config["pitch"]),
            yaw=float(yaw if yaw is not None else reg_config["yaw"]),
            trans_range=float(trans_range if trans_range is not None else reg_config["trans_range"]),
            random_seed=reg_config.get("random_seed"),
        )

        points = self.io.point_cloud_to_numpy(result.recovered)
        mesh = self.io.load_mesh(mesh_file)
        _, distance_stats = self.evaluator.point_to_mesh_distance_torch(
            points,
            mesh,
            chunk_size=int(self.config["evaluation"]["chunk_size"]),
            triangle_chunk_size=int(self.config["evaluation"]["triangle_chunk_size"]),
        )

        print("Registration result")
        print(f"Fitness: {result.fitness:.4f}")
        print(f"RMSE: {result.rmse:.6f}")
        print(f"Transformation:\n{result.transformation}")
        self._print_distance_stats(distance_stats)

        if self._should_visualize(no_visualize):
            self.visualizer.show_registration(source, result.target, result.recovered)

        return {
            "source": source,
            "registration": result,
            "distance_stats": distance_stats,
        }

    def run_contour(
        self,
        pointcloud_path: str | None = None,
        cad_path: str | None = None,
        mesh_path: str | None = None,
        mesh_deflection: float | None = None,
        point_chunk_size: int | None = None,
        triangle_chunk_size: int | None = None,
        save_path: str | None = None,
        excel_output_path: str | None = None,
        no_visualize: bool = False,
    ) -> dict[str, Any]:
        paths = self.config["paths"]
        eval_config = self.config["evaluation"]

        if mesh_path:
            pointcloud = self.io.load_point_cloud(pointcloud_path or paths["pointcloud"])
            mesh = self.io.load_mesh(mesh_path)
            points = self.io.point_cloud_to_numpy(pointcloud)
            distances, stats = self.evaluator.point_to_mesh_distance_torch(
                points,
                mesh,
                chunk_size=int(point_chunk_size if point_chunk_size is not None else eval_config["chunk_size"]),
                triangle_chunk_size=int(
                    triangle_chunk_size if triangle_chunk_size is not None else eval_config["triangle_chunk_size"]
                ),
            )
            self._print_distance_stats(stats)
            return {"pointcloud": pointcloud, "mesh": mesh, "distances": distances, "distance_stats": stats}

        return self.run_contour_distance(
            cad_path=cad_path or paths["cad"],
            pointcloud_path=pointcloud_path or paths["pointcloud"],
            mesh_deflection=mesh_deflection,
            point_chunk_size=point_chunk_size,
            triangle_chunk_size=triangle_chunk_size,
            save_path=save_path,
            excel_output_path=excel_output_path,
            no_visualize=no_visualize,
        )

    def run_contour_distance(
        self,
        cad_path: str | None = None,
        pointcloud_path: str | None = None,
        mesh_deflection: float | None = None,
        point_chunk_size: int | None = None,
        triangle_chunk_size: int | None = None,
        save_path: str | None = None,
        excel_output_path: str | None = None,
        no_visualize: bool = False,
    ) -> dict[str, Any]:
        paths = self.config["paths"]
        eval_config = self.config["evaluation"]
        input_step = self.io.require_cad(cad_path or paths["cad"])
        pointcloud = self.io.load_point_cloud(pointcloud_path or paths["noisy_pointcloud"] or paths["pointcloud"])
        points = self.io.point_cloud_to_numpy(pointcloud)

        vertices, faces = self.cad_converter.step_to_numpy_mesh(
            input_step,
            mesh_deflection=float(
                mesh_deflection if mesh_deflection is not None else eval_config["cad_mesh_deflection"]
            ),
        )
        print(f"CAD mesh vertices: {vertices.shape[0]}")
        print(f"CAD mesh triangles: {faces.shape[0]}")
        distances, stats = self.evaluator.point_to_triangles_distance_torch(
            points,
            vertices,
            faces,
            point_chunk_size=int(point_chunk_size if point_chunk_size is not None else eval_config["chunk_size"]),
            triangle_chunk_size=int(
                triangle_chunk_size if triangle_chunk_size is not None else eval_config["triangle_chunk_size"]
            ),
        )
        self._print_distance_stats(stats)

        heatmap = self.evaluator.colorize_distances(points, distances)
        output = save_path
        if output:
            self.io.save_point_cloud(heatmap, output)
            print(f"Saved contour heatmap: {output}")

        excel_output = None
        if excel_output_path:
            excel_output = self._save_contour_excel(
                excel_output_path,
                cad_path=input_step,
                stats=stats,
            )
            print(f"Saved contour Excel: {excel_output}")

        if self._should_visualize(no_visualize):
            self.visualizer.show_curvature(heatmap, window_name="Point to CAD Distance Heatmap")

        return {
            "pointcloud": heatmap,
            "distances": distances,
            "distance_stats": stats,
            "cad_vertices": vertices,
            "cad_faces": faces,
            "excel_output": excel_output,
        }

    def run_contour_batch(
        self,
        task_ids: str | list[int | str] | None = None,
        all_tasks: bool = False,
        data_root: str | Path = "datasets",
        pointcloud_kind: str = "noise",
        mesh_deflection: float | None = None,
        point_chunk_size: int | None = None,
        triangle_chunk_size: int | None = None,
        save_path: str | None = None,
        excel_output_path: str | None = None,
        no_visualize: bool = False,
    ) -> dict[str, Any]:
        data_root = Path(data_root)
        task_id_list = self._resolve_contour_task_ids(task_ids, all_tasks, data_root)
        rows: list[dict[str, Any]] = []
        results: dict[str, Any] = {}

        for task_id in task_id_list:
            cad_path = data_root / "CAD" / f"{task_id}.step"
            pointcloud_path = self._contour_pointcloud_path(data_root, task_id, pointcloud_kind)
            task_save_path = self._contour_task_save_path(save_path, task_id)

            print(f"===== contour task {task_id} =====")
            missing = [str(path) for path in (cad_path, pointcloud_path) if not path.exists()]
            if missing:
                print("Missing input: " + "; ".join(missing))
                rows.append({"task_id": task_id, "max": None, "mean": None, "status": "失败"})
                continue

            try:
                result = self.run_contour_distance(
                    cad_path=cad_path,
                    pointcloud_path=pointcloud_path,
                    mesh_deflection=mesh_deflection,
                    point_chunk_size=point_chunk_size,
                    triangle_chunk_size=triangle_chunk_size,
                    save_path=str(task_save_path) if task_save_path else None,
                    excel_output_path=None,
                    no_visualize=True if len(task_id_list) > 1 else no_visualize,
                )
            except Exception as exc:
                print(f"Task {task_id} failed: {exc}")
                rows.append({"task_id": task_id, "max": None, "mean": None, "status": "失败"})
                continue

            stats = result["distance_stats"]
            rows.append({"task_id": task_id, "max": stats.max, "mean": stats.mean, "status": "成功"})
            results[task_id] = result

        excel_output = None
        if excel_output_path:
            excel_output = self._save_contour_excel(excel_output_path, rows=rows)
            print(f"Saved contour Excel: {excel_output}")

        return {"rows": rows, "results": results, "excel_output": excel_output}

    def run_curvature(
        self,
        pointcloud_path: str | None = None,
        radius: float | None = None,
        max_nn: int | None = None,
        color_scale: float | None = None,
        save_path: str | None = None,
        no_visualize: bool = False,
    ) -> dict[str, Any]:
        config = self.config["curvature"]
        pointcloud = self.io.load_point_cloud(pointcloud_path or self.config["paths"]["pointcloud"])
        curvatures = self.curvature.compute(
            pointcloud,
            radius=float(radius if radius is not None else config["radius"]),
            max_nn=int(max_nn if max_nn is not None else config["max_nn"]),
        )
        colored = self.curvature.colorize(
            pointcloud,
            curvatures,
            color_scale=float(color_scale if color_scale is not None else config["color_scale"]),
        )
        output = save_path or self.config["paths"].get("curvature_output")
        if output:
            self.io.save_point_cloud(colored, output)
            print(f"Saved curvature point cloud: {output}")
        print(f"Curvature mean={np.mean(curvatures):.8f}, max={np.max(curvatures):.8f}")
        if self._should_visualize(no_visualize):
            self.visualizer.show_curvature(colored)
        return {"pointcloud": colored, "curvatures": curvatures}

    def run_plane(
        self,
        pointcloud_path: str | None = None,
        threshold_mm: float | None = None,
        max_iterations: int | None = None,
        no_visualize: bool = False,
    ) -> dict[str, Any]:
        config = self.config["plane"]
        if pointcloud_path:
            pointcloud = self.io.load_point_cloud(pointcloud_path)
        else:
            pointcloud = self.plane.create_square_plane_with_noise(
                size_mm=float(config["size_mm"]),
                noise_std_mm=float(config["noise_std_mm"]),
                num_points=int(config["num_points"]),
            )

        result = self.plane.estimate(
            pointcloud,
            threshold_mm=float(threshold_mm if threshold_mm is not None else config["threshold_mm"]),
            max_iterations=int(max_iterations if max_iterations is not None else config["max_iterations"]),
        )

        print(f"Plane model: {result.model}")
        print(f"Inlier ratio: {result.inlier_ratio:.2f}%")
        print(f"Normal angle error: {result.normal_angle_deg:.6f} deg")
        print(f"Plane RMSE: {result.rmse:.8f} m")
        if self._should_visualize(no_visualize):
            self.visualizer.show_plane(result)
        return {"plane_result": result}

    def run_full(self, no_visualize: bool = False) -> dict[str, Any]:
        outputs: dict[str, Any] = {}
        if self.config["cad"].get("enabled", False):
            outputs["cad2mesh"] = self.run_cad2mesh()
        if self.config["noise"].get("enabled", True):
            outputs["denoise"] = self.run_denoise(no_visualize=True)
            mesh_file = str(self.io.require_file(self.config["paths"]["mesh"], "Mesh"))
            outputs["registration"] = self._run_registration_with_pointcloud(
                source=outputs["denoise"]["noisy"],
                mesh_file=mesh_file,
                voxel_size=None,
                roll=None,
                pitch=None,
                yaw=None,
                trans_range=None,
                no_visualize=no_visualize,
                reg_config=self.config["registration"],
            )
        else:
            registration_input = self.config["paths"].get("noisy_pointcloud") or self.config["paths"]["pointcloud"]
            outputs["registration"] = self.run_registration(pointcloud_path=registration_input, no_visualize=no_visualize)
        outputs["contour"] = self.run_contour()
        return outputs

    def _should_visualize(self, no_visualize: bool) -> bool:
        return bool(self.config["visualization"].get("enabled", True)) and not no_visualize

    def _print_distance_stats(self, stats: Any) -> None:
        print("Distance stats")
        print(f"Mean: {stats.mean:.8f}")
        print(f"Max: {stats.max:.8f}")
        print(f"RMSE: {stats.rmse:.8f}")
        print(f"Min: {stats.min:.8f}")
        print(f"Std: {stats.std:.8f}")

    def _save_contour_excel(
        self,
        output_path: str | Path,
        *,
        rows: list[dict[str, Any]] | None = None,
        cad_path: str | Path | None = None,
        stats: Any | None = None,
    ) -> Path:
        output = self._normalize_excel_output_path(output_path)
        if rows is None:
            if cad_path is None or stats is None:
                raise ValueError("cad_path and stats are required when rows is not provided.")
            rows = [{"task_id": Path(cad_path).stem, "max": stats.max, "mean": stats.mean, "status": "成功"}]

        table = [
            [
                "任务号",
                "最大值",
                "平均值",
                "状态",
            ]
        ]
        for row in rows:
            table.append([
                row.get("task_id"),
                row.get("max"),
                row.get("mean"),
                row.get("status"),
            ])
        self._write_xlsx(table, output, sheet_name="contour")
        return output

    def _resolve_contour_task_ids(
        self,
        task_ids: str | list[int | str] | None,
        all_tasks: bool,
        data_root: Path,
    ) -> list[str]:
        if all_tasks:
            cad_dir = data_root / "CAD"
            if not cad_dir.exists():
                raise FileNotFoundError(f"CAD directory not found: {cad_dir}")
            return sorted(path.stem for path in cad_dir.glob("*.step"))

        if task_ids is None:
            raise ValueError("Batch contour requires --tasks or --all.")

        if isinstance(task_ids, str):
            values = [item.strip() for item in task_ids.split(",")]
        else:
            values = [str(item).strip() for item in task_ids]
        return [item for item in values if item]

    def _contour_pointcloud_path(self, data_root: Path, task_id: str, pointcloud_kind: str) -> Path:
        if pointcloud_kind == "clean":
            return data_root / "PointCloud" / f"{task_id}.ply"
        if pointcloud_kind == "noise":
            return data_root / "Noise_PointCloud" / f"{task_id}_noisy.ply"
        raise ValueError("pointcloud_kind must be 'noise' or 'clean'.")

    def _contour_task_save_path(self, save_path: str | None, task_id: str) -> Path | None:
        if not save_path:
            return None

        output = Path(save_path)
        if output.suffix.lower() != ".ply":
            output.mkdir(parents=True, exist_ok=True)
            return output / f"{task_id}_contour.ply"

        output.parent.mkdir(parents=True, exist_ok=True)
        if "{task_id}" in output.name:
            return output.with_name(output.name.replace("{task_id}", task_id))
        return output.with_name(f"{output.stem}_{task_id}{output.suffix}")

    def run_translate_align_pointcloud_to_cad(
        self,
        cad_path: str | None = None,
        pointcloud_path: str | None = None,
        output_dir: str = "results",
        axis: str = "x",
        crop_ratio: float = 0.01,
        mesh_deflection: float = 0.1,
        cad_sample_points: int = 30000,
        local_sample_points: int = 1000,
        scale_pointcloud: float = 1.0,
        no_visualize: bool = False,
    ) -> dict[str, Any]:
        """Run the archived front/tail centroid translation alignment workflow."""
        paths = self.config["paths"]
        input_step = self.io.require_cad(cad_path or paths["cad"])
        input_pointcloud = self._resolve_translate_align_pointcloud(
            pointcloud_path or paths["noisy_pointcloud"] or paths["pointcloud"]
        )

        module = self._load_archive_module("translate_align_pointcloud_to_cad.py")
        return module.translate_align_by_front_tail(
            cad_path=str(input_step),
            pointcloud_path=str(input_pointcloud),
            output_dir=output_dir,
            axis=axis,
            crop_ratio=float(crop_ratio),
            mesh_deflection=float(mesh_deflection),
            cad_sample_points=int(cad_sample_points),
            local_sample_points=int(local_sample_points),
            scale_pointcloud=float(scale_pointcloud),
            visualize=not no_visualize,
        )

    def _resolve_translate_align_pointcloud(self, pointcloud_path: str | Path) -> Path:
        path = Path(pointcloud_path)
        if path.is_dir():
            preferred = path / "shifted_whole_pointcloud.ply"
            if preferred.exists():
                return preferred

            candidates = sorted(path.glob("*.ply"))
            if len(candidates) == 1:
                return candidates[0]
            if not candidates:
                raise FileNotFoundError(f"No PLY point cloud found in directory: {path}")
            names = ", ".join(candidate.name for candidate in candidates)
            raise ValueError(f"Multiple PLY files found in {path}; specify one file explicitly. Found: {names}")

        return self.io.require_file(path, "PointCloud")

    def _load_archive_module(self, filename: str) -> Any:
        script_path = Path(__file__).resolve().parent / "archive" / filename
        if not script_path.exists():
            raise FileNotFoundError(f"Archive script not found: {script_path}")

        module_name = f"_archive_{script_path.stem}"
        spec = importlib.util.spec_from_file_location(module_name, script_path)
        if spec is None or spec.loader is None:
            raise ImportError(f"Cannot load archive script: {script_path}")

        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)
        return module

    def _normalize_excel_output_path(self, output_path: str | Path) -> Path:
        output = Path(output_path)
        if output.suffix.lower() != ".xlsx":
            output = output / "contour.xlsx"
        elif output.parent == Path("."):
            output = Path("excel") / output
        output.parent.mkdir(parents=True, exist_ok=True)
        return output

    def _write_xlsx(self, rows: list[list[Any]], output_path: Path, sheet_name: str) -> None:
        sheet_rows = []
        for row_index, row_values in enumerate(rows, start=1):
            cells = []
            for col_index, value in enumerate(row_values, start=1):
                cell_ref = self._xlsx_cell_ref(row_index, col_index)
                if isinstance(value, (int, float, np.integer, np.floating)) and not isinstance(value, bool):
                    cells.append(f'<c r="{cell_ref}"><v>{value}</v></c>')
                else:
                    text = escape("" if value is None else str(value))
                    cells.append(f'<c r="{cell_ref}" t="inlineStr"><is><t>{text}</t></is></c>')
            sheet_rows.append(f'<row r="{row_index}">{"".join(cells)}</row>')

        workbook_xml = (
            '<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
            '<workbook xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main" '
            'xmlns:r="http://schemas.openxmlformats.org/officeDocument/2006/relationships">'
            f'<sheets><sheet name="{escape(sheet_name[:31])}" sheetId="1" r:id="rId1"/></sheets>'
            '</workbook>'
        )
        worksheet_xml = (
            '<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
            '<worksheet xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main">'
            '<sheetData>'
            f'{"".join(sheet_rows)}'
            '</sheetData>'
            '</worksheet>'
        )
        content_types_xml = (
            '<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
            '<Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types">'
            '<Default Extension="rels" ContentType="application/vnd.openxmlformats-package.relationships+xml"/>'
            '<Default Extension="xml" ContentType="application/xml"/>'
            '<Override PartName="/xl/workbook.xml" '
            'ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet.main+xml"/>'
            '<Override PartName="/xl/worksheets/sheet1.xml" '
            'ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.worksheet+xml"/>'
            '</Types>'
        )
        root_rels_xml = (
            '<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
            '<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">'
            '<Relationship Id="rId1" '
            'Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/officeDocument" '
            'Target="xl/workbook.xml"/>'
            '</Relationships>'
        )
        workbook_rels_xml = (
            '<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
            '<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">'
            '<Relationship Id="rId1" '
            'Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/worksheet" '
            'Target="worksheets/sheet1.xml"/>'
            '</Relationships>'
        )
        with ZipFile(output_path, "w", ZIP_DEFLATED) as xlsx:
            xlsx.writestr("[Content_Types].xml", content_types_xml)
            xlsx.writestr("_rels/.rels", root_rels_xml)
            xlsx.writestr("xl/workbook.xml", workbook_xml)
            xlsx.writestr("xl/_rels/workbook.xml.rels", workbook_rels_xml)
            xlsx.writestr("xl/worksheets/sheet1.xml", worksheet_xml)

    def _xlsx_cell_ref(self, row_index: int, col_index: int) -> str:
        letters = ""
        col = col_index
        while col:
            col, remainder = divmod(col - 1, 26)
            letters = chr(65 + remainder) + letters
        return f"{letters}{row_index}"