import json
import argparse
import os
from pathlib import Path
from tqdm import tqdm
import random


def convert_nq_to_invarrag(input_path: str, output_path: str, max_samples: int = None):
    print(f"Converting NQ dataset from {input_path}")
    
    import csv
    data = []
    
    with open(input_path, 'r', encoding='utf-8') as f:
        reader = csv.reader(f, delimiter='\t')
        for idx, row in enumerate(tqdm(reader)):
            if max_samples and idx >= max_samples:
                break
            
            if len(row) >= 2:
                question = row[0]
                answers = eval(row[1])
                
                data.append({
                    'question': question,
                    'answer': answers,
                    'answers': answers,
                    'positive_ctxs': [],
                    'negative_ctxs': [],
                })
    
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(data, f, indent=2, ensure_ascii=False)
    
    print(f"Converted {len(data)} examples to {output_path}")


def convert_triviaqa_to_invarrag(input_path: str, output_path: str, max_samples: int = None):
    print(f"Converting TriviaQA dataset from {input_path}")
    
    import csv
    data = []
    
    with open(input_path, 'r', encoding='utf-8') as f:
        reader = csv.reader(f, delimiter='\t')
        for idx, row in enumerate(tqdm(reader)):
            if max_samples and idx >= max_samples:
                break
            
            if len(row) >= 2:
                question = row[0]
                answers = eval(row[1])
                
                data.append({
                    'question': question,
                    'answer': answers,
                    'answers': answers,
                    'positive_ctxs': [],
                    'negative_ctxs': [],
                })
    
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(data, f, indent=2, ensure_ascii=False)
    
    print(f"Converted {len(data)} examples to {output_path}")


def build_corpus_from_wikipedia(
    wikipedia_path: str,
    output_path: str,
    max_docs: int = 100000,
    doc_length: int = 500
):
    print(f"Building corpus from {wikipedia_path}")
    
    with open(output_path, 'w', encoding='utf-8') as out_f:
        doc_count = 0
        
        with open(wikipedia_path, 'r', encoding='utf-8') as in_f:
            for line in tqdm(in_f):
                if doc_count >= max_docs:
                    break
                
                doc = json.loads(line)
                text = doc.get('text', '')
                
                words = text.split()
                if len(words) > doc_length:
                    text = ' '.join(words[:doc_length])
                
                if len(text) > 100:
                    out_doc = {
                        'id': f"doc_{doc_count}",
                        'text': text,
                        'title': doc.get('title', ''),
                    }
                    out_f.write(json.dumps(out_doc, ensure_ascii=False) + '\n')
                    doc_count += 1
    
    print(f"Built corpus with {doc_count} documents at {output_path}")


