"""Tests for lockstock_core.auditor — operator decryption tool."""

import base64
import hashlib
import json
import pytest
import respx
import httpx

from lockstock_core.keyring import (
    generate_account_key,
    generate_dek,
    derive_kek,
    wrap_dek,
    derive_blob_key,
)
from lockstock_core.encrypt import encrypt_payload_envelope
from lockstock_core.auditor import (
    audit_agent,
    fetch_audit_entries,
    fetch_wrapped_dek,
    fetch_payload_blob,
    decrypt_blob,
    format_text,
    main,
    _extract_readable,
)

SERVER = "https://test.lockstock.example.com"
AGENT_ID = "agent_test_auditor"
API_KEY = "lsk_test_key_abc123"


def _make_test_crypto():
    """Create a full set of test crypto materials."""
    account_key = generate_account_key()
    kek = derive_kek(account_key, AGENT_ID)
    dek = generate_dek()
    wrapped_dek = wrap_dek(kek, dek)
    return account_key, kek, dek, wrapped_dek


def _make_encrypted_blob(dek, stamp_hash, payload_hash, plaintext):
    """Encrypt a blob and return the server response dict."""
    ciphertext = encrypt_payload_envelope(plaintext, stamp_hash, payload_hash, dek)
    return {
        "payload_hash": payload_hash,
        "stamp_hash": stamp_hash,
        "agent_id": AGENT_ID,
        "ciphertext_b64": base64.b64encode(ciphertext).decode("ascii"),
        "plaintext_length": len(plaintext),
        "key_version": 1,
        "created_at": "2026-02-15T23:00:00Z",
    }


class TestFetchFunctions:
    """Test individual HTTP fetch functions."""

    @respx.mock
    def test_fetch_audit_entries(self):
        entries = [{"action": "CHATCOMPLETION", "details": {"sequence_id": 1}}]
        respx.get(f"{SERVER}/api/guard/audit").mock(
            return_value=httpx.Response(200, json=entries)
        )
        result = fetch_audit_entries(SERVER, AGENT_ID, API_KEY, limit=10)
        assert len(result) == 1
        assert result[0]["action"] == "CHATCOMPLETION"

    @respx.mock
    def test_fetch_wrapped_dek_latest(self):
        respx.get(f"{SERVER}/v1/agents/{AGENT_ID}/keys/latest").mock(
            return_value=httpx.Response(200, json={"wrapped_dek_b64": "AAAA", "key_version": 1})
        )
        result = fetch_wrapped_dek(SERVER, AGENT_ID, API_KEY)
        assert result["key_version"] == 1

    @respx.mock
    def test_fetch_wrapped_dek_specific_version(self):
        respx.get(f"{SERVER}/v1/agents/{AGENT_ID}/keys/2").mock(
            return_value=httpx.Response(200, json={"wrapped_dek_b64": "BBBB", "key_version": 2})
        )
        result = fetch_wrapped_dek(SERVER, AGENT_ID, API_KEY, key_version=2)
        assert result["key_version"] == 2

    @respx.mock
    def test_fetch_payload_blob_found(self):
        ph = "a" * 64
        respx.get(f"{SERVER}/v1/payload/{ph}").mock(
            return_value=httpx.Response(200, json={"payload_hash": ph, "ciphertext_b64": "xx"})
        )
        result = fetch_payload_blob(SERVER, ph, API_KEY)
        assert result is not None
        assert result["payload_hash"] == ph

    @respx.mock
    def test_fetch_payload_blob_not_found(self):
        ph = "b" * 64
        respx.get(f"{SERVER}/v1/payload/{ph}").mock(
            return_value=httpx.Response(404, json={"error": "not_found"})
        )
        result = fetch_payload_blob(SERVER, ph, API_KEY)
        assert result is None


class TestDecryptBlob:
    """Test blob decryption with real crypto."""

    def test_decrypt_roundtrip(self):
        _, _, dek, _ = _make_test_crypto()
        stamp_hash = "aa" * 32
        plaintext = b"d7:content5:helloe"
        payload_hash = hashlib.sha256(plaintext).hexdigest()

        blob = _make_encrypted_blob(dek, stamp_hash, payload_hash, plaintext)
        result = decrypt_blob(blob, dek)
        assert result == plaintext

    def test_decrypt_wrong_dek_fails(self):
        _, _, dek, _ = _make_test_crypto()
        wrong_dek = generate_dek()
        stamp_hash = "cc" * 32
        plaintext = b"d7:content5:helloe"
        payload_hash = hashlib.sha256(plaintext).hexdigest()

        blob = _make_encrypted_blob(dek, stamp_hash, payload_hash, plaintext)
        with pytest.raises(Exception):
            decrypt_blob(blob, wrong_dek)


