"""End-to-end tests for KindredSolver using internal API.

These tests exercise the solver bridge without requiring kcsolve pybind11
bindings (which need a FreeCAD build). Instead we test the internal
pipeline directly.
"""

import math

import pytest
from kindred_solver.constraints import (
    CoincidentConstraint,
    DistancePointPointConstraint,
    FixedConstraint,
)
from kindred_solver.dof import count_dof
from kindred_solver.entities import RigidBody
from kindred_solver.newton import newton_solve
from kindred_solver.params import ParamTable
from kindred_solver.prepass import single_equation_pass, substitution_pass


def _solve_system(bodies, constraint_objs):
    """Helper: run the full solve pipeline on a list of bodies + constraints."""
    params = bodies[0].tx  # hack to get the shared ParamTable
    # Actually, we need the param table. Let's require it as arg.
    raise NotImplementedError


class TestCoincidentSolve:
    def test_two_bodies_coincident(self):
        """Ground body A at origin, body B at (10,0,0).
        Coincident constraint → B moves to origin."""
        pt = ParamTable()
        body_a = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        body_b = RigidBody("b", pt, (10, 0, 0), (1, 0, 0, 0))

        c = CoincidentConstraint(body_a, (0, 0, 0), body_b, (0, 0, 0))

        residuals = list(c.residuals())
        residuals.append(body_b.quat_norm_residual())
        quat_groups = [body_b.quat_param_names()]

        residuals = substitution_pass(residuals, pt)
        residuals = single_equation_pass(residuals, pt)

        converged = newton_solve(residuals, pt, quat_groups=quat_groups)
        assert converged

        env = pt.get_env()
        pos = body_b.extract_position(env)
        assert abs(pos[0]) < 1e-8
        assert abs(pos[1]) < 1e-8
        assert abs(pos[2]) < 1e-8

    def test_coincident_with_markers(self):
        """Body A at origin, body B at (10,0,0).
        Marker on A at local (5,0,0), marker on B at local (-5,0,0).
        After solve, B should be at (10,0,0) since 5+0 == 10-5."""
        pt = ParamTable()
        body_a = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        body_b = RigidBody("b", pt, (20, 0, 0), (1, 0, 0, 0))

        c = CoincidentConstraint(body_a, (5, 0, 0), body_b, (-5, 0, 0))

        residuals = list(c.residuals())
        residuals.append(body_b.quat_norm_residual())
        quat_groups = [body_b.quat_param_names()]

        residuals = substitution_pass(residuals, pt)
        converged = newton_solve(residuals, pt, quat_groups=quat_groups)
        assert converged

        env = pt.get_env()
        pos = body_b.extract_position(env)
        assert abs(pos[0] - 10.0) < 1e-8


class TestDistanceSolve:
    def test_distance_from_origin(self):
        """Ground A at origin, B starts at (1,0,0).
        Distance constraint d=5 → B moves to (5,0,0)."""
        pt = ParamTable()
        body_a = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        body_b = RigidBody("b", pt, (1, 0, 0), (1, 0, 0, 0))

        c = DistancePointPointConstraint(
            body_a,
            (0, 0, 0),
            body_b,
            (0, 0, 0),
            5.0,
        )

        residuals = list(c.residuals())
        residuals.append(body_b.quat_norm_residual())
        quat_groups = [body_b.quat_param_names()]

        residuals = substitution_pass(residuals, pt)
        converged = newton_solve(residuals, pt, quat_groups=quat_groups)
        assert converged

        env = pt.get_env()
        pos = body_b.extract_position(env)
        dist = math.sqrt(pos[0] ** 2 + pos[1] ** 2 + pos[2] ** 2)
        assert abs(dist - 5.0) < 1e-8


class TestFixedSolve:
    def test_fixed_weld(self):
        """Ground A at origin, B at (5,0,0) with identity rotation.
        Fixed constraint with identity markers → B stays at same pose as A."""
        pt = ParamTable()
        body_a = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        body_b = RigidBody("b", pt, (5, 0, 0), (1, 0, 0, 0))

        c = FixedConstraint(
            body_a,
            (0, 0, 0),
            (1, 0, 0, 0),
            body_b,
            (0, 0, 0),
            (1, 0, 0, 0),
        )

        residuals = list(c.residuals())
        residuals.append(body_b.quat_norm_residual())
        quat_groups = [body_b.quat_param_names()]

        residuals = substitution_pass(residuals, pt)
        converged = newton_solve(residuals, pt, quat_groups=quat_groups)
        assert converged

        env = pt.get_env()
        pos = body_b.extract_position(env)
        assert abs(pos[0]) < 1e-8
        assert abs(pos[1]) < 1e-8
        assert abs(pos[2]) < 1e-8

        quat = body_b.extract_quaternion(env)
        assert abs(quat[0] - 1.0) < 1e-8  # w
        assert abs(quat[1]) < 1e-8  # x
        assert abs(quat[2]) < 1e-8  # y
        assert abs(quat[3]) < 1e-8  # z


class TestDOFCounting:
    def test_single_body_unconstrained(self):
        """One non-grounded body: 7 params - 1 quat norm = 6 DOF."""
        pt = ParamTable()
        body = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0))
        residuals = [body.quat_norm_residual()]
        dof = count_dof(residuals, pt)
        assert dof == 6

    def test_two_bodies_coincident(self):
        """Grounded A + free B with coincident: 7 - (3 coincident + 1 quat) = 3 DOF."""
        pt = ParamTable()
        body_a = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        body_b = RigidBody("b", pt, (0, 0, 0), (1, 0, 0, 0))
        c = CoincidentConstraint(body_a, (0, 0, 0), body_b, (0, 0, 0))
        residuals = list(c.residuals()) + [body_b.quat_norm_residual()]
        residuals = substitution_pass(residuals, pt)
        dof = count_dof(residuals, pt)
        assert dof == 3  # 3 rotation DOF remain

    def test_two_bodies_fixed(self):
        """Grounded A + free B with fixed: 7 - (6 fixed + 1 quat) = 0 DOF."""
        pt = ParamTable()
        body_a = RigidBody("a", pt, (0, 0, 0), (1, 0, 0, 0), grounded=True)
        body_b = RigidBody("b", pt, (0, 0, 0), (1, 0, 0, 0))
        c = FixedConstraint(
            body_a,
            (0, 0, 0),
            (1, 0, 0, 0),
            body_b,
            (0, 0, 0),
            (1, 0, 0, 0),
        )
        residuals = list(c.residuals()) + [body_b.quat_norm_residual()]
        residuals = substitution_pass(residuals, pt)
        dof = count_dof(residuals, pt)
        assert dof == 0