import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread
import re

MODEL_PATH = "VDrontV2-0.1b-Title"
TEMPERATURE = 0.4
MAX_NEW_TOKENS = 64
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

STOP_TOKENS = ["<|endoftext|>", "(end_title)", "(title)", "\n\n\n"]
STOP_PATTERNS = [
    r"\(end_title\)",
    r"\(title\)",
    r"<\|endoftext\|>",
]


def load_model_and_tokenizer(model_path):
    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=False)
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
        device_map="auto",
        trust_remote_code=False
    )
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    if tokenizer.eos_token is None:
        tokenizer.eos_token = "<|endoftext|>"
        tokenizer.eos_token_id = tokenizer.convert_tokens_to_ids("<|endoftext|>")

    return model, tokenizer


def should_stop(text, accumulated_text=""):
    full_text = accumulated_text + text

    for stop_token in STOP_TOKENS:
        if stop_token in full_text:
            return True, full_text.split(stop_token)[0]

    if len(full_text) > 100:
        last_part = full_text[-50:]
        if len(set(last_part.split())) <= 3:
            return True, full_text

    if len(full_text) > 60:
        for i in range(10, 30):
            pattern = full_text[-i:]
            if full_text.count(pattern) > 3:
                return True, full_text

    return False, full_text


def generate_stream(model, tokenizer, prompt, temperature=0.4, max_new_tokens=256):
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048)
    inputs = {k: v.to(model.device) for k, v in inputs.items()}

    streamer = TextIteratorStreamer(
        tokenizer,
        skip_prompt=True,
        skip_special_tokens=False,
        timeout=30.0
    )

    generation_kwargs = dict(
        **inputs,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        do_sample=True,
        top_p=0.96,
        repetition_penalty=1.15,
        pad_token_id=tokenizer.pad_token_id,
        eos_token_id=tokenizer.eos_token_id,
        streamer=streamer,
    )

    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()

    accumulated_text = ""

    try:
        for new_text in streamer:
            if new_text:
                accumulated_text += new_text

                stop, clean_text = should_stop(new_text,
                                               accumulated_text[:-len(new_text)] if len(accumulated_text) > len(
                                                   new_text) else "")

                if stop:
                    if clean_text and len(clean_text) > len(accumulated_text) - len(new_text):
                        yield clean_text
                    break

                partial_stop = False
                for stop_token in STOP_TOKENS:
                    for i in range(1, len(stop_token)):
                        if accumulated_text.endswith(stop_token[:i]):
                            partial_stop = True
                            break
                    if partial_stop:
                        break

                if not partial_stop:
                    yield new_text

                if len(accumulated_text) > 50 and len(accumulated_text.split()) > 10:
                    if any(pattern in accumulated_text for pattern in ["(end_title)", "(title)"]):
                        break
    except Exception as e:
        print(f"\n[Stream error: {e}]")

    thread.join(timeout=5)


def text_continuation(model, tokenizer, temperature):
    print(f"Instruct Mode (temp={temperature})")
    print("Enter your text and the model will continue it.")
    print("Commands: 'exit' or 'quit' — exit.")
    print(f"Stop tokens: {', '.join(STOP_TOKENS)}")
    print("=" * 60)

    while True:
        try:
            user_input = input("\nUser: ").strip()
        except (KeyboardInterrupt, EOFError):
            print("\nGoodbye!")
            break

        if user_input.lower() in ["exit", "quit"]:
            print("Goodbye!")
            break

        if not user_input:
            continue

        print("Title: ", end="", flush=True)

        formatted_prompt = f"(user){user_input}(title)"

        response_text = ""
        for token in generate_stream(model, tokenizer, formatted_prompt, temperature=temperature,
                                     max_new_tokens=MAX_NEW_TOKENS):
            print(token, end="", flush=True)
            response_text += token

            if len(response_text.split()) > 50:
                break

        print()

        if response_text:
            print(f"[Generated {len(response_text.split())} words]")


if __name__ == "__main__":
    print(f"Loading model from {MODEL_PATH}...")
    model, tokenizer = load_model_and_tokenizer(MODEL_PATH)
    print(f"Model loaded. Device: {DEVICE}")
    print(f"Tokenizer eos_token: '{tokenizer.eos_token}'")
    print(f"Tokenizer eos_token_id: {tokenizer.eos_token_id}")

    text_continuation(model, tokenizer, temperature=TEMPERATURE)