#!/root/venvs/train/bin/python
"""
Script for fine-tuning Mistral models for SIA using the Mistral API.
"""
from dataclasses import dataclass
from pathlib import Path
import argparse
import json
import os
import sys
import tempfile
import requests

# Import from our shared library
from .util import TrainingParams, DatasetCreator

@dataclass
class Config:
    def __init__(self):
        parser = argparse.ArgumentParser(description='Train SIA model using Mistral API')
        parser.add_argument(
            '--config',
            type=Path,
            default=Path('/root/sia/training/config.yaml'),
            help='Path to config file'
        )
        parser.add_argument(
            '--model',
            type=str,
            default='mistral-large-latest',
            help='Base model for fine-tuning'
        )
        parser.add_argument(
            '--api-key',
            type=str,
            default=os.environ.get('SIA_MISTRAL_API_KEY'),
            help='Mistral API key'
        )
        self.args = parser.parse_args()

    @property
    def config_path(self) -> Path:
        return self.args.config

    @property
    def model(self) -> str:
        return self.args.model
    
    @property
    def api_key(self) -> str:
        return self.args.api_key

def upload_file(api_key: str, file_path: Path) -> str:
    """Upload a file to the Mistral API and return the file ID"""
    url = "https://api.mistral.ai/v1/files"
    headers = {
        "Authorization": f"Bearer {api_key}"
    }
    files = {
        "file": ("dataset.jsonl", open(file_path, "rb"), "application/jsonl"),
        "purpose": (None, "fine-tune")
    }
    
    response = requests.post(url, headers=headers, files=files)
    if response.status_code != 200:
        print(f"Error uploading file: {response.text}")
        sys.exit(1)
    
    return response.json()["id"]

def start_finetune_job(api_key: str, model: str, file_id: str, params: sia_train_lib.TrainingParams):
    """Start a fine-tuning job on the Mistral API"""
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    data = {
        "model": model,
        "training_files": [{"file_id": file_id, "weight": 1}],
        "hyperparameters": {
            "learning_rate": params.learning_rate,
            "epochs": params.epochs
        }
    }
    
    response = requests.post(
        "https://api.mistral.ai/v1/fine_tuning/jobs",
        headers=headers,
        json=data
    )
    
    if response.status_code != 200:
        print(f"Error creating fine-tuning job: {response.text}")
        return None
        
    return response.json()["id"]

def main():
    config = Config()
    if not config.api_key:
        print("Error: Mistral API key not found. Set SIA_MISTRAL_API_KEY environment variable.")
        return 1
    
    training_data, train_params, commit_hash = sia_train_lib.prepare_training_data(config.config_path)
    
    if not training_data:
        print("No valid training data found. Exiting.")
        return 1
    
    model_name = f"sia_{commit_hash}"
    
    # Create temp file and upload
    with tempfile.NamedTemporaryFile(mode='w', suffix='.jsonl', delete=False) as f:
        for sample in training_data:
            json.dump(sample, f, ensure_ascii=False)
            f.write('\n')
    
    try:
        file_id = upload_file(config.api_key, Path(f.name))
        
        # Start fine-tuning job
        job_id = start_finetune_job(
            api_key=config.api_key,
            model=config.model,
            file_id=file_id,
            params=train_params
        )
        
        if not job_id:
            return 1
            
        print(f"Started fine-tuning job: {model_name}")
        print(f"Job ID: {job_id}")
        print(f"Check status: curl -H 'Authorization: Bearer {config.api_key}' https://api.mistral.ai/v1/fine_tuning/jobs/{job_id}")
    finally:
        os.unlink(f.name)
    
    return 0

if __name__ == "__main__":
    exit(main())