import json
import threading
from types import SimpleNamespace
import numpy as np
import pytest

pytest.importorskip('small_gicp')
from k1link.missions.registration import register, transform, rigid, path_hint, angle_deg
from k1link.device_plugins.xgrids_k1.localization_source import extract_submap
from k1link.device_plugins.xgrids_k1.planning_source import export_planning_source
from k1link.missions.registration_runs import RegistrationRuns
from test_stream_summary import _write_capture, _pcl_payload, _pose_payload


def geometry():
    r = np.random.default_rng(41)
    return np.vstack([r.normal([0,0,0],[4,3,.1],(3000,3)),
                      r.normal([3,2,2],[.1,2,2],(1500,3)),
                      r.normal([-2,-1,2],[2,.1,1],(1500,3))])


def test_recovers_known_rigid_transform_in_source_to_target_convention():
    p = geometry(); t = np.eye(4); a = .1
    t[:2,:2] = [[np.cos(a),-np.sin(a)],[np.sin(a),np.cos(a)]]
    t[:3,3] = [.7,-.4,.2]
    result = register(p, transform(p, np.linalg.inv(t)), np.eye(4))
    found = np.array(result['T_reference_query'])
    assert result['status'] == 'candidate'
    assert np.linalg.norm(found[:3,3]-t[:3,3]) < .01
    assert angle_deg(found[:3,:3] @ t[:3,:3].T) < .1
    assert not result['localization_confirmed'] and not result['vehicle_control']


def test_no_overlap_and_uninformative_plane_are_rejected():
    p = geometry()
    assert register(p, p+[100,100,100], np.eye(4))['status'] == 'rejected'
    p[:,2] = 0
    result = register(p, p, np.eye(4))
    assert result['status'] == 'rejected' and result['shape_ratio'] == 0


def test_nonrigid_and_nonfinite_inputs_rejected():
    with pytest.raises(ValueError): rigid(np.ones((4,4)))
    p = geometry(); p[0,0] = np.nan
    with pytest.raises(ValueError): register(p, geometry(), np.eye(4))
    with pytest.raises(ValueError): register(geometry()[:20], geometry(), np.eye(4))


def test_path_hint_maps_query_entry_and_heading_to_selected_route():
    t = path_hint([[8,9,1],[8,14,1]], [[-1,-2,0],[4,-2,0]])
    assert np.allclose(transform(np.array([[-1,-2,0],[4,-2,0]]),t), [[8,9,1],[8,14,1]])
    with pytest.raises(ValueError): path_hint([[0,0,0],[0,0,0]], [[0,0,0],[4,0,0]])


def test_extractor_uses_only_selected_interval_and_no_second_pose_transform(tmp_path):
    raw = tmp_path/'mqtt.raw.k1mqtt'
    _write_capture(raw, [('x/lio_pose', _pose_payload((5,0,0))),
                         ('x/lio_pcl', _pcl_payload(scaler=1000, point_count=4)),
                         ('x/lio_pose', _pose_payload((8,0,0))),
                         ('x/lio_pcl', b'bad-future-cloud')])
    out = tmp_path/'planning.json'; export_planning_source(raw,out)
    points, meta = extract_submap(raw,json.loads(out.read_text()),0,1)
    assert len(points) >= 1 and np.allclose(points[0], [1,-2,.5])
    assert meta['available_frames'] == 1 and meta['frames'][0]['message_index'] == 1
    assert meta['message_interval'] == [0,2]


def test_run_rejects_overlapping_self_comparison_and_revision_race(tmp_path):
    source = {'poses':[{'position':[i,0,0],'distance_m':i} for i in range(31)]}
    draft = {'id':'a', 'revision':1, 'zone':{'session_id':'A'},
             'route':{'length_m':30,'start_index':0,'end_index':30}}
    drafts = SimpleNamespace(database=tmp_path/'drafts.sqlite', get=lambda _:draft,
                             sources=SimpleNamespace(bound=lambda *_:source))
    runs = RegistrationRuns(drafts)
    req = {'revision':1,'session_id':'A','generation':'x','start_index':0,'end_index':30}
    try:
        with pytest.raises(ValueError, match='не должны пересекаться'): runs.start('a',req)
        with pytest.raises(ValueError, match='изменён'): runs.start('a',{**req,'revision':2})
        assert not runs.lock.locked() and not list(runs.root.glob('*/report.json'))
    finally: runs.close()


def test_run_is_single_owner_persisted_and_bound_to_submitted_revision(tmp_path, monkeypatch):
    import k1link.missions.registration_runs as module
    entered, release = threading.Event(), threading.Event()
    source = {'poses':[{'position':[i,0,0], 'distance_m':i} for i in range(31)]}
    draft = {'id':'draft', 'revision':1, 'zone':{'session_id':'A','generation':'a'},
             'route':{'length_m':30,'start_index':0,'end_index':30,
                      'points':[{'position':[i,0,0]} for i in range(31)]}}
    def submap(*args):
        entered.set(); assert release.wait(5)
        return geometry(), {'session_id':args[0]}
    drafts = SimpleNamespace(database=tmp_path/'db', get=lambda _:draft,
        sources=SimpleNamespace(bound=lambda *_:source, submap=submap))
    monkeypatch.setattr(module, 'write_scene', lambda p,*a: p.write_bytes(b'test-rrd'))
    runs = RegistrationRuns(drafts)
    req = {'revision':1,'session_id':'B','generation':'b','start_index':0,'end_index':30}
    first = runs.start('draft',req)
    try:
        assert entered.wait(5)
        with pytest.raises(ValueError, match='ещё выполняется'): runs.start('draft',req)
        # Real draft reads are decoded snapshots. A later edit does not relabel an existing report.
        draft = {**draft, 'revision':2}
    finally:
        release.set(); runs.close()
    report = runs.get(first['id'])
    assert report['state'] == 'ready' and report['revision'] == 1
    assert report['reference']['session_id'] == 'A' and report['query']['session_id'] == 'B'
    assert not report['localization_confirmed'] and not report['vehicle_control']
    restarted = RegistrationRuns(drafts)
    try:
        assert restarted.get(first['id']) == report
        assert restarted.list('draft')[0]['id'] == first['id']
    finally: restarted.close()