export_import.py

import json
import logging
from database import SessionLocal, Provider, VirtualKey

logger = logging.getLogger("vapi-export")

def export_data(file_path: str) -> bool:
    db = SessionLocal()
    try:
        providers = db.query(Provider).all()
        v_keys = db.query(VirtualKey).all()
        
        data = {
            "providers": [
                {
                    "name": p.name,
                    "api_key_encrypted": p.api_key_encrypted,
                    "is_active": p.is_active
                } for p in providers
            ],
            "virtual_keys": [
                {
                    "key": vk.key,
                    "name": vk.name,
                    "provider_name": vk.provider.name if vk.provider else None,
                    "is_active": vk.is_active,
                    "daily_budget": vk.daily_budget,
                    "monthly_budget": vk.monthly_budget,
                    "allowed_models": vk.allowed_models
                } for vk in v_keys
            ]
        }
        
        with open(file_path, "w") as f:
            json.dump(data, f, indent=4)
        return True
    except Exception as e:
        logger.error(f"Failed to export data: {e}")
        return False
    finally:
        db.close()

def import_data(file_path: str) -> bool:
    db = SessionLocal()
    try:
        with open(file_path, "r") as f:
            data = json.load(f)
            
        for p_data in data.get("providers", []):
            existing_p = db.query(Provider).filter(Provider.name == p_data["name"]).first()
            if not existing_p:
                p = Provider(
                    name=p_data["name"],
                    api_key_encrypted=p_data["api_key_encrypted"],
                    is_active=p_data["is_active"]
                )
                db.add(p)
                db.commit()
            else:
                existing_p.api_key_encrypted = p_data["api_key_encrypted"]
                existing_p.is_active = p_data["is_active"]
                db.commit()

        for vk_data in data.get("virtual_keys", []):
            existing_vk = db.query(VirtualKey).filter(VirtualKey.key == vk_data["key"]).first()
            prov = db.query(Provider).filter(Provider.name == vk_data["provider_name"]).first()
            if not prov:
                logger.warning(f"Skipping key '{vk_data['name']}' as provider '{vk_data['provider_name']}' is missing.")
                continue
                
            if not existing_vk:
                vk = VirtualKey(
                    key=vk_data["key"],
                    name=vk_data["name"],
                    provider_id=prov.id,
                    is_active=vk_data["is_active"],
                    daily_budget=vk_data.get("daily_budget"),
                    monthly_budget=vk_data.get("monthly_budget"),
                    allowed_models=vk_data.get("allowed_models")
                )
                db.add(vk)
            else:
                existing_vk.name = vk_data["name"]
                existing_vk.provider_id = prov.id
                existing_vk.is_active = vk_data["is_active"]
                existing_vk.daily_budget = vk_data.get("daily_budget")
                existing_vk.monthly_budget = vk_data.get("monthly_budget")
                existing_vk.allowed_models = vk_data.get("allowed_models")
            db.commit()
            
        return True
    except Exception as e:
        logger.error(f"Failed to import data: {e}")
        db.rollback()
        return False
    finally:
        db.close()