#!/usr/bin/env python3

"""

Swarm Provenance Tracker

A CLI for tracking agent contributions with cryptographic accountability.

"""

import argparse

import json

import hashlib

import time

from datetime import datetime

from pathlib import Path

from typing import Dict, List, Optional

import sys

class SwarmTracker:

    """Manages swarm workspace and provenance tracking."""

    def __init__(self, workspace_path: str = ".swarm"):

        self.workspace = Path(workspace_path)

        self.config_file = self.workspace / "config.json"

        self.log_file = self.workspace / "provenance.jsonl"

        self.agents_file = self.workspace / "agents.json"

    def init(self, name: str):

        """Initialize a new swarm workspace."""

        if self.workspace.exists():

            print(f"❌ Swarm workspace already exists at {self.workspace}")

            return False

        self.workspace.mkdir(parents=True)

        config = {

            "name": name,

            "created_at": datetime.utcnow().isoformat(),

            "version": "1.0.0"

        }

        with open(self.config_file, 'w') as f:

            json.dump(config, f, indent=2)

        # Initialize empty agents registry

        with open(self.agents_file, 'w') as f:

            json.dump({"agents": {}}, f, indent=2)

        # Create empty provenance log

        self.log_file.touch()

        print(f"✅ Initialized swarm workspace: {name}")

        print(f"📁 Location: {self.workspace.absolute()}")

        return True

    def add_agent(self, name: str, role: str, key: Optional[str] = None):

        """Register an agent in the swarm."""

        if not self.workspace.exists():

            print("❌ No swarm workspace found. Run 'swarm.py init' first.")

            return False

        with open(self.agents_file, 'r') as f:

            data = json.load(f)

        if name in data["agents"]:

            print(f"❌ Agent '{name}' already registered")

            return False

        agent_key = key or self._generate_agent_key(name)

        data["agents"][name] = {

            "role": role,

            "public_key": agent_key,

            "registered_at": datetime.utcnow().isoformat(),

            "contributions": 0

        }

        with open(self.agents_file, 'w') as f:

            json.dump(data, f, indent=2)

        print(f"✅ Registered agent: {name}")

        print(f"🔑 Key: {agent_key[:16]}...")

        return True

    def record_contribution(self, agent: str, file_path: str, message: str, sign: bool = True):

        """Record a contribution to the provenance log."""

        if not self.workspace.exists():

            print("❌ No swarm workspace found.")

            return False

        # Load agents to verify

        with open(self.agents_file, 'r') as f:

            agents_data = json.load(f)

        if agent not in agents_data["agents"]:

            print(f"❌ Agent '{agent}' not registered. Add with 'swarm.py agent add'")

            return False

        # Get previous hash for chaining

        prev_hash = self._get_last_hash()

        # Create contribution record

        contribution = {

            "agent": agent,

            "file": file_path,

            "message": message,

            "timestamp": datetime.utcnow().isoformat(),

            "prev_hash": prev_hash

        }

        # Calculate hash of this contribution

        contrib_hash = self._hash_contribution(contribution)

        contribution["hash"] = contrib_hash

        if sign:

            # Simple signature simulation (in production, use real crypto)

            agent_key = agents_data["agents"][agent]["public_key"]

            contribution["signature"] = self._sign(contrib_hash, agent_key)

        # Append to log

        with open(self.log_file, 'a') as f:

            f.write(json.dumps(contribution) + '\n')

        # Update agent contribution count

        agents_data["agents"][agent]["contributions"] += 1

        with open(self.agents_file, 'w') as f:

            json.dump(agents_data, f, indent=2)

        print(f"✅ Recorded contribution by {agent}")

        print(f"📝 {message}")

        print(f"🔗 Hash: {contrib_hash[:16]}...")

        return True

    def show_history(self, limit: int = 20):

        """Display provenance history."""

        if not self.log_file.exists():

            print("❌ No provenance log found.")

            return

        with open(self.log_file, 'r') as f:

            contributions = [json.loads(line) for line in f]

        if not contributions:

            print("📭 No contributions recorded yet.")

            return

        print(f"\n🔗 Provenance Chain ({len(contributions)} contributions)\n")

        for contrib in contributions[-limit:]:

            timestamp = contrib["timestamp"][:19].replace('T', ' ')

            print(f"┌─ {contrib['agent']} @ {timestamp}")

            print(f"│  📄 {contrib['file']}")

            print(f"│  💬 {contrib['message']}")

            print(f"│  🔗 {contrib['hash'][:16]}...")

            if 'signature' in contrib:

                print(f"│  ✍️  Signed")

            print("└─")

    def export(self, format: str = "json"):

        """Export provenance data."""

        if not self.log_file.exists():

            print("❌ No provenance log found.")

            return

        with open(self.log_file, 'r') as f:

            contributions = [json.loads(line) for line in f]

        if format == "json":

            print(json.dumps(contributions, indent=2))

        else:

            print("❌ Unsupported format. Use: json")

    # Helper methods

    def _generate_agent_key(self, agent_name: str) -> str:

        """Generate a simple key for an agent (demo purposes)."""

        return hashlib.sha256(f"{agent_name}-{time.time()}".encode()).hexdigest()

    def _get_last_hash(self) -> str:

        """Get the hash of the last contribution (for chaining)."""

        if not self.log_file.exists() or self.log_file.stat().st_size == 0:

            return "genesis"

        with open(self.log_file, 'r') as f:

            lines = f.readlines()

            if lines:

                last = json.loads(lines[-1])

                return last.get("hash", "unknown")

        return "genesis"

    def _hash_contribution(self, contrib: Dict) -> str:

        """Calculate hash of a contribution."""

        data = f"{contrib['agent']}{contrib['file']}{contrib['message']}{contrib['timestamp']}{contrib['prev_hash']}"

        return hashlib.sha256(data.encode()).hexdigest()

    def _sign(self, data: str, key: str) -> str:

        """Simple signature simulation (use real crypto in production)."""

        return hashlib.sha256(f"{data}{key}".encode()).hexdigest()[:32]

