import json
import os
from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments

# Step 1: Convert text files to JSONL
def convert_to_jsonl(input_dir, output_file):
    with open(output_file, 'w', encoding='utf-8') as outfile:
        for filename in os.listdir(input_dir):
            if filename.endswith('.txt'):
                with open(os.path.join(input_dir, filename), 'r', encoding='utf-8') as infile:
                    text = infile.read()
                    # Split text into manageable chunks (e.g., 512 tokens)
                    chunks = [text[i:i+512] for i in range(0, len(text), 512)]
                    for chunk in chunks:
                        entry = {"text": chunk}
                        json.dump(entry, outfile, ensure_ascii=False)
                        outfile.write('\n')

# Step 2: Fine-tune Mistral model
def fine_tune_mistral(model_name, train_file, output_dir):
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)

    def tokenize_function(examples):
        return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=512)

    # Load JSONL dataset
    with open(train_file, 'r', encoding='utf-8') as f:
        lines = [json.loads(line) for line in f]
    dataset = {"text": [line["text"] for line in lines]}
    tokenized_dataset = tokenize_function(dataset)

    training_args = TrainingArguments(
        output_dir=output_dir,
        per_device_train_batch_size=4,
        num_train_epochs=3,
        save_steps=1000,
        save_total_limit=2,
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_dataset,
    )
    trainer.train()
    model.save_pretrained(output_dir)
    tokenizer.save_pretrained(output_dir)

# Step 3: Generate text as Harry Potter
def generate_as_harry(model_dir, prompt):
    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    model = AutoModelForCausalLM.from_pretrained(model_dir)
    input_ids = tokenizer.encode(
        f"You are Harry Potter. {prompt}",
        return_tensors="pt",
        max_length=512,
        truncation=True
    )
    outputs = model.generate(input_ids, max_length=200, num_return_sequences=1)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Example usage
if __name__ == "__main__":
    input_dir = "path/to/harry_potter_books"
    output_jsonl = "harry_potter.jsonl"
    model_name = "mistralai/Mixtral-8x7B-Instruct-v0.1"
    output_dir = "harry_potter_model"

    # Convert books to JSONL
    convert_to_jsonl(input_dir, output_jsonl)

    # Fine-tune model
    fine_tune_mistral(model_name, output_jsonl, output_dir)

    # Generate response as Harry Potter
    prompt = "How do you feel about facing Voldemort again?"
    response = generate_as_harry(output_dir, prompt)
    print(response)