"""Published automatic fits must bind the real factory and respect saved intent."""

import json
from pathlib import Path
from types import SimpleNamespace
import struct
import unittest

from opendbc.car.car_helpers import interfaces
from opendbc.car.gm.values import CAR as GM
from opendbc.car.hyundai.values import CAR as HKG
from opendbc.car.toyota.values import CAR as TOYOTA
from openpilot.cereal import messaging
from openpilot.common.params import Params
from openpilot.common.prefix import OpenpilotPrefix
from openpilot.selfdrive.controls.lib.latcontrol_torque import LatControlTorque
from openpilot.starpilot.flm.learning_source import (
  AdmittedPolicy, AutomaticLearningConsumer, MAX_AGE_NS, RESET_KEY, SERVICE, conversion_identity, identity_for, learning_payload,
)
from openpilot.starpilot.lateral.controller_selection import ControllerMode, DOCUMENT_KEY, policy_for, replace_mode
from openpilot.starpilot.lateral.torque_runtime import TorqueHost
from openpilot.starpilot.lateral.torque_settings import DOCUMENT_KEY as TORQUE_KEY, replace_field, serialize_document

VEHICLES = (GM.CHEVROLET_BOLT_CC_2018_2021, HKG.HYUNDAI_IONIQ_6, TOYOTA.TOYOTA_COROLLA_TSS2)
NOW = 20_000_000_000


def float32(value):
  """The controller's CarParams torque fields are Float32, not Python doubles."""
  return struct.unpack('<f', struct.pack('<f', value))[0]


def context(vehicle):
  cp = interfaces[vehicle].get_non_essential_params(vehicle).as_reader()
  params = Params()
  if policy_for(cp) is not None:
    params.put(DOCUMENT_KEY, json.loads(replace_mode(None, cp, ControllerMode.STANDARD)), block=True)
  ci = interfaces[vehicle](cp)
  source = AutomaticLearningConsumer(cp, ci, ControllerMode.STANDARD, params,
                                    policies={'test-contract': AdmittedPolicy(.1)})
  return cp, ci, params, source


def wire(cp, params, *, now=NOW, sequence=1, **changes):
  identity, enabled, manual = identity_for(cp, ControllerMode.STANDARD, params, 'opendbc.linear_inverse_v1')
  assert identity is not None and enabled and not manual
  event = messaging.new_message(SERVICE, valid=True)
  state = learning_payload(getattr(event, SERVICE))
  state.version, state.status, state.family = 1, 'qualifiedFit', 'linearFactor'
  state.identity = identity
  state.session, state.sequence = 'consumer-test', sequence
  state.observedMonoTime, state.validUntilMonoTime = now, now + MAX_AGE_NS
  state.sourceWindowStartMonoTime, state.sourceWindowEndMonoTime = now - 10_000_000_000, now - 1_000_000_000
  state.latAccelFactor = cp.lateralTuning.torque.latAccelFactor * 1.05
  state.acceptancePolicyId = 'test-contract'
  state.trainWindows, state.heldoutWindows = 10, 8
  state.trainBaselineLoss, state.trainCandidateLoss = .1, .09
  state.heldoutBaselineLoss, state.heldoutCandidateLoss = .1, .09
  for key, value in changes.items():
    setattr(state, key, value)
  # Exercise the actual wire reader; these declared transport flags are input
  # doubles, not production frequency/IPC qualification.
  reader = messaging.log_from_bytes(event.to_bytes())
  sm = {SERVICE: getattr(reader, SERVICE)}
  return SimpleNamespace(seen={SERVICE: True}, alive={SERVICE: True}, valid={SERVICE: True}, freq_ok={SERVICE: True},
                         recv_time={SERVICE: now / 1e9}), sm


class InputState(SimpleNamespace):
  def __getitem__(self, key):
    return self.states[key]

  def all_checks(self, names):
    return all(self.seen.get(name, False) and self.alive.get(name, False) and
               self.valid.get(name, False) and self.freq_ok.get(name, False) for name in names)