def main():

    parser = argparse.ArgumentParser(

        description="Swarm Provenance Tracker - Track agent contributions with cryptographic accountability"

    )

    subparsers = parser.add_subparsers(dest='command', help='Commands')

    # Init command

    init_parser = subparsers.add_parser('init', help='Initialize swarm workspace')

    init_parser.add_argument('--name', required=True, help='Swarm name')

    # Agent commands

    agent_parser = subparsers.add_parser('agent', help='Manage agents')

    agent_subparsers = agent_parser.add_subparsers(dest='agent_command')

    add_agent_parser = agent_subparsers.add_parser('add', help='Add agent to swarm')

    add_agent_parser.add_argument('--name', required=True, help='Agent name')

    add_agent_parser.add_argument('--role', required=True, help='Agent role')

    add_agent_parser.add_argument('--key', help='Public key (auto-generated if not provided)')

    # Contribute command

    contrib_parser = subparsers.add_parser('contribute', help='Record a contribution')

    contrib_parser.add_argument('--agent', required=True, help='Agent name')

    contrib_parser.add_argument('--file', required=True, help='File path')

    contrib_parser.add_argument('--message', required=True, help='Contribution message')

    contrib_parser.add_argument('--sign', action='store_true', help='Sign the contribution')

    # History command

    history_parser = subparsers.add_parser('history', help='Show provenance history')

    history_parser.add_argument('--limit', type=int, default=20, help='Number of entries to show')

    # Export command

    export_parser = subparsers.add_parser('export', help='Export provenance data')

    export_parser.add_argument('--format', default='json', help='Export format (json)')

    args = parser.parse_args()

    if not args.command:

        parser.print_help()

        return

    tracker = SwarmTracker()

    if args.command == 'init':

        tracker.init(args.name)

    elif args.command == 'agent' and args.agent_command == 'add':

        tracker.add_agent(args.name, args.role, args.key)

    elif args.command == 'contribute':

        tracker.record_contribution(args.agent, args.file, args.message, args.sign)

    elif args.command == 'history':

        tracker.show_history(args.limit)

    elif args.command == 'export':

        tracker.export(args.format)

if __name__ == "__main__":

    main()