import json
import struct
import os
import datetime

class OAM:
    def __init__(self):
        self.metadata = {
            "created_at": datetime.datetime.now().isoformat(),
            "model_name": "Unknown",
            "description": "OAM Model",
            "params_count": 0,
            "params_count_readable": "0",
            "layers": 0,
            "architecture": "custom"
        }
        self.vocab = {}
        self.weights = {}

    def save(self, filepath):
        meta_bytes = json.dumps(self.metadata).encode('utf-8')
        vocab_bytes = json.dumps(self.vocab).encode('utf-8')
        weights_bytes = json.dumps(self.weights).encode('utf-8')
        
        with open(filepath, 'wb') as f:
            f.write(b'OAM1')
            f.write(struct.pack('I', len(meta_bytes)))
            f.write(meta_bytes)
            f.write(struct.pack('I', len(vocab_bytes)))
            f.write(vocab_bytes)
            f.write(struct.pack('I', len(weights_bytes)))
            f.write(weights_bytes)

    @staticmethod
    def load(filepath):
        model = OAM()
        with open(filepath, 'rb') as f:
            magic = f.read(4)
            if magic != b'OAM1':
                raise ValueError("Invalid OAM file format")
            
            meta_len = struct.unpack('I', f.read(4))[0]
            model.metadata = json.loads(f.read(meta_len).decode('utf-8'))
            
            vocab_len = struct.unpack('I', f.read(4))[0]
            model.vocab = json.loads(f.read(vocab_len).decode('utf-8'))
            
            weights_len = struct.unpack('I', f.read(4))[0]
            model.weights = json.loads(f.read(weights_len).decode('utf-8'))
        return model