sekuire 0.1.0

use ed25519_dalek::{Signature, Signer, SigningKey, Verifier, VerifyingKey};
use rand::rngs::OsRng;
use serde::{Deserialize, Serialize};
use thiserror::Error;

#[derive(Debug, Serialize, Deserialize)]
pub struct HandshakeHello {
    pub client_nonce: String,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct HandshakeWelcome {
    pub agent_id: String,
    pub agent_nonce: String,
    pub signature_c: String,
    pub credentials: Vec<String>,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct HandshakeAuth {
    pub signature_a: String,
}

// Core modules
pub mod builtin_tools;
pub mod config;
pub mod identity;
pub mod llm;
pub mod logger;
pub mod memory;
pub mod new_agent;
pub mod policy;
pub mod worker;

// Legacy modules - temporarily disabled due to API changes
// TODO: Update these modules to use new LLM provider interface
// pub mod compliance;
// pub mod legacy_agent;
// pub mod tools;

// Config-first API (recommended)
pub use config::{AgentConfig as Config, LLMConfig, SekuireConfig};
pub use llm::{ChatOptions, ChatResponse, LLMProvider, Message};
pub use new_agent::{get_agent, get_agents, get_system_prompt, get_tools, SekuireAgent};

// Built-in tools
pub use builtin_tools::{
    create_default_tool_registry, CalculatorTool, CreateDirectoryTool, FileReadTool,
    HttpRequestTool, ListDirectoryTool, Tool, ToolInput, ToolMetadata, ToolParameter, ToolRegistry,
    WebSearchTool, WriteFileTool,
};

// Memory storage
pub use memory::{create_memory_storage, InMemoryStorage, MemoryStorage};

// LLM providers
pub use llm::{
    create_llm_provider, AnthropicProvider, GeminiProvider, OllamaProvider, OpenAIProvider,
};

// Identity
pub use identity::{calculate_sekuire_id, AgentIdentity};
pub use policy::{ActivePolicy, PolicyClient, PolicyEnforcer};

// Task Worker (Async Mesh)
pub use worker::{create_worker, TaskContext, TaskEvent, TaskHandler, TaskWorker, WorkerConfig};

#[derive(Error, Debug)]
pub enum HandshakeError {
    #[error("Network error: {0}")]
    Network(#[from] reqwest::Error),
    #[error("Protocol error: {0}")]
    Protocol(String),
    #[error("Crypto error: {0}")]
    Crypto(String),
}

/// A Client that can connect to a Sekuire Agent and perform the handshake.
pub struct SekuireClient {
    signing_key: SigningKey,
    registry_url: String, // To verify VCs
}

impl SekuireClient {
    pub fn new(signing_key: SigningKey, registry_url: String) -> Self {
        Self {
            signing_key,
            registry_url,
        }
    }

    /// Perform the full handshake with a remote agent.
    pub async fn connect(
        &self,
        agent_url: &str,
        expected_agent_id: Option<&str>,
    ) -> Result<(), HandshakeError> {
        let client = reqwest::Client::new();

        // 1. Generate Nonce C
        let mut csprng = OsRng;
        let nonce_c_bytes: [u8; 32] = rand::Rng::gen(&mut csprng);
        let nonce_c = hex::encode(nonce_c_bytes);

        // 2. Send HELLO
        let hello = HandshakeHello {
            client_nonce: nonce_c.clone(),
        };
        let res = client
            .post(format!("{}/sekuire/hello", agent_url))
            .json(&hello)
            .send()
            .await?;

        if !res.status().is_success() {
            return Err(HandshakeError::Protocol(format!(
                "Agent rejected HELLO: {}",
                res.status()
            )));
        }

        let welcome: HandshakeWelcome = res.json().await?;

        // 3. Verify Agent Identity
        if let Some(expected) = expected_agent_id {
            if welcome.agent_id != expected {
                return Err(HandshakeError::Protocol(format!(
                    "Agent ID Mismatch. Expected {}, got {}",
                    expected, welcome.agent_id
                )));
            }
        }

        // 4. Verify Signature C (Agent signed our nonce)
        // We need the Agent's Public Key. In a real flow, we fetch this from Registry.
        // For now, we assume we fetch it via GET /agents/:id from the Registry.
        let agent_pub_key = self.fetch_public_key(&welcome.agent_id).await?;

        let sig_c_bytes = hex::decode(&welcome.signature_c)
            .map_err(|_| HandshakeError::Crypto("Invalid Signature C hex".into()))?;
        let sig_c = Signature::from_bytes(&sig_c_bytes.try_into().unwrap());

        if let Err(e) = agent_pub_key.verify(&nonce_c_bytes, &sig_c) {
            return Err(HandshakeError::Crypto(format!(
                "Agent failed to sign Client Nonce: {}",
                e
            )));
        }

        // 5. Sign Nonce A (We sign their nonce)
        let sig_a = self.signing_key.sign(welcome.agent_nonce.as_bytes());
        let auth = HandshakeAuth {
            signature_a: hex::encode(sig_a.to_bytes()),
        };

        // 6. Send AUTH
        let res = client
            .post(format!("{}/sekuire/auth", agent_url))
            .json(&auth)
            .send()
            .await?;

        if !res.status().is_success() {
            return Err(HandshakeError::Protocol("Agent rejected AUTH".into()));
        }

        Ok(())
    }

    async fn fetch_public_key(&self, agent_id: &str) -> Result<VerifyingKey, HandshakeError> {
        let url = format!("{}/api/v1/agents/{}", self.registry_url, agent_id);
        let res = reqwest::get(&url).await?;

        if !res.status().is_success() {
            return Err(HandshakeError::Protocol(format!(
                "Registry lookup failed for {}",
                agent_id
            )));
        }

        let agent_data: serde_json::Value = res.json().await?;
        let key_hex = agent_data["author_key"]
            .as_str()
            .ok_or(HandshakeError::Protocol(
                "Registry response missing author_key".into(),
            ))?;

        let key_bytes = hex::decode(key_hex)
            .map_err(|_| HandshakeError::Crypto("Invalid Public Key Hex from Registry".into()))?;

        VerifyingKey::from_bytes(&key_bytes.try_into().unwrap())
            .map_err(|_| HandshakeError::Crypto("Invalid Ed25519 Public Key".into()))
    }
}