#!/usr/bin/env python3
"""
Script untuk setup dan menggunakan Novita AI
"""

import os
import sys
import requests
import json
from pathlib import Path
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class NovitaAIClient:
    def __init__(self, api_key):
        self.api_key = api_key
        self.base_url = "https://api.novita.ai"
        self.headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }
    
    def test_connection(self):
        """Test koneksi ke Novita AI API"""
        try:
            response = requests.get(
                f"{self.base_url}/v1/models",
                headers=self.headers
            )
            if response.status_code == 200:
                logger.info("✅ Koneksi ke Novita AI berhasil!")
                return True
            else:
                logger.error(f"❌ Error: {response.status_code} - {response.text}")
                return False
        except Exception as e:
            logger.error(f"❌ Error koneksi: {e}")
            return False
    
    def get_available_models(self):
        """Dapatkan daftar model yang tersedia"""
        try:
            response = requests.get(
                f"{self.base_url}/v1/models",
                headers=self.headers
            )
            if response.status_code == 200:
                models = response.json()
                logger.info("📋 Model yang tersedia:")
                for model in models.get('data', []):
                    logger.info(f"  - {model.get('id', 'Unknown')}: {model.get('name', 'Unknown')}")
                return models
            else:
                logger.error(f"❌ Error: {response.status_code}")
                return None
        except Exception as e:
            logger.error(f"❌ Error: {e}")
            return None
    
    def create_fine_tuning_job(self, model_name, training_file, validation_file=None):
        """Buat fine-tuning job"""
        try:
            payload = {
                "model": model_name,
                "training_file": training_file,
                "validation_file": validation_file,
                "hyperparameters": {
                    "n_epochs": 3,
                    "batch_size": 4,
                    "learning_rate_multiplier": 1.0
                }
            }
            
            response = requests.post(
                f"{self.base_url}/v1/fine_tuning/jobs",
                headers=self.headers,
                json=payload
            )
            
            if response.status_code == 200:
                job = response.json()
                logger.info(f"✅ Fine-tuning job created: {job.get('id')}")
                return job
            else:
                logger.error(f"❌ Error: {response.status_code} - {response.text}")
                return None
                
        except Exception as e:
            logger.error(f"❌ Error: {e}")
            return None
    
    def list_fine_tuning_jobs(self):
        """List semua fine-tuning jobs"""
        try:
            response = requests.get(
                f"{self.base_url}/v1/fine_tuning/jobs",
                headers=self.headers
            )
            
            if response.status_code == 200:
                jobs = response.json()
                logger.info("📋 Fine-tuning jobs:")
                for job in jobs.get('data', []):
                    status = job.get('status', 'unknown')
                    model = job.get('model', 'unknown')
                    job_id = job.get('id', 'unknown')
                    logger.info(f"  - {job_id}: {model} ({status})")
                return jobs
            else:
                logger.error(f"❌ Error: {response.status_code}")
                return None
                
        except Exception as e:
            logger.error(f"❌ Error: {e}")
            return None
    
    def get_fine_tuning_job(self, job_id):
        """Dapatkan detail fine-tuning job"""
        try:
            response = requests.get(
                f"{self.base_url}/v1/fine_tuning/jobs/{job_id}",
                headers=self.headers
            )
            
            if response.status_code == 200:
                job = response.json()
                logger.info(f"📋 Job {job_id}:")
                logger.info(f"  Status: {job.get('status')}")
                logger.info(f"  Model: {job.get('model')}")
                logger.info(f"  Created: {job.get('created_at')}")
                return job
            else:
                logger.error(f"❌ Error: {response.status_code}")
                return None
                
        except Exception as e:
            logger.error(f"❌ Error: {e}")
            return None

def setup_novita_environment():
    """Setup environment untuk Novita AI"""
    print("🚀 Setup Novita AI Environment")
    print("=" * 40)
    
    # Check API key
    api_key = os.getenv('NOVITA_API_KEY')
    if not api_key:
        print("⚠️  NOVITA_API_KEY tidak ditemukan")
        api_key = input("Masukkan Novita AI API key: ").strip()
        if api_key:
            os.environ['NOVITA_API_KEY'] = api_key
        else:
            print("❌ API key diperlukan untuk melanjutkan")
            return None
    
    # Test connection
    client = NovitaAIClient(api_key)
    if not client.test_connection():
        print("❌ Gagal koneksi ke Novita AI")
        return None
    
    return client

def create_sample_dataset():
    """Buat sample dataset untuk fine-tuning"""
    data_dir = Path("data")
    data_dir.mkdir(exist_ok=True)
    
    # Sample data untuk fine-tuning
    sample_data = [
        {
            "messages": [
                {"role": "system", "content": "Anda adalah asisten AI yang membantu dengan pertanyaan dalam bahasa Indonesia."},
                {"role": "user", "content": "Apa itu machine learning?"},
                {"role": "assistant", "content": "Machine learning adalah cabang dari artificial intelligence yang memungkinkan komputer belajar dari data tanpa diprogram secara eksplisit."}
            ]
        },
        {
            "messages": [
                {"role": "system", "content": "Anda adalah asisten AI yang membantu dengan pertanyaan dalam bahasa Indonesia."},
                {"role": "user", "content": "Jelaskan tentang deep learning"},
                {"role": "assistant", "content": "Deep learning adalah subset dari machine learning yang menggunakan neural network dengan banyak layer untuk memproses data kompleks."}
            ]
        },
        {
            "messages": [
                {"role": "system", "content": "Anda adalah asisten AI yang membantu dengan pertanyaan dalam bahasa Indonesia."},
                {"role": "user", "content": "Bagaimana cara kerja neural network?"},
                {"role": "assistant", "content": "Neural network bekerja dengan menerima input, memproses melalui hidden layers, dan menghasilkan output berdasarkan weights yang telah dilatih."}
            ]
        }
    ]
    
    # Save sebagai JSONL
    output_file = data_dir / "training_data.jsonl"
    with open(output_file, 'w', encoding='utf-8') as f:
        for item in sample_data:
            json.dump(item, f, ensure_ascii=False)
            f.write('\n')
    
    print(f"✅ Sample dataset created: {output_file}")
    return str(output_file)

def main():
    print("🤖 Novita AI Setup & Fine-tuning")
    print("=" * 50)
    
    # Setup environment
    client = setup_novita_environment()
    if not client:
        return
    
    # Get available models
    print("\n1️⃣ Getting available models...")
    models = client.get_available_models()
    
    # Create sample dataset
    print("\n2️⃣ Creating sample dataset...")
    training_file = create_sample_dataset()
    
    # Show menu
    while True:
        print("\n📋 Menu:")
        print("1. List fine-tuning jobs")
        print("2. Create fine-tuning job")
        print("3. Check job status")
        print("4. Exit")
        
        choice = input("\nPilihan (1-4): ").strip()
        
        if choice == "1":
            client.list_fine_tuning_jobs()
        elif choice == "2":
            if models and models.get('data'):
                model_id = input("Masukkan model ID: ").strip()
                job = client.create_fine_tuning_job(model_id, training_file)
                if job:
                    print(f"✅ Job created: {job.get('id')}")
            else:
                print("❌ Tidak ada model tersedia")
        elif choice == "3":
            job_id = input("Masukkan job ID: ").strip()
            client.get_fine_tuning_job(job_id)
        elif choice == "4":
            print("👋 Goodbye!")
            break
        else:
            print("❌ Pilihan tidak valid")

if __name__ == "__main__":
    main()