"""Tests for the pipeline certification package."""

import json
import tempfile
from pathlib import Path

import pytest
from rdflib import Dataset

from cf_pipeline_cert import (
    generate_keypair,
    sign_data,
    verify_signature,
    compute_distribution_hash,
    validate_step_package,
    verify_certificate,
)
from cf_pipeline_cert.cli import cmd_certify, cmd_import
from cf_pipeline_cert.crypto import create_signable_json

_STRUCTURED_SUFFIX = "." + "json" + "ld"
_STRUCTURED_FORMAT = "json" + "-ld"


def _build_steps_document(
    *,
    plugin_name: str = "test-plugin",
    plugin_version: str = "1.0.0",
    operation_name: str = "test_op",
    include_operation: bool = True,
    include_contract: bool = True,
) -> dict:
    return {
        "@context": {
            "cf": "https://cogniflow.odea-project.org/cf#",
            "rdfs": "http://www.w3.org/2000/01/rdf-schema#",
            "rdf": "http://www.w3.org/1999/02/22-rdf-syntax-ns#",
        },
        "@graph": [
            {
                "@id": "pkg:TestPackage",
                "@type": "cf:StepPackage",
                "cf:packageId": plugin_name,
                "cf:packageVersion": plugin_version,
                "rdfs:label": "Test Package",
                "cf:description": "Test package",
            },
            {
                "@id": "pkg:TestPlugin",
                "@type": "cf:PluginArtifact",
                "cf:pluginName": plugin_name,
                "cf:pluginVersion": plugin_version,
                "cf:artifactPath": "bin",
                "cf:pluginListSymbol": "cf_list_steps",
                "cf:pluginResolveSymbol": "cf_resolve_step",
            },
            {
                "@id": "pkg:TestImpl",
                "@type": "cf:StepImplementation",
                "cf:implementationArtifact": {"@id": "pkg:TestPlugin"},
                **({"cf:operationName": operation_name} if include_operation else {}),
                **(
                    {"cf:operationContract": {"kind": "test"}}
                    if include_contract and include_operation
                    else {}
                ),
            },
            {
                "@id": "pkg:TestStep",
                "@type": "cf:ProcessingStep",
                "cf:belongsToStepPackage": {"@id": "pkg:TestPackage"},
                "cf:hasImplementation": {"@id": "pkg:TestImpl"},
            },
        ],
    }


def _prepare_document_for_nquads(document: dict) -> dict:
    payload = json.loads(json.dumps(document))
    for node in payload.get("@graph", []):
        if not isinstance(node, dict):
            continue
        for key in ("cf:operationContract", "cf:pluginCertificate"):
            value = node.get(key)
            if isinstance(value, dict) and "@value" not in value:
                node[key] = {
                    "@type": "rdf:JSON",
                    "@value": json.dumps(value, sort_keys=True),
                }
    return payload


def _write_step_document(
    plugin_path: Path,
    *,
    file_name: str = "steps.json",
    plugin_name: str = "test-plugin",
    plugin_version: str = "1.0.0",
    operation_name: str = "test_op",
    include_operation: bool = True,
    include_contract: bool = True,
) -> Path:
    steps_doc = _build_steps_document(
        plugin_name=plugin_name,
        plugin_version=plugin_version,
        operation_name=operation_name,
        include_operation=include_operation,
        include_contract=include_contract,
    )
    path = plugin_path / file_name
    if path.suffix.lower() in {_STRUCTURED_SUFFIX, ".json"}:
        path.write_text(json.dumps(steps_doc, indent=2), encoding="utf-8")
        return path

    dataset = Dataset()
    dataset.parse(data=json.dumps(_prepare_document_for_nquads(steps_doc)), format=_STRUCTURED_FORMAT)
    payload = dataset.serialize(format="nquads")
    if isinstance(payload, bytes):
        payload = payload.decode("utf-8")
    path.write_text(payload, encoding="utf-8")
    return path


def _load_document(path: Path) -> dict:
    if path.suffix.lower() in {_STRUCTURED_SUFFIX, ".json"}:
        return json.loads(path.read_text(encoding="utf-8"))
    dataset = Dataset()
    dataset.parse(data=path.read_text(encoding="utf-8"), format="nquads")
    payload = dataset.serialize(
        format=_STRUCTURED_FORMAT,
        context={
            "cf": "https://cogniflow.odea-project.org/cf#",
            "rdf": "http://www.w3.org/1999/02/22-rdf-syntax-ns#",
            "rdfs": "http://www.w3.org/2000/01/rdf-schema#",
            "schema": "https://schema.org/",
        },
        auto_compact=True,
    )
    if isinstance(payload, bytes):
        payload = payload.decode("utf-8")
    return json.loads(payload)


