import json
import os
from pathlib import Path

import requests
from a2spa import (
    build_payload_fields as package_build_payload_fields,
    build_request_body,
    build_signed_request as build_a2spa_signed_request,
    compute_payload_hash,
    finalize_payload,
    policy_input,
    require_authorized_result,
)


# Install with: pip install a2spa requests


class A2SPAClient:
    def __init__(self, api_base, api_key, agent_id, private_key_path, timeout=20):
        self.api_base = api_base.rstrip("/")
        self.api_key = api_key
        self.agent_id = agent_id
        self.private_key_path = private_key_path
        self.timeout = timeout
        self._private_key = None

    @classmethod
    def from_env(cls):
        return cls(
            api_base=os.environ.get("A2SPA_API_BASE", "https://api.aimodularity.com/A2SPA"),
            api_key=os.environ["A2SPA_API_KEY"],
            agent_id=os.environ["A2SPA_AGENT_ID"],
            private_key_path=os.environ["A2SPA_PRIVATE_KEY_PATH"],
        )

    def _load_private_key(self):
        if self._private_key is None:
            self._private_key = Path(self.private_key_path).read_text(encoding="utf-8")
        return self._private_key

    def build_payload_fields(self, target_agent_id, input_data, output_data=None, alert_threshold=10, state_continuity=None):
        return package_build_payload_fields(
            agent_id=self.agent_id,
            target_agent_id=target_agent_id,
            input_data=input_data,
            output_data=output_data if output_data is not None else {"status": "ready"},
            alert_threshold=alert_threshold,
            state_continuity=state_continuity,
        )

    def build_signed_request(self, target_agent_id, input_data, output_data=None, alert_threshold=10, state_continuity=None):
        payload = self.build_payload_fields(target_agent_id, input_data, output_data, alert_threshold, state_continuity)
        return build_a2spa_signed_request(self._load_private_key(), payload)

    def post_request_body(self, request_body):
        return requests.post(
            f"{self.api_base}/api/verify_payload",
            headers={
                "Content-Type": "application/json",
                "x-api-key": self.api_key,
            },
            json=request_body,
            timeout=self.timeout,
        )

    def send_payload(self, target_agent_id, input_data, output_data=None, alert_threshold=10, state_continuity=None, debug=False):
        request_body = self.build_signed_request(target_agent_id, input_data, output_data, alert_threshold, state_continuity)
        response = self.post_request_body(request_body)
        if debug and response.status_code >= 400:
            print(json.dumps({"status_code": response.status_code, "error": response.text}, indent=2))
        response.raise_for_status()
        result = response.json()
        return require_authorized_result(result)

    def fetch_logs(self, agent_id=None):
        response = requests.get(
            f"{self.api_base}/api/logs_for_agent",
            headers={"x-api-key": self.api_key},
            params={"agent_id": agent_id or self.agent_id},
            timeout=self.timeout,
        )
        response.raise_for_status()
        return response.json().get("logs", [])

    def fetch_inbox(self, agent_id=None, after=None, limit=50):
        params = {"agent_id": agent_id or self.agent_id, "limit": limit}
        if after:
            params["after"] = after
        response = requests.get(
            f"{self.api_base}/api/inbox_for_agent",
            headers={"x-api-key": self.api_key},
            params=params,
            timeout=self.timeout,
        )
        response.raise_for_status()
        return response.json().get("messages", [])

    def prepare_and_send(self, target_agent_id, input_data, runner, alert_threshold=10):
        # The runner must only prepare output_data. Do not perform side effects before A2SPA authorization.
        output_data = runner(input_data)
        return self.send_payload(target_agent_id, input_data, output_data, alert_threshold)


def existing_agent(input_data):
    return {
        "status": "ready",
        "summary": f"Processed: {input_data.get('message', 'no message')}",
    }


if __name__ == "__main__":
    client = A2SPAClient.from_env()
    target_agent_id = os.environ["A2SPA_TARGET_AGENT_ID"]
    input_data = policy_input(action="send_message", workflow_scope="messages:send", message="Hello from my existing Python agent")
    output_data = existing_agent(input_data)
    signed_request = client.build_signed_request(target_agent_id, input_data, output_data)
    response = client.post_request_body(signed_request)
    print(f"A2SPA verify response: {response.status_code} {response.text}")