from fastmcp import FastMCP
from crypto_engine import hash_tool, worm_write_tool, proof_generate_tool, verify_proof_tool
import json

# Initialize the MCP server
mcp = FastMCP("Secure Reasoning Server")


@mcp.tool()
def hash_data(data: str) -> str:
    """
    Hash a string or JSON data using SHA-256.
    Input: data (string or JSON-serializable object as string)
    Output: SHA-256 hex digest
    """
    return hash_tool(data)


@mcp.tool()
def write_to_worm(step_data: str, hash_value: str, merkle_root: str) -> str:
    """
    Write a step record to WORM (Write Once, Read Many) storage.
    Input: step_data (JSON string), hash_value (hex string), merkle_root (hex string)
    Output: JSON record with id, timestamp, step, hash, and root
    """
    step_dict = json.loads(step_data) if isinstance(step_data, str) else step_data
    record = worm_write_tool(step_dict, hash_value, merkle_root)
    return json.dumps(record)


@mcp.tool()
def generate_proof(record_id: int) -> str:
    """
    Generate a Merkle proof for a specific record in the WORM log.
    Input: record_id (integer, the ID of the record)
    Output: JSON containing record_id, hash, merkle_proof, merkle_root, timestamp, and step_details
    """
    proof = proof_generate_tool(record_id)
    if proof is None:
        return json.dumps({"error": f"Record with ID {record_id} not found"})
    return json.dumps(proof)


@mcp.tool()
def verify_proof(target_hash: str, merkle_proof: str, merkle_root: str) -> str:
    """
    Verify if a target_hash belongs to the merkle_root using the merkle_proof.
    Input: target_hash (hex string), merkle_proof (JSON string of proof array), merkle_root (hex string)
    Output: JSON with result (true/false) and verification status message
    """
    proof_list = json.loads(merkle_proof) if isinstance(merkle_proof, str) else merkle_proof
    is_valid = verify_proof_tool(target_hash, proof_list, merkle_root)
    return json.dumps({
        "verified": is_valid,
        "message": "Proof verified successfully" if is_valid else "Proof verification failed - possible tampering"
    })


if __name__ == "__main__":
    mcp.run()