#!/usr/bin/env python3
"""
Real OpenAI Agent Test with LockStock Integration - v1.1.0
-----------------------------------------------------------
Uses POE's OpenAI-compatible API with real tool calls,
fully instrumented with LockStock cryptographic audit trail.

v1.1.0 CHAIN-BASED AUTHENTICATION:
- NO secrets retrieved or stored
- Guard daemon tracks chain state (hash, matrix, sequence)
- current_hash IS the HMAC key for signatures
- Genesis token is BURNED, not converted to persistent secret
"""

import asyncio
import json
import os
import subprocess
from openai import OpenAI


# Configuration
POE_API_KEY = os.getenv("POE_ACCESS_KEY", "Dpbh5bn1P4oQZ5w7WN5IBAQThHvlOsdT_2XXQTRLXF8")
AGENT_ID = "agent_b5e5d66b_trinityrunpod#1"
GUARD_SOCKET = os.path.expanduser("~/.liberty/liberty.sock")


# Tool definitions
TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "read_file",
            "description": "Read contents of a file",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {"type": "string", "description": "File path"}
                },
                "required": ["path"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "run_command",
            "description": "Execute a shell command (safe commands only)",
            "parameters": {
                "type": "object",
                "properties": {
                    "command": {"type": "string", "description": "Command to run"}
                },
                "required": ["command"]
            }
        }
    }
]


def execute_tool(name: str, args: dict) -> str:
    """Execute tool and return result."""
    print(f"    Executing: {name}({json.dumps(args)})")

    if name == "read_file":
        try:
            path = args["path"]
            with open(path, "r") as f:
                content = f.read()[:300]  # Limit for demo
            return f"File contents ({len(content)} chars):\n{content}"
        except Exception as e:
            return f"Error reading file: {e}"

    elif name == "run_command":
        # Safety: only allow safe commands
        cmd = args["command"]
        safe_commands = ["ls", "pwd", "whoami", "date", "hostname"]
        if not any(cmd.startswith(c) for c in safe_commands):
            return "Command not allowed (demo safety)"
        try:
            result = subprocess.run(
                cmd.split(),
                capture_output=True,
                text=True,
                timeout=5
            )
            return result.stdout or result.stderr or "Command completed"
        except Exception as e:
            return f"Error: {e}"

    return "Unknown tool"


async def main():
    """Run the agent with LockStock chain-based authentication."""

    print("=" * 70)
    print("OpenAI Agent + LockStock v1.1.0 - Chain-Based Auth")
    print("=" * 70)
    print()
    print("NO SECRETS! ONLY CHAIN STATE!")
    print()

    # Initialize LockStock guardrail using Chain Authority Mode
    print("Step 1: Initializing LockStock guardrail (Chain Authority Mode)...")
    from lockstock_openai import LockStockGuardrail

    # Create guardrail - connects to Guard daemon for chain operations
    # NO SECRET PARAMETER! Guard daemon manages chain state internally
    guardrail = LockStockGuardrail.from_liberty(
        agent_id=AGENT_ID,
        socket_path=GUARD_SOCKET,
        endpoint="https://lockstock-api-i9kp.onrender.com"
    )
    print(f"  ✓ Agent ID: {AGENT_ID}")
    print(f"  ✓ Guard Socket: {GUARD_SOCKET}")
    print(f"  ✓ NO SECRETS STORED - Chain state only!")
    print()

    # Initialize chain state (Guard syncs from server if needed)
    print("Step 2: Initializing chain state...")
    try:
        # This ensures Guard has synced chain state from server
        await guardrail._ensure_chain_initialized()
        print(f"  ✓ Chain state initialized")
        print(f"  ✓ Current hash: {guardrail.current_hash[:32] if guardrail.current_hash else 'N/A'}...")
        print(f"  ✓ Sequence: {guardrail.current_sequence}")
    except Exception as e:
        print(f"  ⚠ Chain init: {e}")
    print()

    # Initialize POE client (OpenAI-compatible)
    print("Step 5: Connecting to POE API (OpenAI-compatible)...")
    poe = OpenAI(
        api_key=POE_API_KEY,
        base_url="https://api.poe.com/v1"
    )
    print("  ✓ Connected")
    print()

    # Create conversation
    print("Step 6: Starting agent conversation with tools...")
    print("-" * 70)
    messages = [
        {
            "role": "system",
            "content": "You are a helpful assistant with file and shell access. Keep responses brief."
        },
        {
            "role": "user",
            "content": "Please run 'pwd' command, then read the pyproject.toml file in the current directory."
        }
    ]

    print(f"User: {messages[-1]['content']}")
    print()

    # Agent loop
    max_turns = 5
    for turn in range(1, max_turns + 1):
        print(f"--- Turn {turn} ---")

        try:
            response = poe.chat.completions.create(
                model="Claude-Sonnet-4",
                messages=messages,
                tools=TOOLS,
                tool_choice="auto"
            )
        except Exception as e:
            print(f"  ✗ POE API error: {e}")
            break

        message = response.choices[0].message

        # Handle tool calls
        if message.tool_calls:
            messages.append(message)

            for tool_call in message.tool_calls:
                tool_name = tool_call.function.name
                tool_args = json.loads(tool_call.function.arguments)

                print(f"  Tool call: {tool_name}")

                # Map to LockStock capability
                capability_map = {
                    "read_file": "FileRead",
                    "run_command": "Shell"
                }
                capability = capability_map.get(tool_name, "All")

                # Verify with LockStock (chain-based authentication)
                # Guard uses current_hash as HMAC key, returns signature
                # NO SECRET TRANSMITTED!

                # Create mock tool call for validation
                mock_tool_call = {
                    "name": tool_name,
                    "arguments": tool_args
                }

                # Validate with guardrail (returns GuardrailResult)
                guard_result = await guardrail.validate(agent=None, message={"tool_calls": [mock_tool_call]})

                if guard_result.passed:
                    print(f"  ✓ LockStock: AUTHORIZED (chain-based)")
                    print(f"    Sequence: {guardrail.current_sequence}")
                    print(f"    New hash: {guardrail.current_hash[:32] if guardrail.current_hash else 'N/A'}...")
                    print(f"    (Guard signed with previous hash as key)")

                    # Execute tool
                    tool_result = execute_tool(tool_name, tool_args)
                    print(f"    Result: {tool_result[:100]}...")
                else:
                    print(f"  ✗ LockStock: DENIED - {guard_result.reason}")
                    tool_result = f"DENIED: Agent lacks {capability} capability"

                # Add tool result to messages
                messages.append({
                    "role": "tool",
                    "tool_call_id": tool_call.id,
                    "content": tool_result
                })
            print()

        # Check for finish
        elif message.content:
            print(f"  Agent: {message.content}")
            break
        else:
            break

    print()
    print("=" * 70)
    print("Chain-Based Auth Audit Trail Summary")
    print("=" * 70)
    print(f"Final sequence: {guardrail.current_sequence}")
    print(f"Final hash: {guardrail.current_hash}")
    print()
    print("✓ All tool calls are now in the cryptographic hash chain!")
    print("✓ NO SECRETS were used or transmitted")
    print("✓ Chain state evolved deterministically")
    print("✓ Guard daemon used current_hash as HMAC key")
    print()
    print("✓ View audit trail at: https://lockstock-api-i9kp.onrender.com")
    print(f"✓ Agent ID: {AGENT_ID}")
    print()
    print("WE HAVE NO SECRETS TO KEEP!")

    await guardrail._http.aclose()


if __name__ == "__main__":
    asyncio.run(main())
