#!/usr/bin/env python3
"""
POE + LockStock Authenticated Agent
------------------------------------
Demonstrates cryptographic hash chain authentication on every tool call
using POE's OpenAI-compatible API.

Every action the agent takes is:
1. Verified against LockStock capabilities
2. Signed with HMAC-SHA256
3. Added to an immutable hash chain

Prerequisites:
    pip install openai httpx

Usage:
    python poe_lockstock_agent.py
"""

import asyncio
import json
import os
import subprocess
import sys
from typing import Optional

# Add parent to path for imports
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from openai import OpenAI
from lockstock_core import LockStockClient


def get_liberty_secret(name: str) -> Optional[str]:
    """Get a secret from Liberty vault."""
    try:
        result = subprocess.run(
            ["liberty", "show", name],
            capture_output=True, text=True
        )
        if result.returncode == 0:
            return result.stdout.strip().split('\n')[-1]
    except Exception:
        pass
    return None


# Tool definitions for the agent
TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "read_file",
            "description": "Read the contents of a file",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {"type": "string", "description": "File path to read"}
                },
                "required": ["path"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "run_command",
            "description": "Run a shell command",
            "parameters": {
                "type": "object",
                "properties": {
                    "command": {"type": "string", "description": "Command to execute"}
                },
                "required": ["command"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "deploy",
            "description": "Deploy the application",
            "parameters": {
                "type": "object",
                "properties": {
                    "target": {"type": "string", "description": "Deployment target"}
                },
                "required": ["target"]
            }
        }
    }
]


def execute_tool(name: str, args: dict) -> str:
    """Execute a tool and return result."""
    if name == "read_file":
        try:
            with open(args["path"], "r") as f:
                content = f.read()[:500]  # Limit for demo
            return f"File contents:\n{content}"
        except Exception as e:
            return f"Error reading file: {e}"

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

    elif name == "deploy":
        return f"Deployed to {args['target']} successfully (simulated)"

    return "Unknown tool"


async def run_agent():
    """Run the POE agent with LockStock authentication."""

    print("=" * 70)
    print("POE + LockStock Authenticated Agent")
    print("=" * 70)
    print()

    # Load credentials
    poe_api_key = get_liberty_secret("POE_API_KEY")
    agent_id = get_liberty_secret("LOCKSTOCK_OPENAI_TEST_V2_ID")
    agent_secret = get_liberty_secret("LOCKSTOCK_OPENAI_TEST_V2_SECRET")

    if not poe_api_key:
        print("ERROR: POE_API_KEY not found in Liberty vault")
        sys.exit(1)

    if not agent_id or not agent_secret:
        print("ERROR: Agent credentials not found in Liberty vault")
        sys.exit(1)

    print(f"Agent ID: {agent_id}")
    print(f"Secret: {agent_secret[:16]}...")
    print()

    # Initialize LockStock client
    lockstock = LockStockClient(
        agent_id=agent_id,
        secret=agent_secret,
        endpoint="https://lockstock-api-i9kp.onrender.com"
    )

    # Initialize POE client (OpenAI-compatible)
    # See: https://creator.poe.com/docs/external-applications/openai-compatible-api
    poe = OpenAI(
        api_key=poe_api_key,
        base_url="https://api.poe.com/v1"
    )

    try:
        # Bootstrap LockStock identity
        print("Bootstrapping LockStock identity...")
        bootstrap = await lockstock.bootstrap()
        print(f"  Root hash: {bootstrap.get('root_hash', 'N/A')[:32]}...")
        print(f"  Matrix: {bootstrap.get('root_matrix', {})}")
        print()

        # Authenticate initial heartbeat
        print("Sending heartbeat...")
        result = await lockstock.verify(task="Heartbeat")
        print(f"  Status: {'OK' if result.authorized else 'DENIED'}")
        print(f"  Hash: {lockstock._current_hash[:32]}...")
        print()

        # Agent conversation
        messages = [
            {
                "role": "system",
                "content": """You are a secure agent with LockStock authentication.
Every tool you use is cryptographically verified and logged.
Available tools: read_file, run_command, deploy.
Keep responses brief."""
            },
            {
                "role": "user",
                "content": "Please check the current directory with ls, then read the pyproject.toml file."
            }
        ]

        print("=" * 70)
        print("Starting Agent Conversation")
        print("=" * 70)
        print()
        print(f"User: {messages[-1]['content']}")
        print()

        # POE model names - see https://poe.com for available models
        # Examples: Claude-Sonnet-4, GPT-4o, Gemini-2.5-Pro, Grok-4
        model = "Claude-Sonnet-4"  # Claude model for tool use

        max_turns = 5
        turn = 0

        while turn < max_turns:
            turn += 1
            print(f"--- Turn {turn} ---")

            try:
                response = poe.chat.completions.create(
                    model=model,
                    messages=messages,
                    tools=TOOLS,
                    tool_choice="auto"
                )
            except Exception as e:
                print(f"POE API error: {e}")
                print("Note: POE API endpoint may require different configuration")
                break

            message = response.choices[0].message

            # Check for 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: {tool_name}({tool_args})")

                    # LockStock authentication for this tool call
                    capability = {
                        "read_file": "FileRead",
                        "run_command": "Shell",
                        "deploy": "Deploy"
                    }.get(tool_name, "All")

                    auth_result = await lockstock.verify(task=capability)

                    status = "AUTHORIZED" if auth_result.authorized else "DENIED"
                    print(f"  LockStock: [{status}]")
                    print(f"  Hash: {lockstock._current_hash[:32]}...")
                    print(f"  Sequence: {lockstock._sequence}")

                    if not auth_result.authorized:
                        tool_result = f"DENIED: Agent lacks {capability} capability"
                    else:
                        tool_result = execute_tool(tool_name, tool_args)

                    print(f"  Result: {tool_result[:100]}...")
                    print()

                    messages.append({
                        "role": "tool",
                        "tool_call_id": tool_call.id,
                        "content": tool_result
                    })

            else:
                # No tool calls, just a response
                print(f"  Assistant: {message.content}")
                print()

                if message.content:
                    break  # Conversation complete

        print("=" * 70)
        print("Session Summary")
        print("=" * 70)
        print(f"Total actions: {lockstock._sequence}")
        print(f"Final hash: {lockstock._current_hash}")
        print(f"Final matrix: {lockstock._state_matrix.to_dict()}")
        print()
        print("Every tool call was cryptographically authenticated.")
        print("The hash chain provides irrefutable audit trail.")

    finally:
        await lockstock.close()


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