def add_negative_samples(
    data_path: str,
    corpus_path: str,
    output_path: str,
    num_negatives: int = 5
):
    print(f"Adding negative samples to {data_path}")
    
    with open(data_path, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    corpus = []
    with open(corpus_path, 'r', encoding='utf-8') as f:
        for line in f:
            doc = json.loads(line)
            corpus.append(doc)
    
    print(f"Loaded {len(corpus)} documents from corpus")
    
    for item in tqdm(data):
        negative_docs = random.sample(corpus, min(num_negatives, len(corpus)))
        item['negative_ctxs'] = [
            {'text': doc['text'], 'title': doc.get('title', '')}
            for doc in negative_docs
        ]
    
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(data, f, indent=2, ensure_ascii=False)
    
    print(f"Saved augmented data to {output_path}")


def split_dataset(
    input_path: str,
    output_dir: str,
    train_ratio: float = 0.8,
    val_ratio: float = 0.1,
    test_ratio: float = 0.1,
    seed: int = 42
):
    print(f"Splitting dataset {input_path}")
    
    with open(input_path, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    random.seed(seed)
    random.shuffle(data)
    
    total = len(data)
    train_end = int(total * train_ratio)
    val_end = train_end + int(total * val_ratio)
    
    train_data = data[:train_end]
    val_data = data[train_end:val_end]
    test_data = data[val_end:]
    
    os.makedirs(output_dir, exist_ok=True)
    
    with open(os.path.join(output_dir, 'train.json'), 'w') as f:
        json.dump(train_data, f, indent=2, ensure_ascii=False)
    
    with open(os.path.join(output_dir, 'val.json'), 'w') as f:
        json.dump(val_data, f, indent=2, ensure_ascii=False)
    
    with open(os.path.join(output_dir, 'test.json'), 'w') as f:
        json.dump(test_data, f, indent=2, ensure_ascii=False)
    
    print(f"Split into: train={len(train_data)}, val={len(val_data)}, test={len(test_data)}")


def create_sample_data(output_dir: str, num_samples: int = 100):
    print(f"Creating sample data with {num_samples} examples")
    
    os.makedirs(output_dir, exist_ok=True)
    
    sample_qa = [
        ("What is the capital of France?", ["Paris"]),
        ("Who invented the telephone?", ["Alexander Graham Bell", "Bell"]),
        ("What is the largest planet?", ["Jupiter"]),
        ("Who wrote Romeo and Juliet?", ["William Shakespeare", "Shakespeare"]),
        ("What is the speed of light?", ["299,792,458 m/s", "approximately 300,000 km/s"]),
    ]
    
    sample_docs = [
        "Paris is the capital and most populous city of France.",
        "Alexander Graham Bell was a Scottish-born inventor who is credited with inventing the first practical telephone.",
        "Jupiter is the largest planet in our Solar System.",
        "William Shakespeare wrote Romeo and Juliet in the early years of his career.",
        "The speed of light in vacuum is exactly 299,792,458 metres per second.",
        "London is the capital of England and the United Kingdom.",
        "Thomas Edison was an American inventor who developed many devices.",
    ]
    
    train_data = []
    for i in range(num_samples):
        qa = sample_qa[i % len(sample_qa)]
        train_data.append({
            'question': qa[0],
            'answer': qa[1],
            'answers': qa[1],
            'positive_ctxs': [
                {'text': sample_docs[i % len(sample_docs)]}
            ],
            'negative_ctxs': [
                {'text': sample_docs[(i + j) % len(sample_docs)]}
                for j in range(1, 3)
            ]
        })
    
    with open(os.path.join(output_dir, 'train.json'), 'w') as f:
        json.dump(train_data, f, indent=2, ensure_ascii=False)
    
    val_data = train_data[:max(10, num_samples // 10)]
    with open(os.path.join(output_dir, 'val.json'), 'w') as f:
        json.dump(val_data, f, indent=2, ensure_ascii=False)
    
    with open(os.path.join(output_dir, 'corpus.jsonl'), 'w') as f:
        for i, doc in enumerate(sample_docs):
            f.write(json.dumps({
                'id': f'doc_{i}',
                'text': doc
            }, ensure_ascii=False) + '\n')
    
    print(f"Created sample data at {output_dir}")


def main():
    parser = argparse.ArgumentParser(description="Prepare Invar-RAG training data")
    
    subparsers = parser.add_subparsers(dest='command', help='Command to run')
    
    # Convert NQ
    nq_parser = subparsers.add_parser('convert_nq', help='Convert NQ dataset')
    nq_parser.add_argument('--input', type=str, required=True)
    nq_parser.add_argument('--output', type=str, required=True)
    nq_parser.add_argument('--max_samples', type=int, default=None)
    
    # Convert TriviaQA
    tqa_parser = subparsers.add_parser('convert_triviaqa', help='Convert TriviaQA dataset')
    tqa_parser.add_argument('--input', type=str, required=True)
    tqa_parser.add_argument('--output', type=str, required=True)
    tqa_parser.add_argument('--max_samples', type=int, default=None)
    
    # Build corpus
    corpus_parser = subparsers.add_parser('build_corpus', help='Build document corpus')
    corpus_parser.add_argument('--input', type=str, required=True)
    corpus_parser.add_argument('--output', type=str, required=True)
    corpus_parser.add_argument('--max_docs', type=int, default=100000)
    
    # Add negatives
    neg_parser = subparsers.add_parser('add_negatives', help='Add negative samples')
    neg_parser.add_argument('--data', type=str, required=True)
    neg_parser.add_argument('--corpus', type=str, required=True)
    neg_parser.add_argument('--output', type=str, required=True)
    neg_parser.add_argument('--num_negatives', type=int, default=5)
    
    # Split dataset
    split_parser = subparsers.add_parser('split', help='Split dataset')
    split_parser.add_argument('--input', type=str, required=True)
    split_parser.add_argument('--output_dir', type=str, required=True)
    split_parser.add_argument('--train_ratio', type=float, default=0.8)
    split_parser.add_argument('--val_ratio', type=float, default=0.1)
    
    # Create sample
    sample_parser = subparsers.add_parser('create_sample', help='Create sample data')
    sample_parser.add_argument('--output_dir', type=str, required=True)
    sample_parser.add_argument('--num_samples', type=int, default=100)
    
    args = parser.parse_args()
    
    if args.command == 'convert_nq':
        convert_nq_to_invarrag(args.input, args.output, args.max_samples)
    elif args.command == 'convert_triviaqa':
        convert_triviaqa_to_invarrag(args.input, args.output, args.max_samples)
    elif args.command == 'build_corpus':
        build_corpus_from_wikipedia(args.input, args.output, args.max_docs)
    elif args.command == 'add_negatives':
        add_negative_samples(args.data, args.corpus, args.output, args.num_negatives)
    elif args.command == 'split':
        split_dataset(args.input, args.output_dir, args.train_ratio, args.val_ratio)
    elif args.command == 'create_sample':
        create_sample_data(args.output_dir, args.num_samples)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()