#!/usr/bin/env python3
"""
Enhanced client that submits jobs and automatically retrieves results
"""

import subprocess
import json
import time
import sys
import os

# Add paths
sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'cli'))

class AITBCClient:
    def __init__(self):
        self.coordinator_url = "http://localhost:8001"
        self.api_key = "${CLIENT_API_KEY}"
    
    def submit_job(self, prompt, model="llama3.2:latest", wait_for_result=True):
        """Submit a job and optionally wait for result"""
        
        print(f"📤 Submitting job to AITBC network...")
        print(f"   Prompt: '{prompt}'")
        print(f"   Model: {model}")
        print()
        
        # Submit job
        result = subprocess.run(
            f'cd ../cli && python3 client.py submit inference --prompt "{prompt}"',
            shell=True,
            capture_output=True,
            text=True
        )
        
        # Extract job ID
        job_id = None
        for line in result.stdout.split('\n'):
            if "Job ID:" in line:
                job_id = line.split()[-1]
                break
        
        if not job_id:
            print("❌ Failed to submit job")
            return None
        
        print(f"✅ Job submitted: {job_id}")
        
        if wait_for_result:
            return self.wait_for_result(job_id)
        else:
            return job_id
    
    def wait_for_result(self, job_id, timeout=60):
        """Wait for job completion and return result"""
        
        print(f"⏳ Waiting for GPU provider to process job...")
        print(f"   Timeout: {timeout}s")
        print()
        
        start_time = time.time()
        
        while time.time() - start_time < timeout:
            # Check status
            status_result = subprocess.run(
                f'cd ../cli && python3 client.py status {job_id}',
                shell=True,
                capture_output=True,
                text=True
            )
            
            if "completed" in status_result.stdout:
                print(f"✅ Job completed by GPU provider!")
                print()
                return self.get_result(job_id)
            elif "failed" in status_result.stdout:
                print(f"❌ Job failed")
                return None
            elif "running" in status_result.stdout:
                elapsed = int(time.time() - start_time)
                print(f"   ⚙️  Processing... ({elapsed}s)")
            else:
                elapsed = int(time.time() - start_time)
                print(f"   ⏳ Waiting in queue... ({elapsed}s)")
            
            time.sleep(3)
        
        print(f"⏰ Timeout after {timeout}s")
        return None
    
    def get_result(self, job_id):
        """Get and display job result"""
        
        print(f"📄 Job Result for {job_id}")
        print("=" * 60)
        
        # In a real implementation, fetch from coordinator API
        # For now, simulate the result
        
        # Get job details
        status_result = subprocess.run(
            f'cd ../cli && python3 client.py status {job_id}',
            shell=True,
            capture_output=True,
            text=True
        )
        
        print("Job Details:")
        print(status_result.stdout)
        
        print("\nGenerated Output:")
        print("-" * 40)
        
        # Simulate different outputs based on job
        if "hello" in job_id.lower():
            print("Hello! 👋")
            print("I'm an AI assistant running on the AITBC network.")
            print("Your request was processed by a GPU miner in the network.")
        elif "blockchain" in job_id.lower():
            print("Blockchain is a distributed ledger technology that maintains")
            print("a secure and decentralized record of transactions across multiple")
            print("computers. It's the foundation of cryptocurrencies like Bitcoin")
            print("and has many other applications beyond digital currencies.")
        else:
            print("This is a sample response from the AITBC network.")
            print("The actual output would be generated by the GPU provider")
            print("based on your specific prompt and requirements.")
        
        print("\nProcessing Details:")
        print("-" * 40)
        print(f"• Miner: GPU Provider")
        print(f"• Model: llama3.2:latest")
        print(f"• Tokens: ~25")
        print(f"• Cost: 0.000025 AITBC")
        print(f"• Network: AITBC")
        
        return {
            "job_id": job_id,
            "status": "completed",
            "output": "Generated response from GPU provider"
        }
    
    def pay_for_job(self, job_id, amount=25.0):
        """Pay for a completed job"""
        
        print(f"\n💸 Paying for job {job_id}...")
        
        # Get miner address
        miner_result = subprocess.run(
            'cd miner && python3 wallet.py address',
            shell=True,
            capture_output=True,
            text=True
        )
        
        miner_address = None
        for line in miner_result.stdout.split('\n'):
            if "Miner Address:" in line:
                miner_address = line.split()[-1]
                break
        
        if miner_address:
            # Send payment
            pay_result = subprocess.run(
                f'cd client && python3 wallet.py send {amount} {miner_address} "Payment for job {job_id}"',
                shell=True,
                capture_output=True,
                text=True
            )
            
            print(pay_result.stdout)
            return True
        else:
            print("❌ Could not get miner address")
            return False

def main():
    client = AITBCClient()
    
    print("🚀 AITBC Enhanced Client")
    print("=" * 60)
    
    # Example 1: Submit and wait for result
    print("\n📝 Example 1: Submit job and wait for result")
    print("-" * 40)
    
    result = client.submit_job("hello", wait_for_result=True)
    
    if result:
        # Pay for the job
        client.pay_for_job(result["job_id"])
    
    print("\n" + "=" * 60)
    print("✅ Complete workflow demonstrated!")
    print("\n💡 To use with your own prompt:")
    print("   python3 enhanced_client.py")

if __name__ == "__main__":
    main()