class TestAuditAgent:
    """Test the full audit_agent flow with mocked HTTP."""

    @respx.mock
    def test_full_audit_flow(self):
        account_key, kek, dek, wrapped_dek = _make_test_crypto()
        stamp_hash_1 = "11" * 32
        stamp_hash_2 = "22" * 32
        plaintext_1 = b"d7:content13:Hello Trinity!e"
        plaintext_2 = b"d7:content18:I see the markets.e"
        payload_hash_1 = hashlib.sha256(plaintext_1).hexdigest()
        payload_hash_2 = hashlib.sha256(plaintext_2).hexdigest()

        # Mock audit entries (server returns newest first)
        entries = [
            {
                "action": "CHATCOMPLETIONRESPONSE",
                "timestamp": "2026-02-15T23:01:00Z",
                "details": {"sequence_id": 2, "payload_hash": payload_hash_2},
            },
            {
                "action": "CHATCOMPLETION",
                "timestamp": "2026-02-15T23:00:00Z",
                "details": {"sequence_id": 1, "payload_hash": payload_hash_1},
            },
        ]
        respx.get(f"{SERVER}/api/guard/audit").mock(
            return_value=httpx.Response(200, json=entries)
        )

        # Mock wrapped DEK
        respx.get(f"{SERVER}/v1/agents/{AGENT_ID}/keys/latest").mock(
            return_value=httpx.Response(200, json={
                "wrapped_dek_b64": base64.b64encode(wrapped_dek).decode("ascii"),
                "key_version": 1,
            })
        )

        # Mock blobs
        blob_1 = _make_encrypted_blob(dek, stamp_hash_1, payload_hash_1, plaintext_1)
        blob_2 = _make_encrypted_blob(dek, stamp_hash_2, payload_hash_2, plaintext_2)
        respx.get(f"{SERVER}/v1/payload/{payload_hash_1}").mock(
            return_value=httpx.Response(200, json=blob_1)
        )
        respx.get(f"{SERVER}/v1/payload/{payload_hash_2}").mock(
            return_value=httpx.Response(200, json=blob_2)
        )

        # Run audit
        transcript = audit_agent(
            account_key=account_key,
            agent_id=AGENT_ID,
            api_key=API_KEY,
            server_url=SERVER,
        )

        # Oldest first
        assert len(transcript) == 2
        assert transcript[0]["sequence"] == 1
        assert transcript[0]["task"] == "CHATCOMPLETION"
        assert transcript[0]["decrypted"] is True
        assert b"Hello Trinity!" in transcript[0]["content"].encode()

        assert transcript[1]["sequence"] == 2
        assert transcript[1]["task"] == "CHATCOMPLETIONRESPONSE"
        assert transcript[1]["decrypted"] is True
        assert b"I see the markets." in transcript[1]["content"].encode()

    @respx.mock
    def test_audit_with_missing_blob(self):
        account_key, kek, dek, wrapped_dek = _make_test_crypto()
        payload_hash = "55" * 32

        entries = [
            {
                "action": "CHATCOMPLETION",
                "timestamp": "2026-02-15T23:00:00Z",
                "details": {"sequence_id": 1, "payload_hash": payload_hash},
            },
        ]
        respx.get(f"{SERVER}/api/guard/audit").mock(
            return_value=httpx.Response(200, json=entries)
        )
        respx.get(f"{SERVER}/v1/agents/{AGENT_ID}/keys/latest").mock(
            return_value=httpx.Response(200, json={
                "wrapped_dek_b64": base64.b64encode(wrapped_dek).decode("ascii"),
                "key_version": 1,
            })
        )
        respx.get(f"{SERVER}/v1/payload/{payload_hash}").mock(
            return_value=httpx.Response(404, json={"error": "not_found"})
        )

        transcript = audit_agent(
            account_key=account_key,
            agent_id=AGENT_ID,
            api_key=API_KEY,
            server_url=SERVER,
        )

        assert len(transcript) == 1
        assert transcript[0]["decrypted"] is False
        assert "not found" in transcript[0]["error"]

    @respx.mock
    def test_audit_sequence_filter(self):
        account_key, kek, dek, wrapped_dek = _make_test_crypto()
        stamp_hash = "66" * 32
        plaintext = b"d7:content4:teste"
        payload_hash = hashlib.sha256(plaintext).hexdigest()

        entries = [
            {
                "action": "CHATCOMPLETIONRESPONSE",
                "timestamp": "2026-02-15T23:02:00Z",
                "details": {"sequence_id": 3, "payload_hash": "88" * 32},
            },
            {
                "action": "CHATCOMPLETION",
                "timestamp": "2026-02-15T23:01:00Z",
                "details": {"sequence_id": 2, "payload_hash": payload_hash},
            },
            {
                "action": "CHATCOMPLETION",
                "timestamp": "2026-02-15T23:00:00Z",
                "details": {"sequence_id": 1, "payload_hash": "99" * 32},
            },
        ]
        respx.get(f"{SERVER}/api/guard/audit").mock(
            return_value=httpx.Response(200, json=entries)
        )
        respx.get(f"{SERVER}/v1/agents/{AGENT_ID}/keys/latest").mock(
            return_value=httpx.Response(200, json={
                "wrapped_dek_b64": base64.b64encode(wrapped_dek).decode("ascii"),
                "key_version": 1,
            })
        )

        blob = _make_encrypted_blob(dek, stamp_hash, payload_hash, plaintext)
        respx.get(f"{SERVER}/v1/payload/{payload_hash}").mock(
            return_value=httpx.Response(200, json=blob)
        )

        transcript = audit_agent(
            account_key=account_key,
            agent_id=AGENT_ID,
            api_key=API_KEY,
            server_url=SERVER,
            from_seq=2,
            to_seq=2,
        )

        assert len(transcript) == 1
        assert transcript[0]["sequence"] == 2


