import io
import os
import unittest
import zipfile
from engine import cosine, inspect_content, content_kind

def policy(): return {"dlp":{"maxBytes":1048576,"rules":[{"id":"secret","kind":"literal","value":"Secret Roadmap"}],"semantic":{"enabled":False}}}

class InspectionTests(unittest.TestCase):
    def test_text_and_binary(self):
        self.assertEqual(inspect_content(b"Public information",policy())["decision"],"allow")
        self.assertEqual(inspect_content(b"Secret Roadmap",policy())["ruleId"],"secret")
        self.assertEqual(inspect_content(b"\xff\x00",policy())["decision"],"block")
    def test_office_extraction(self):
        output=io.BytesIO()
        with zipfile.ZipFile(output,"w") as archive:
            archive.writestr("word/document.xml",'<document xmlns="urn:test"><t>Secret Roadmap</t></document>')
        result=inspect_content(output.getvalue(),policy())
        self.assertEqual(result["contentType"],"office");self.assertEqual(result["ruleId"],"secret")
    def test_pdf_extraction(self):
        from pypdf import PdfWriter
        from pypdf.generic import DecodedStreamObject,NameObject,DictionaryObject
        writer=PdfWriter();page=writer.add_blank_page(300,300)
        font=DictionaryObject({NameObject('/Type'):NameObject('/Font'),NameObject('/Subtype'):NameObject('/Type1'),NameObject('/BaseFont'):NameObject('/Helvetica')})
        page[NameObject('/Resources')]=DictionaryObject({NameObject('/Font'):DictionaryObject({NameObject('/F1'):writer._add_object(font)})})
        content=DecodedStreamObject();content.set_data(b'BT /F1 12 Tf 10 100 Td (Secret Roadmap) Tj ET');page[NameObject('/Contents')]=writer._add_object(content)
        output=io.BytesIO();writer.write(output)
        result=inspect_content(output.getvalue(),policy());self.assertEqual(result['contentType'],'pdf');self.assertEqual(result['ruleId'],'secret')
    def test_semantic_threshold_and_missing_model(self):
        p=policy();p['dlp']['semantic']={'enabled':True,'model':'google/embeddinggemma-2','revision':'test-revision','categories':[{'id':'semantic-secret','threshold':.8,'prototype':[1.,0.]}]}
        os.environ['EMBEDDING_MODEL_REVISION']='test-revision'
        result=inspect_content(b'Confidential paraphrase',p,encoder=lambda *_:[1.,0.]);self.assertEqual(result['ruleId'],'semantic-secret')
        result=inspect_content(b'Confidential paraphrase',p,encoder=lambda *_: (_ for _ in ()).throw(ValueError()));self.assertEqual(result['decision'],'block')
    def test_cosine_validates(self):
        self.assertAlmostEqual(cosine([1,0],[1,0]),1)
        with self.assertRaises(ValueError):cosine([0,0],[0,0])
    def test_video_has_no_false_safe_coverage(self):
        video=b'\x00\x00\x00\x18ftypisom'+b'\x00'*20
        self.assertEqual(content_kind(video),'video');result=inspect_content(video,policy());self.assertEqual(result['decision'],'block');self.assertFalse(result['coverage']['complete'])

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