import ast
from pathlib import Path
from unittest import TestCase

TREE = ast.parse((Path(__file__).resolve().parents[1] / 'handler.py').read_text())


class StructuralControl(TestCase):
    def setUp(self):
        namespace = {}
        definitions = [node for node in TREE.body
                       if isinstance(node, ast.FunctionDef) and node.name in ('_control_strength', '_validate_job_keys')
                       or isinstance(node, ast.Assign) and any(isinstance(t, ast.Name) and
                           (t.id.startswith('MODE_') or t.id in ('CONTROL_LORA_FILES', '_COMMON_KEYS', '_CONDITIONING_KEYS',
                            '_SIZE_KEYS', '_GUIDED_KEYS', '_TWO_STAGE_SIGMA_KEYS', '_MODE_KEYS')) for t in node.targets)]
        exec(compile(ast.Module(body=definitions, type_ignores=[]), '<control>', 'exec'), namespace)
        self.namespace = namespace

    def test_strength_boundaries_and_bad_input(self):
        check = self.namespace['_control_strength']
        self.assertEqual(check({'video_url': 'https://example.test/guide.mp4'}), 1.0)
        for value in (0, .6, 1):
            self.assertEqual(check({'video_url': 'guide', 'conditioning_strength': value}), value)
        for value in (-.1, 1.1, float('nan'), float('inf'), None):
            with self.subTest(value=value), self.assertRaises((ValueError, TypeError)):
                check({'video_url': 'guide', 'conditioning_strength': value})
        for url in (None, '', ' ', 123):
            with self.subTest(url=url), self.assertRaises(ValueError):
                check({'video_url': url})

    def test_control_payload_uses_guide_and_images_without_ignored_guidance(self):
        validate = self.namespace['_validate_job_keys']
        for mode in ('motion', 'union'):
            validate({'mode': mode, 'video_url': 'guide', 'image_url': 'anchor',
                      'conditioning_strength': .8, 'width': 768, 'height': 1344}, mode)
            for ignored in ('negative_prompt', 'num_inference_steps', 'stage_2_sigmas', 'auto_duration'):
                with self.subTest(mode=mode, ignored=ignored), self.assertRaises(ValueError):
                    validate({'mode': mode, ignored: 1}, mode)
        with self.assertRaises(ValueError):
            validate({'conditioning_strength': .6}, 'i2v')