def _write_document(path: Path, document: dict) -> None:
    if path.suffix.lower() in {_STRUCTURED_SUFFIX, ".json"}:
        path.write_text(json.dumps(document, indent=2), encoding="utf-8")
        return
    dataset = Dataset()
    dataset.parse(data=json.dumps(_prepare_document_for_nquads(document)), format=_STRUCTURED_FORMAT)
    payload = dataset.serialize(format="nquads")
    if isinstance(payload, bytes):
        payload = payload.decode("utf-8")
    path.write_text(payload, encoding="utf-8")


class _Args:
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)

class TestCrypto:
    """Test cryptographic operations."""
    
    def test_generate_keypair(self):
        """Test keypair generation."""
        private_key, public_key = generate_keypair()
        
        assert len(private_key) == 44  # Base64 encoded 32 bytes
        assert len(public_key) == 44
        assert private_key != public_key
    
    def test_sign_and_verify(self):
        """Test signing and verification."""
        private_key, public_key = generate_keypair()
        data = b"Hello, CogniFlow!"
        
        signature = sign_data(data, private_key)
        assert verify_signature(data, signature, public_key)
    
    def test_tampered_data_fails_verification(self):
        """Test that tampered data fails verification."""
        private_key, public_key = generate_keypair()
        data = b"Original message"
        
        signature = sign_data(data, private_key)
        tampered_data = b"Tampered message"
        
        assert not verify_signature(tampered_data, signature, public_key)
    
    def test_wrong_key_fails_verification(self):
        """Test that wrong key fails verification."""
        private_key1, public_key1 = generate_keypair()
        _, public_key2 = generate_keypair()
        
        data = b"Test message"
        signature = sign_data(data, private_key1)
        
        assert verify_signature(data, signature, public_key1)
        assert not verify_signature(data, signature, public_key2)
    
    def test_create_signable_json_is_deterministic(self):
        """Test that signable JSON is deterministic."""
        data1 = {"b": 2, "a": 1, "signature": "should_be_removed"}
        data2 = {"a": 1, "b": 2}
        
        assert create_signable_json(data1) == create_signable_json(data2)


class TestValidation:
    """Test step package validation."""
    
    @pytest.fixture
    def temp_plugin(self):
        """Create a temporary step package."""
        with tempfile.TemporaryDirectory() as tmpdir:
            plugin_path = Path(tmpdir)

            _write_step_document(plugin_path, plugin_name="test-plugin", plugin_version="1.0.0")
            
            # Create a Python file for hash computation
            with open(plugin_path / "ops.py", "w") as f:
                f.write("def test_op(x): return x * 2\n")
            
            yield plugin_path
    
    def test_validate_valid_plugin(self, temp_plugin):
        """Test validation of a valid plugin."""
        result = validate_step_package(temp_plugin)
        
        assert result.valid
        assert result.plugin_name == "test-plugin"
        assert result.plugin_version == "1.0.0"
        assert len(result.operations) == 1
        assert result.distribution_hash

    def test_validate_valid_plugin_nquads_manifest(self):
        """Test validation of a valid plugin with steps.nq."""
        with tempfile.TemporaryDirectory() as tmpdir:
            plugin_path = Path(tmpdir)
            _write_step_document(plugin_path, file_name="steps.nq")
            (plugin_path / "ops.py").write_text("def test_op(x): return x * 2\n", encoding="utf-8")

            result = validate_step_package(plugin_path)

            assert result.valid
            assert result.plugin_name == "test-plugin"
            assert result.plugin_version == "1.0.0"
    
    def test_validate_missing_step_document(self):
        """Test validation fails without a step document."""
        with tempfile.TemporaryDirectory() as tmpdir:
            result = validate_step_package(Path(tmpdir))
            
            assert not result.valid
            assert any("step document" in e for e in result.errors)

    def test_validate_prefers_steps_nq_over_structured_manifest(self):
        """Test package-directory lookup prefers steps.nq when both manifests exist."""
        with tempfile.TemporaryDirectory() as tmpdir:
            plugin_path = Path(tmpdir)
            _write_step_document(plugin_path, file_name="steps.nq")
            (plugin_path / "steps.json").write_text("{\n", encoding="utf-8")
            (plugin_path / "ops.py").write_text("def test_op(x): return x * 2\n", encoding="utf-8")

            result = validate_step_package(plugin_path)

            assert result.valid
    
    def test_validate_missing_operations(self):
        """Test validation fails without operations."""
        with tempfile.TemporaryDirectory() as tmpdir:
            plugin_path = Path(tmpdir)

            _write_step_document(
                plugin_path,
                plugin_name="test",
                plugin_version="1.0.0",
                include_operation=False,
            )
            
            result = validate_step_package(plugin_path)
            assert not result.valid


