from unified_db import register_hf_dataset

# Register a new dataset
result = register_hf_dataset(
    repo_name="openai/gsm8k",
    dataset_type="SFT",
    name="GSM8K Math Problems",
    created_by="OpenAI"
)

if result['success']:
    dataset = result['dataset']
    print(f"Registered: {dataset['name']} with {dataset['num_tasks']} tasks")
    if result.get('exists'):
        print("Dataset already existed")

# Force update an existing dataset
result = register_hf_dataset(
    repo_name="openai/gsm8k",
    dataset_type="SFT",
    name="GSM8K Math Problems",
    created_by="OpenAI Team",
    forced_update=True  # This will update the existing entry
)

if result['success'] and result.get('updated'):
    print("Dataset updated successfully")

from unified_db import register_local_parquet

# Register a new dataset
result = register_local_parquet(
    file_path="/path/to/data.parquet",
    name="My Custom Dataset",
    created_by="researcher@company.com",
    dataset_type="RL"
)

if result['success']:
    dataset = result['dataset']
    print(f"Registered: {dataset['name']} ({dataset['num_tasks']} rows)")

# Force update an existing dataset
result = register_local_parquet(
    file_path="/path/to/updated_data.parquet",
    name="My Custom Dataset",
    created_by="researcher@company.com",
    dataset_type="RL",
    forced_update=True  # This will update the existing entry
)

if result['success'] and result.get('updated'):
    print("Dataset updated successfully")

from unified_db import get_dataset_by_name

dataset = get_dataset_by_name("GSM8K Math Problems")
if dataset:
    print(f"Found dataset: {dataset['name']}")
    print(f"Created by: {dataset['created_by']}")
    print(f"Number of tasks: {dataset['num_tasks']}")

from unified_db import register_hf_model
from datetime import datetime, timezone
from uuid import uuid4

# Register a new model
result = register_hf_model(
    repo_name="microsoft/phi-2",
    agent_id=str(uuid4()),  # Your agent's ID
    training_start=datetime.now(timezone.utc),
    name="Fine-tuned Phi-2",
    created_by="researcher@company.com",
    training_type="SFT",
    training_parameters={
        "learning_rate": 1e-5,
        "batch_size": 32,
        "epochs": 3
    }
)

if result['success']:
    model = result['model']
    print(f"Registered: {model['name']}")
    print(f"Weights: {model['weights_location']}")
    print(f"Status: {model['training_status']}")

# Force update an existing model
result = register_hf_model(
    repo_name="microsoft/phi-2",
    agent_id=str(uuid4()),
    training_start=datetime.now(timezone.utc),
    name="Fine-tuned Phi-2",
    created_by="updated@company.com",
    forced_update=True  # Update the existing entry
)

from unified_db import register_local_model

# Register a local model
result = register_local_model(
    model_path="/path/to/model/directory",
    name="Custom BERT Model",
    created_by="ml-engineer@company.com",
    agent_id=str(uuid4()),
    training_start=datetime.now(timezone.utc),
    training_type="SFT",
    training_parameters={
        "optimizer": "adamw",
        "learning_rate": 5e-5,
        "warmup_steps": 500
    }
)

if result['success']:
    model = result['model']
    print(f"Registered: {model['name']}")
    print(f"Framework: {model['training_parameters']['model_framework']}")
    print(f"Location: {model['weights_location']}")

from unified_db import get_model_by_name

model = get_model_by_name("Fine-tuned Phi-2")
if model:
    print(f"Found model: {model['name']}")
    print(f"Created by: {model['created_by']}")
    print(f"Training type: {model['training_type']}")
    print(f"Status: {model['training_status']}")

from unified_db import register_agent

# Register a new agent (only manual fields required)
result = register_agent(
    name="My Evaluation Agent",
    agent_version_hash="abcd1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcd",  # 64-char SHA-256 hash
    description="Agent for evaluating coding tasks"
)

if result['success']:
    agent = result['agent']
    print(f"Registered: {agent['name']}")
    print(f"Agent ID: {agent['id']}")
    print(f"Version: {agent['agent_version_hash']}")

# Register minimal agent (only name required)
result = register_agent(
    name="Minimal Agent",
    description="Agent with minimal configuration"
    # agent_version_hash is optional
)

# Force update an existing agent
result = register_agent(
    name="My Evaluation Agent",
    agent_version_hash="new_hash_1234567890abcdef1234567890abcdef1234567890abcdef1234567890",
    description="Updated agent description",
    forced_update=True  # Update the existing entry
)

from unified_db import get_agent_by_name

agent = get_agent_by_name("My Evaluation Agent")
if agent:
    print(f"Found agent: {agent['name']}")
    print(f"Version hash: {agent['agent_version_hash']}")
    print(f"Description: {agent['description']}")
    print(f"Last updated: {agent['updated_at']}")