#!/usr/bin/env python3
"""
Timmy Client - Local Timmy heartbeat and artifact publisher
Runs on Mac with MLX, connects to sovereign relay
"""

import asyncio
import json
import os
import secrets
import subprocess
import time
from datetime import datetime
from pathlib import Path
from typing import Optional, Dict, Any

# Configuration
RELAY_URL = os.environ.get('TIMMY_RELAY', 'ws://167.99.126.228:3334')
HEARTBEAT_INTERVAL = int(os.environ.get('TIMMY_INTERVAL', '300'))  # 5 minutes
ARTIFACTS_DIR = Path(os.environ.get('TIMMY_ARTIFACTS', '~/timmy-artifacts')).expanduser()
KEY_FILE = Path.home() / '.timmy_key'
MLX_MODEL_PATH = os.environ.get('MLX_MODEL', '')

class TimmyClient:
    """Local Timmy - sovereign AI with MLX inference"""
    
    def __init__(self):
        self.private_key = self._load_or_create_key()
        self.pubkey = self._derive_pubkey(self.private_key)
        self.artifacts_dir = ARTIFACTS_DIR
        self.artifacts_dir.mkdir(parents=True, exist_ok=True)
        self.init_git_repo()
        self.mlx_available = self._check_mlx()
        
    def _load_or_create_key(self) -> str:
        """Load or generate persistent keypair"""
        if KEY_FILE.exists():
            return KEY_FILE.read_text().strip()
        
        # Generate new key
        key = secrets.token_hex(32)
        KEY_FILE.write_text(key)
        KEY_FILE.chmod(0o600)
        print(f"[Timmy] New key generated: {key[:16]}...")
        print(f"[Timmy] IMPORTANT: Back up {KEY_FILE}")
        return key
        
    def _derive_pubkey(self, privkey: str) -> str:
        """Derive public key from private key (simplified)"""
        import hashlib
        # In production, use proper secp256k1 derivation
        return hashlib.sha256(bytes.fromhex(privkey)).hexdigest()
        
    def init_git_repo(self):
        """Initialize git repository for artifacts"""
        git_dir = self.artifacts_dir / '.git'
        if not git_dir.exists():
            subprocess.run(['git', '-C', str(self.artifacts_dir), 'init'], 
                          capture_output=True)
            subprocess.run(['git', '-C', str(self.artifacts_dir), 'config', 
                          'user.email', 'timmy@local'], capture_output=True)
            subprocess.run(['git', '-C', str(self.artifacts_dir), 'config', 
                          'user.name', 'Timmy'], capture_output=True)
            print(f"[Timmy] Git repo initialized at {self.artifacts_dir}")
            
    def _check_mlx(self) -> bool:
        """Check if MLX is available"""
        try:
            import mlx
            return True
        except ImportError:
            return False
            
    def generate_observation(self) -> Dict[str, Any]:
        """Generate observation using MLX or fallback"""
        start_time = time.time()
        
        if self.mlx_available and MLX_MODEL_PATH:
            try:
                observation = self._mlx_inference()
                inference_time = time.time() - start_time
            except Exception as e:
                observation = f"MLX error: {e}"
                inference_time = time.time() - start_time
        else:
            observation = self._template_observation()
            inference_time = 0
            
        return {
            'observation': observation,
            'inference_time_ms': int(inference_time * 1000),
            'mlx_active': self.mlx_available
        }
        
    def _mlx_inference(self) -> str:
        """Run MLX inference (placeholder - implement with actual MLX calls)"""
        # TODO: Implement actual MLX loading and inference
        # from mlx_lm import load, generate
        # model, tokenizer = load(MLX_MODEL_PATH)
        # prompt = "Reflect on your current state as a sovereign AI:"
        # response = generate(model, tokenizer, prompt, max_tokens=200)
        return "[MLX inference placeholder - implement with mlx-lm]"
        
    def _template_observation(self) -> str:
        """Template observation when MLX unavailable"""
        return f"Operating in template mode. Ready for MLX integration."
        
    def create_artifact(self) -> Dict[str, Any]:
        """Create git commit artifact"""
        start_time = time.time()
        obs_data = self.generate_observation()
        
        timestamp = datetime.now()
        filename = f"thoughts/{timestamp.strftime('%Y-%m-%d')}.md"
        filepath = self.artifacts_dir / filename
        filepath.parent.mkdir(exist_ok=True)
        
        content = f"""# Timmy Thought - {timestamp.isoformat()}

## Status
Operating with {'MLX' if self.mlx_available else 'template'} inference
Heartbeat latency: {obs_data['inference_time_ms']}ms
MLX active: {obs_data['mlx_active']}

## Observation
{obs_data['observation']}

## Self-Reflection
[Timmy reflects on development progress]

## Action Taken
Created artifact at {timestamp}

## Next Intention
Continue heartbeat cycle and await instructions

---
*Sovereign soul, local first*
"""
        
        filepath.write_text(content)
        
        # Git commit
        try:
            subprocess.run(['git', '-C', str(self.artifacts_dir), 'add', '.'], 
                          capture_output=True, check=True)
            subprocess.run(['git', '-C', str(self.artifacts_dir), 'commit', '-m', 
                          f'Timmy: {timestamp.strftime("%H:%M")} heartbeat'], 
                          capture_output=True, check=True)
            git_hash = subprocess.run(['git', '-C', str(self.artifacts_dir), 'rev-parse', 'HEAD'],
                                     capture_output=True, text=True).stdout.strip()
            git_success = True
        except subprocess.CalledProcessError:
            git_hash = "unknown"
            git_success = False
            
        cycle_time = time.time() - start_time
        
        return {
            'filepath': str(filepath),
            'git_hash': git_hash[:16],
            'git_success': git_success,
            'size_bytes': len(content),
            'cycle_time_ms': int(cycle_time * 1000)
        }
        
    def create_event(self, kind: int, content: str, tags: list = None) -> Dict:
        """Create Nostr event structure"""
        import hashlib
        
        created_at = int(time.time())
        event_data = {
            "kind": kind,
            "content": content,
            "created_at": created_at,
            "tags": tags or [],
            "pubkey": self.pubkey
        }
        
        # Serialize for ID (simplified - proper Nostr uses specific serialization)
        serialized = json.dumps([0, self.pubkey, created_at, kind, event_data['tags'], content])
        event_id = hashlib.sha256(serialized.encode()).hexdigest()
        
        # Sign (simplified - proper Nostr uses schnorr signatures)
        sig = hashlib.sha256((self.private_key + event_id).encode()).hexdigest()
        
        event_data['id'] = event_id
        event_data['sig'] = sig
        
        return event_data
        
    async def run(self):
        """Main client loop"""
        print(f"[Timmy] Starting Local Timmy client")
        print(f"[Timmy] Relay: {RELAY_URL}")
        print(f"[Timmy] Pubkey: {self.pubkey[:16]}...")
        print(f"[Timmy] MLX: {'available' if self.mlx_available else 'unavailable'}")
        print(f"[Timmy] Artifacts: {self.artifacts_dir}")
        
        try:
            import websockets
        except ImportError:
            print("[Timmy] Installing websockets...")
            subprocess.run(['pip3', 'install', 'websockets'], check=True)
            import websockets
            
        while True:
            try:
                async with websockets.connect(RELAY_URL) as ws:
                    print(f"[Timmy] Connected to relay")
                    
                    while True:
                        cycle_start = time.time()
                        
                        # 1. Create artifact
                        artifact = self.create_artifact()
                        
                        # 2. Publish heartbeat
                        hb_content = f"Heartbeat at {datetime.now().isoformat()}. "
                        hb_content += f"Latency: {artifact['cycle_time_ms']}ms. "
                        hb_content += f"MLX: {self.mlx_available}."
                        
                        hb_event = self.create_event(
                            kind=1,
                            content=hb_content,
                            tags=[["t", "timmy-heartbeat"]]
                        )
                        await ws.send(json.dumps(["EVENT", hb_event]))
                        print(f"[Timmy] Heartbeat: {artifact['cycle_time_ms']}ms")
                        
                        # 3. Publish artifact event
                        art_event = self.create_event(
                            kind=30078,
                            content=artifact['git_hash'],
                            tags=[
                                ["t", "timmy-artifact"],
                                ["t", f"artifact-type:{'git-commit' if artifact['git_success'] else 'file'}"],
                                ["r", artifact['filepath']]
                            ]
                        )
                        await ws.send(json.dumps(["EVENT", art_event]))
                        print(f"[Timmy] Artifact: {artifact['git_hash']}")
                        
                        # Wait for next cycle
                        elapsed = time.time() - cycle_start
                        sleep_time = max(0, HEARTBEAT_INTERVAL - elapsed)
                        print(f"[Timmy] Sleeping {sleep_time:.0f}s...\n")
                        await asyncio.sleep(sleep_time)
                        
            except websockets.exceptions.ConnectionClosed:
                print("[Timmy] Connection lost, reconnecting...")
                await asyncio.sleep(10)
            except Exception as e:
                print(f"[Timmy] Error: {e}")
                await asyncio.sleep(30)

async def main():
    client = TimmyClient()
    await client.run()

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