class TestCertification:
    """Test full certification workflow."""
    
    @pytest.fixture
    def temp_plugin_with_cert(self):
        """Create a temp plugin with certificate."""
        private_key, public_key = generate_keypair()
        
        with tempfile.TemporaryDirectory() as tmpdir:
            plugin_path = Path(tmpdir)

            manifest_path = _write_step_document(
                plugin_path,
                plugin_name="certified-plugin",
                plugin_version="1.0.0",
                operation_name="certified_op",
            )
            
            with open(plugin_path / "ops.py", "w") as f:
                f.write("def certified_op(x): return x\n")
            
            # Validate and create certificate
            result = validate_step_package(plugin_path)
            cert = result.to_certificate(sign_key_b64=private_key)

            steps_doc = _load_document(manifest_path)
            for node in steps_doc.get("@graph", []):
                if isinstance(node, dict) and node.get("@id") == "pkg:TestPlugin":
                    node["cf:pluginCertificate"] = cert
            _write_document(manifest_path, steps_doc)
            
            yield plugin_path, public_key, private_key

    @pytest.fixture
    def temp_plugin_with_cert_nquads(self):
        """Create a temp plugin with certificate stored in steps.nq."""
        private_key, public_key = generate_keypair()

        with tempfile.TemporaryDirectory() as tmpdir:
            plugin_path = Path(tmpdir)

            manifest_path = _write_step_document(
                plugin_path,
                file_name="steps.nq",
                plugin_name="certified-plugin-nq",
                plugin_version="1.0.0",
                operation_name="certified_op",
            )
            (plugin_path / "ops.py").write_text("def certified_op(x): return x\n", encoding="utf-8")

            result = validate_step_package(plugin_path)
            cert = result.to_certificate(sign_key_b64=private_key)

            steps_doc = _load_document(manifest_path)
            graph = steps_doc.get("@graph") if isinstance(steps_doc.get("@graph"), list) else [steps_doc]
            for node in graph:
                if isinstance(node, dict) and node.get("@id") == "pkg:TestPlugin":
                    node["cf:pluginCertificate"] = cert
            _write_document(manifest_path, steps_doc)

            yield plugin_path, public_key
    
    def test_verify_valid_certificate(self, temp_plugin_with_cert):
        """Test verification of a valid certificate."""
        plugin_path, public_key, _ = temp_plugin_with_cert
        
        result = verify_certificate(plugin_path, public_key)
        
        assert result["valid"]
        assert result["hash_matches"]
        assert result["signature_valid"]

    def test_verify_valid_certificate_nquads_manifest(self, temp_plugin_with_cert_nquads):
        """Test verification of a valid certificate stored in steps.nq."""
        plugin_path, public_key = temp_plugin_with_cert_nquads

        result = verify_certificate(plugin_path, public_key)

        assert result["valid"]
        assert result["hash_matches"]
        assert result["signature_valid"]
    
    def test_detect_code_tampering(self, temp_plugin_with_cert):
        """Test that code tampering is detected."""
        plugin_path, public_key, _ = temp_plugin_with_cert
        
        # Tamper with the code
        with open(plugin_path / "ops.py", "w") as f:
            f.write("def certified_op(x): return x * 1000  # BACKDOOR!\n")
        
        result = verify_certificate(plugin_path, public_key)
        
        assert not result["valid"]
        assert not result["hash_matches"]
        assert result["signature_valid"]  # Signature still valid, but hash doesn't match
    
    def test_detect_certificate_tampering(self, temp_plugin_with_cert):
        """Test that certificate tampering is detected."""
        plugin_path, public_key, _ = temp_plugin_with_cert
        
        # Tamper with the certificate
        manifest_path = plugin_path / "steps.json"
        data = _load_document(manifest_path)

        for node in data.get("@graph", []):
            if isinstance(node, dict) and node.get("@id") == "pkg:TestPlugin":
                node["cf:pluginCertificate"]["plugin_version"] = "9.9.9"

        _write_document(manifest_path, data)
        
        result = verify_certificate(plugin_path, public_key)
        
        assert not result["valid"]
        assert not result["signature_valid"]