def input_state(cp, params, **kwargs):
  flags, states = wire(cp, params, **kwargs)
  return InputState(**vars(flags), states=states)


class TestAutomaticLearningSource(unittest.TestCase):
  def test_actual_linear_factories_and_slewed_numerical_application_without_manual_document(self):
    for vehicle in VEHICLES:
      with self.subTest(vehicle=vehicle), OpenpilotPrefix():
        cp, ci, params, source = context(vehicle)
        self.assertTrue(source.supported)
        self.assertIsNone(params.get(TORQUE_KEY))
        sm = input_state(cp, params)
        host = TorqueHost(params, cp, allow_learning=False, automatic_interface=ci, automatic_controller=ControllerMode.STANDARD)
        controller = LatControlTorque(cp, ci, .01, controller_mode=ControllerMode.STANDARD)
        previous = controller.torque_from_lateral_accel(1., controller.torque_params)
        for frame in range(20):
          now = NOW + frame * 10_000_000
          tune = source.sample(sm, now_ns=now, lat_active=True)
          self.assertIsNotNone(tune)
          host.apply(controller, host.sample(sm, now_ns=now, lat_active=True, automatic=tune))
        self.assertGreater(controller.torque_params.latAccelFactor, cp.lateralTuning.torque.latAccelFactor)
        self.assertLess(controller.torque_from_lateral_accel(1., controller.torque_params), previous)
        self.assertEqual(controller.torque_params.latAccelOffset, cp.lateralTuning.torque.latAccelOffset)
        self.assertEqual(controller.torque_params.friction, cp.lateralTuning.torque.friction)
        host.apply(controller, host.sample(sm, now_ns=NOW + 210_000_000, lat_active=False))
        self.assertEqual(controller.torque_params.latAccelFactor, cp.lateralTuning.torque.latAccelFactor)

  def test_lazy_adoption_slews_from_applied_native_tune_and_preserves_native_nonfactor_values(self):
    with OpenpilotPrefix():
      cp, ci, params, source = context(HKG.HYUNDAI_IONIQ_6)
      basis = cp.lateralTuning.torque
      controller = LatControlTorque(cp, ci, .01, controller_mode=ControllerMode.STANDARD)
      factor, offset, friction = float32(basis.latAccelFactor * 1.1), float32(.03), float32(basis.friction * 1.1)
      controller.update_torque_parameters(factor, offset, friction)
      host = TorqueHost(params, cp, allow_learning=True, automatic_interface=ci, automatic_controller=ControllerMode.STANDARD)
      host.adopt_linear_controller(controller)
      self.assertFalse(host.apply(controller, host.applied))
      sm = input_state(cp, params)
      for name in ('seen', 'alive', 'valid', 'freq_ok'):
        getattr(sm, name)['lateralTorqueParameters'] = True
      sm.logMonoTime = {'lateralTorqueParameters': NOW}
      sm.recv_time['lateralTorqueParameters'] = NOW / 1e9
      sm.states['lateralTorqueParameters'] = SimpleNamespace(useParams=True, valid=True, version=1, latAccelFactorFiltered=factor,
        latAccelOffsetFiltered=offset, frictionCoefficientFiltered=friction)
      step = .005 * basis.latAccelFactor
      for frame in range(55):
        now = NOW + frame * 10_000_000
        prior = host.applied.lat_accel_factor
        automatic = source.sample(sm, now_ns=now, lat_active=True)
        host.apply(controller, host.sample(sm, now_ns=now, lat_active=True, automatic=automatic))
        self.assertLessEqual(abs(host.applied.lat_accel_factor - prior), step + 1e-12)
        self.assertEqual(controller.torque_params.latAccelFactor, float32(host.applied.lat_accel_factor))
        self.assertEqual(controller.torque_params.latAccelOffset, offset)
        self.assertEqual(controller.torque_params.friction, friction)
        if frame == 0:
          self.assertEqual(controller.torque_params.latAccelFactor, float32(factor - step))
        if frame > 50:
          self.assertIsNone(automatic)
          self.assertGreater(controller.torque_params.latAccelFactor, prior)
      # Missing native updates retain the old owner's scalar base during fallback.
      sm.alive['lateralTorqueParameters'] = False
      prior = host.applied.lat_accel_factor
      host.apply(controller, host.sample(sm, now_ns=NOW + 550_000_000, lat_active=False))
      self.assertLessEqual(abs(host.applied.lat_accel_factor - prior), step + 1e-12)
      self.assertEqual(controller.torque_params.latAccelFactor, float32(host.applied.lat_accel_factor))
      self.assertEqual(controller.torque_params.latAccelOffset, offset)
      self.assertEqual(controller.torque_params.friction, friction)
      # A fresh native update may change its own nonfactor parameters normally.
      sm.alive['lateralTorqueParameters'] = True
      sm.logMonoTime['lateralTorqueParameters'] = NOW + 560_000_000
      sm.recv_time['lateralTorqueParameters'] = (NOW + 560_000_000) / 1e9
      sm.states['lateralTorqueParameters'].latAccelOffsetFiltered = float32(-.02)
      sm.states['lateralTorqueParameters'].frictionCoefficientFiltered = float32(basis.friction * 1.2)
      host.apply(controller, host.sample(sm, now_ns=NOW + 560_000_000, lat_active=False))
      self.assertEqual(controller.torque_params.latAccelOffset, float32(-.02))
      self.assertEqual(controller.torque_params.friction, float32(basis.friction * 1.2))

  def test_wrong_identity_expiry_sequence_and_unadmitted_policy_do_not_apply(self):
    with OpenpilotPrefix():
      cp, _, params, source = context(HKG.HYUNDAI_IONIQ_6)
      sm = input_state(cp, params)
      self.assertIsNotNone(source.sample(sm, now_ns=NOW, lat_active=True))
      wrong = learning_payload(sm[SERVICE]).identity.to_dict()
      wrong['vehicle'] = 'different-vehicle'
      for sequence, overrides in enumerate(({'identity': wrong}, {'version': 2}, {'latAccelFactor': float('nan')}, {'latAccelFactor': 100.},
                        {'acceptancePolicyId': 'unqualified'}, {'trainWindows': 7}, {'heldoutWindows': 7},
                        {'trainCandidateLoss': .1}, {'heldoutCandidateLoss': .1}, {'sequence': 0}, {'validUntilMonoTime': NOW - 1}), start=2):
        with self.subTest(overrides=overrides):
          bad = input_state(cp, params, **({'sequence': sequence} | overrides))
          self.assertIsNone(source.sample(bad, now_ns=NOW, lat_active=True))
      self.assertIsNone(source.sample(sm, now_ns=NOW + MAX_AGE_NS + 1, lat_active=True))

  def test_saved_off_malformed_manual_and_reset_invalidate_existing_result(self):
    for changed_key, raw in (('ForceAutoTuneOff', b'1'), (TORQUE_KEY, b'malformed'), (RESET_KEY, b'1234')):
      with self.subTest(key=changed_key), OpenpilotPrefix():
        cp, _, params, source = context(HKG.HYUNDAI_IONIQ_6)
        sm = input_state(cp, params)
        self.assertIsNotNone(source.sample(sm, now_ns=NOW, lat_active=True))
        Path(params.get_param_path(changed_key)).write_bytes(raw)
        self.assertIsNone(source.sample(sm, now_ns=NOW + 1_000_000_000, lat_active=True))
        self.assertNotEqual(source.identity, learning_payload(sm[SERVICE]).identity.to_dict())

  def test_saved_manual_factor_overrides_automatic_overlay_after_adoption(self):
    with OpenpilotPrefix():
      cp, ci, params, source = context(HKG.HYUNDAI_IONIQ_6)
      basis = cp.lateralTuning.torque
      sm = input_state(cp, params)
      controller = LatControlTorque(cp, ci, .01, controller_mode=ControllerMode.STANDARD)
      controller.update_torque_parameters(basis.latAccelFactor * 1.1, .03, basis.friction * 1.1)
      host = TorqueHost(params, cp, allow_learning=False, automatic_interface=ci, automatic_controller=ControllerMode.STANDARD)
      host.adopt_linear_controller(controller)
      automatic = source.sample(sm, now_ns=NOW, lat_active=True)
      assert automatic is not None
      host.apply(controller, host.sample(sm, now_ns=NOW, lat_active=True, automatic=automatic))
      selected = replace_field({}, str(cp.carFingerprint),
        (basis.latAccelFactor, basis.latAccelOffset, basis.friction), 'factor', 'custom', basis.latAccelFactor * 1.2)
      params.put(TORQUE_KEY, json.loads(serialize_document(selected)), block=True)
      # Refresh discovers the saved manual identity; the old automatic wire cannot match it.
      self.assertIsNone(source.sample(sm, now_ns=NOW + 1_000_000_000, lat_active=True))
      host.sample(sm, now_ns=NOW + 1_000_000_000, lat_active=True, automatic=automatic)
      self.assertAlmostEqual(host.selected.lat_accel_factor, basis.latAccelFactor * 1.2)
      self.assertEqual(host.selected.lat_accel_offset, basis.latAccelOffset)
      self.assertEqual(host.selected.source.value, 'user')
      host.sample(sm, now_ns=NOW + 1_010_000_000, lat_active=True, automatic=automatic)
      self.assertAlmostEqual(host.selected.lat_accel_factor, basis.latAccelFactor * 1.2)

  def test_malformed_source_at_host_refresh_cannot_bypass_with_previously_accepted_auto(self):
    with OpenpilotPrefix():
      cp, ci, params, source = context(GM.CHEVROLET_BOLT_CC_2018_2021)
      sm = input_state(cp, params)
      automatic = source.sample(sm, now_ns=NOW, lat_active=True)
      assert automatic is not None
      for adopted in (False, True):
        with self.subTest(adopted=adopted):
          controller = LatControlTorque(cp, ci, .01, controller_mode=ControllerMode.STANDARD)
          host = TorqueHost(params, cp, allow_learning=False, automatic_interface=ci, automatic_controller=ControllerMode.STANDARD)
          if adopted:
            host.adopt_linear_controller(controller)
          host.sample(sm, now_ns=NOW, lat_active=True, automatic=automatic)
          self.assertTrue(host.settings.valid)  # Genuine absence takes the strict default parser.
          Path(params.get_param_path(TORQUE_KEY)).write_bytes(b'malformed')
          # A due host refresh may precede the consumer's independently phased refresh.
          host.last_refresh_ns = NOW - 1_000_000_000
          selected = host.sample(sm, now_ns=NOW + 10_000_000, lat_active=True, automatic=automatic)
          self.assertFalse(host.settings.valid)
          self.assertEqual(host.selected.lat_accel_factor, cp.lateralTuning.torque.latAccelFactor)
          self.assertNotEqual(selected.lat_accel_factor, automatic.lat_accel_factor)
          params.remove(TORQUE_KEY)

  def test_fixed_custom_and_nonlinear_conversion_are_not_inferred_from_preference(self):
    with OpenpilotPrefix():
      cp, ci, params, _ = context(GM.CHEVROLET_BOLT_CC_2018_2021)
      self.assertIsNone(conversion_identity(cp, ci, ControllerMode.STARPILOT))
      previous_document = Path(params.get_param_path(DOCUMENT_KEY)).read_bytes()
      cp, ci, params, _ = context(GM.CHEVROLET_SILVERADO)
      self.assertEqual(Path(params.get_param_path(DOCUMENT_KEY)).read_bytes(), previous_document)
      self.assertNotIn(str(cp.carFingerprint), json.loads(previous_document)['vehicles'])
      self.assertIsNone(conversion_identity(cp, ci, ControllerMode.STANDARD))
      cp, ci, _, _ = context(HKG.HYUNDAI_IONIQ_6)
      self.assertIsNone(conversion_identity(cp, ci, ControllerMode.STARPILOT))