"""Collect the conversations Sprout learns to chat from -> data/chats.jsonl

  - sft/*.jsonl                    conversations written for Sprout's level
  - everyday-conversations (HF)    short everyday chats with an assistant
  - SimpleStories                  "tell me a story about ..." requests, built from its labels
  - TinyStories                    "continue this story: ..." requests
"""
import argparse
import glob
import json
import random
import re

import pyarrow.parquet as pq

from prepare import REPLACE


def tidy(text):
    text = text.translate(REPLACE).strip()
    return '\n'.join(' '.join(line.split()) for line in text.split('\n'))


def conversation(messages, source):
    msgs = [{'role': m['role'], 'content': tidy(m['content'])} for m in messages if m['role'] in ('user', 'assistant')]
    ok = msgs and msgs[0]['role'] == 'user' and msgs[-1]['role'] == 'assistant' and all(m['content'] for m in msgs)
    ok = ok and all(a['role'] != b['role'] for a, b in zip(msgs, msgs[1:]))
    return {'messages': msgs, 'source': source} if ok else None


STORY_ASKS = [
    'Tell me a story about {subject}.', 'Can you write a story about {subject}?', 'story about {subject} pls',
    'I want a short story about {subject}.', 'Tell me a bedtime story about {subject}.',
    'Write a story where the hero is {subject}.', 'Could you tell me a story about {subject}, please?',
    'Make up a story about {subject}!',
]
# "Once upon a time, there was a little fox named ..." -> "a little fox"
SUBJECT = re.compile(r"\bthere (?:was|lived) (a|an) ([a-z]+(?: [a-z]+){0,3}?)(?= named| who| that| in | with| called|,|\.)")


def story_requests(raw, n, rng, per_noun=120):
    """Ask for a story about exactly the hero the story is about, so request and answer agree.

    (A first version asked for SimpleStories' abstract topic labels like "bygone eras"; the stories
    only loosely matched them, and the model learned that a story need not fit the request.)
    """
    with open(f'{raw}/TinyStoriesV2-GPT4-train.txt', encoding='utf-8') as f:
        text = f.read(150_000_000)
    per = {}
    out = []
    stories = text.split('<|endoftext|>')[1:-1]
    rng.shuffle(stories)
    for s in stories:
        s = tidy(s)
        m = SUBJECT.search(s[:220])
        if not m or len(s.split()) > 170 or not 1 <= len(m.group(2).split()) <= 3:
            continue
        noun = m.group(2).split()[-1]
        if per.get(noun, 0) >= per_noun:
            continue
        per[noun] = per.get(noun, 0) + 1
        ask = rng.choice(STORY_ASKS).format(subject=f'{m.group(1)} {m.group(2)}')
        out.append(conversation([{'role': 'user', 'content': ask}, {'role': 'assistant', 'content': s}], 'story'))
        if len(out) >= n:
            break
    return out


def continuations(raw, n, rng):
    with open(f'{raw}/TinyStoriesV2-GPT4-valid.txt', encoding='utf-8') as f:
        stories = [s.strip() for s in f.read().split('<|endoftext|>') if 300 < len(s.strip()) < 1200]
    out = []
    for s in rng.sample(stories, min(n, len(stories))):
        sentences = re.split(r'(?<=[.!?])\s+', tidy(s))
        if len(sentences) < 4:
            continue
        k = rng.choice([1, 2])
        start, rest = ' '.join(sentences[:k]), ' '.join(sentences[k:])
        ask = rng.choice(['Continue this story: {s}', 'Can you finish this story? "{s}"', 'What happens next? {s}',
                          'Please continue: {s}'])
        out.append(conversation([{'role': 'user', 'content': ask.format(s=start)},
                                 {'role': 'assistant', 'content': rest}], 'continue'))
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--raw', default='raw')
    ap.add_argument('--written', default='sft')
    ap.add_argument('--out', default='data/chats.jsonl')
    ap.add_argument('--stories', type=int, default=2500)
    ap.add_argument('--continues', type=int, default=800)
    ap.add_argument('--everyday', type=int, default=1200)
    ap.add_argument('--written-repeat', type=int, default=2)
    ap.add_argument('--self-repeat', type=int, default=4, help='conversations about Sprout itself (files 08-*)')
    args = ap.parse_args()
    rng = random.Random(7)

    written = []
    for path in sorted(glob.glob(f'{args.written}/*.jsonl')):
        repeat = args.self_repeat if path.split('/')[-1].startswith('08-') else args.written_repeat
        for line in open(path):
            if line.strip():
                written += [conversation(json.loads(line)['messages'], 'written')] * repeat
    everyday = [conversation(m, 'everyday')
                for m in pq.read_table(f'{args.raw}/everyday-train.parquet').column('messages').to_pylist()]
    parts = {
        'written': [c for c in written if c],
        'everyday': rng.sample([c for c in everyday if c], args.everyday),
        'story': [c for c in story_requests(args.raw, args.stories, rng) if c],
        'continue': [c for c in continuations(args.raw, args.continues, rng) if c],
    }
    chats = [c for part in parts.values() for c in part]
    rng.shuffle(chats)
    with open(args.out, 'w') as f:
        for c in chats:
            f.write(json.dumps(c) + '\n')
    print({k: len(v) for k, v in parts.items()}, '->', len(chats))


if __name__ == '__main__':
    main()