import copy
import json
from pathlib import Path
import tempfile
import unittest

import numpy as np
from PIL import Image

from sam_segment import choose_candidate, validate_prompts
from segmentation_outputs import assess_candidate, render_outputs, write_report


class SegmentationTests(unittest.TestCase):
    def setUp(self):
        self.rgba = np.zeros((24, 32, 4), dtype=np.uint8)
        self.rgba[4:20, 5:27] = [120, 50, 200, 255]
        self.rgba[5, 6, 3] = 80
        self.image = Image.fromarray(self.rgba)
        self.mask = self.rgba[:, :, 3] > 0

    def test_background_leak_measured_before_clipping(self):
        raw = np.ones((24, 32), dtype=bool)
        qa = assess_candidate(raw, self.rgba[:, :, 3], [5, 4, 27, 20], [[6, 5]])
        self.assertGreater(qa['outside_alpha_ratio'], .5)
        self.assertEqual(qa['seed_hit_fraction'], 1)
        self.assertEqual(qa['area_px'], 16 * 22)

    def test_rank_prefers_local_mask_and_never_invents_empty_mask(self):
        wrong = self.mask.copy()
        local = np.zeros_like(wrong)
        local[8:13, 8:13] = True
        best, _ = choose_candidate([wrong, local], [.95, .7], self.rgba[:, :, 3], [8, 8, 13, 13], [[10, 10]])
        self.assertEqual(best['index'], 1)
        best, _ = choose_candidate([np.zeros_like(local)], [.99], self.rgba[:, :, 3], [8, 8, 13, 13], [[10, 10]])
        self.assertIsNone(best)

    def test_prompts_reject_transparent_points_mismatched_canvas_and_duplicate_ids(self):
        plan = {'parts': [{'id': 'hand', 'visibility': 'visible'}]}
        prompts = {'width': 32, 'height': 24, 'parts': [{'id': 'hand', 'status': 'candidate',
                   'seed_source': 'alpha_roi_distance_transform', 'bbox_px_xyxy': [0, 0, 32, 24],
                   'positive_points_px': [[10, 10]]}]}
        validate_prompts(plan, prompts, self.image)
        for variant in range(3):
            bad = copy.deepcopy(prompts)
            if variant == 0:
                bad['parts'][0]['positive_points_px'] = [[0, 0]]
            elif variant == 1:
                bad['width'] = 99
            else:
                bad['parts'].append(copy.deepcopy(bad['parts'][0]))
            with self.assertRaises(ValueError):
                validate_prompts(plan, bad, self.image)

    def test_render_preserves_cutout_pixels_alpha_and_measures_overlap(self):
        parts = [{'id': 'hand', 'semantic': 'hand', 'mask': self.mask, 'model_score': .9},
                 {'id': 'forearm', 'semantic': 'forearm', 'mask': self.mask, 'model_score': .8}]
        with tempfile.TemporaryDirectory() as directory:
            result = render_outputs(self.image, parts, directory)
            self.assertEqual(result['coverage']['coverage_ratio'], 1)
            self.assertEqual(result['coverage']['overlap_pixels'], int(self.mask.sum()))
            self.assertEqual(result['near_duplicate_masks'][0]['iou'], 1)
            self.assertIn('NEAR_DUPLICATE_MASK', result['parts'][0]['warnings'])
            cut = np.array(Image.open(Path(directory) / 'parts/hand.png'))
            np.testing.assert_array_equal(cut, self.rgba[4:20, 5:27])
            self.assertEqual(Image.open(Path(directory) / 'annotated.png').size, self.image.size)
            sheet = np.array(Image.open(Path(directory) / 'parts-sheet.png'))
            self.assertTrue(np.all(sheet[0, 0] == 255))
            a, b = [p['sheet_xyxy'] for p in result['parts']]
            self.assertTrue(a[2] <= b[0] or b[2] <= a[0] or a[3] <= b[1] or b[3] <= a[1])
            result['skipped'] = []
            write_report(directory, result)
            self.assertTrue((Path(directory) / 'report.html').exists())
            json.dumps(result)

    def test_renderer_rejects_empty_masks_and_unsafe_ids(self):
        with tempfile.TemporaryDirectory() as directory:
            for name, mask in [('hand', np.zeros_like(self.mask)), ('../bad', self.mask)]:
                with self.assertRaises(ValueError):
                    render_outputs(self.image, [{'id': name, 'mask': mask}], directory)


if __name__ == '__main__':
    unittest.main()