class TestExtractReadable:
    """Test bencoded content extraction."""

    def test_extracts_content_field(self):
        bencoded = "d7:content25:Trinity, what do you see?4:role4:usere"
        result = _extract_readable(bencoded)
        assert result == "Trinity, what do you see?"

    def test_extracts_longest_content(self):
        bencoded = "d7:content4:test7:content20:This is the real onee"
        result = _extract_readable(bencoded)
        assert result == "This is the real one"

    def test_returns_none_for_no_content(self):
        bencoded = "d4:role4:user8:messages5:helloe"
        result = _extract_readable(bencoded)
        assert result is None


class TestFormatText:
    """Test text output formatting."""

    def test_formats_decrypted_entry(self):
        transcript = [{
            "sequence": 17,
            "task": "CHATCOMPLETION",
            "timestamp": "2026-02-15T23:00:00Z",
            "payload_hash": "aa" * 32,
            "stamp_hash": "bb" * 32,
            "key_version": 1,
            "plaintext_length": 100,
            "content": "d7:content25:Trinity, what do you see?e",
            "decrypted": True,
        }]
        output = format_text(transcript, AGENT_ID)
        assert "seq 17" in output
        assert "CHATCOMPLETION" in output
        assert "Trinity, what do you see?" in output

    def test_formats_failed_entry(self):
        transcript = [{
            "sequence": 5,
            "task": "CHATCOMPLETION",
            "timestamp": "2026-02-15T23:00:00Z",
            "payload_hash": "cc" * 32,
            "decrypted": False,
            "error": "blob not found on server",
        }]
        output = format_text(transcript, AGENT_ID)
        assert "NOT DECRYPTED" in output
        assert "not found" in output


class TestCLI:
    """Test CLI argument parsing and execution."""

    def test_invalid_account_key_length(self, capsys):
        code = main(["--account-key", "abcd", "--agent-id", "test", "--api-key", "key"])
        assert code == 1
        captured = capsys.readouterr()
        assert "32 bytes" in captured.err

    def test_invalid_account_key_hex(self, capsys):
        code = main(["--account-key", "zz" * 32, "--agent-id", "test", "--api-key", "key"])
        assert code == 1
        captured = capsys.readouterr()
        assert "hex" in captured.err.lower()