class TestSelfSigned:
    """Test self-signed certificate workflow."""
    
    @pytest.fixture
    def temp_self_signed_plugin(self):
        """Create a temp plugin with self-signed certificate."""
        with tempfile.TemporaryDirectory() as tmpdir:
            plugin_path = Path(tmpdir)

            manifest_path = _write_step_document(
                plugin_path,
                plugin_name="self-signed-plugin",
                plugin_version="1.0.0",
                operation_name="my_op",
            )
            
            with open(plugin_path / "ops.py", "w") as f:
                f.write("def my_op(x): return x * 2\n")
            
            # Validate and create self-signed certificate
            result = validate_step_package(plugin_path)
            cert = result.to_self_signed_certificate()

            steps_doc = _load_document(manifest_path)
            for node in steps_doc.get("@graph", []):
                if isinstance(node, dict) and node.get("@id") == "pkg:TestPlugin":
                    node["cf:pluginCertificate"] = cert
            _write_document(manifest_path, steps_doc)
            
            yield plugin_path
    
    def test_self_signed_verification_without_key(self, temp_self_signed_plugin):
        """Test that self-signed certs can be verified without external key."""
        # No public key provided - should use embedded key
        result = verify_certificate(temp_self_signed_plugin, public_key_b64=None)
        
        assert result["valid"]
        assert result["hash_matches"]
        assert result["signature_valid"]
        assert result["trust_level"] == "self-signed"
    
    def test_self_signed_detects_tampering(self, temp_self_signed_plugin):
        """Test that tampering is detected in self-signed certs."""
        # Tamper with code
        with open(temp_self_signed_plugin / "ops.py", "w") as f:
            f.write("def my_op(x): return x * 999  # TAMPERED\n")
        
        result = verify_certificate(temp_self_signed_plugin, public_key_b64=None)
        
        assert not result["valid"]
        assert not result["hash_matches"]


class TestCliRdfDocuments:
    """Test CLI flows against RDF step documents."""

    def test_cmd_import_preserves_nquads_manifest(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            plugin_path = Path(tmpdir)
            manifest_path = _write_step_document(plugin_path, file_name="steps.nq")
            (plugin_path / "ops.py").write_text("def test_op(x): return x * 2\n", encoding="utf-8")

            result = validate_step_package(plugin_path)
            cert_path = plugin_path / "signed.cert.json"
            cert_path.write_text(
                json.dumps(result.to_self_signed_certificate(), indent=2),
                encoding="utf-8",
            )

            exit_code = cmd_import(
                _Args(certificate=str(cert_path), path=str(plugin_path))
            )

            assert exit_code == 0
            assert manifest_path.suffix == ".nq"
            assert "rdf-syntax-ns#JSON" in manifest_path.read_text(encoding="utf-8")

            verification = verify_certificate(plugin_path, public_key_b64=None)
            assert verification["valid"]

    def test_cmd_certify_preserves_nquads_manifest(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            plugin_path = Path(tmpdir)
            manifest_path = _write_step_document(plugin_path, file_name="steps.nq")
            (plugin_path / "ops.py").write_text("def test_op(x): return x * 2\n", encoding="utf-8")

            exit_code = cmd_certify(
                _Args(path=str(plugin_path), self_signed=True, key=None)
            )

            assert exit_code == 0
            assert manifest_path.suffix == ".nq"

            document = _load_document(manifest_path)
            plugin_node = next(
                node
                for node in document.get("@graph", [])
                if isinstance(node, dict) and node.get("@id") == "pkg:TestPlugin"
            )
            assert isinstance(plugin_node.get("cf:pluginCertificate"), dict)

            verification = verify_certificate(plugin_path, public_key_b64=None)
            assert verification["valid"]


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
