"""Standalone prediction script for MiniRewrite.
Works on any machine with PyTorch and Tokenizers installed (CPU or CUDA).

Usage:
    # 1. Single prediction:
    python predict.py "send me report asap" --tone professional --style polite

    # 2. Friendly direct:
    python predict.py "send me report asap" --tone friendly --style direct

    # 3. Shorten concise:
    python predict.py "In order to complete the project on time, we need more resources." --task shorten --style concise

    # 4. Interactive chat mode:
    python predict.py --interactive
"""

import argparse
import sys
from pathlib import Path
import torch

# Ensure src is on path
SCRIPT_DIR = Path(__file__).resolve().parent
if str(SCRIPT_DIR / "src") not in sys.path:
    sys.path.insert(0, str(SCRIPT_DIR / "src"))

from minirewrite.tokenizer.tokenizer import RewriteTokenizer
from minirewrite.model.transformer import Seq2SeqTransformer, ModelConfig
from minirewrite.inference.generate import rewrite
from minirewrite.training.trainer import load_model_for_inference


def load_model(run_dir: Path, checkpoint_name: str, device: torch.device):
    print(f"Loading model and tokenizer from {run_dir} ({checkpoint_name}) to {device}...")
    return load_model_for_inference(run_dir, device, checkpoint_name)


def run_interactive(model, tokenizer, device):
    print("\n" + "=" * 60)
    print("MiniRewrite Interactive Mode (type 'exit' or 'quit' to stop)")
    print("=" * 60)

    tasks = ["rewrite", "simplify", "shorten"]
    tones = ["professional", "friendly", "neutral", "casual"]
    styles = ["polite", "direct", "concise", "plain"]

    current_task = "rewrite"
    current_tone = "professional"
    current_style = "polite"

    print(f"Current defaults: task={current_task}, tone={current_tone}, style={current_style}")
    print("Type '!settings' to change task/tone/style.\n")

    while True:
        try:
            text = input("Input text > ").strip()
        except (KeyboardInterrupt, EOFError):
            break

        if not text:
            continue
        if text.lower() in ("exit", "quit", "q"):
            print("Goodbye!")
            break
        if text == "!settings":
            print(f"Tasks:  {', '.join(tasks)}")
            t = input(f"Select task [{current_task}]: ").strip()
            if t in tasks:
                current_task = t

            print(f"Tones:  {', '.join(tones)}")
            to = input(f"Select tone [{current_tone}]: ").strip()
            if to in tones:
                current_tone = to

            print(f"Styles: {', '.join(styles)}")
            st = input(f"Select style [{current_style}]: ").strip()
            if st in styles:
                current_style = st
            print(f"Updated: task={current_task}, tone={current_tone}, style={current_style}\n")
            continue

        req = [{"source": text, "task": current_task, "tone": current_tone, "style": current_style}]
        out = rewrite(model, tokenizer, req, max_source_length=128, max_new_tokens=128)[0]
        print(f"Rewritten  > {out}\n")


def main():
    parser = argparse.ArgumentParser(description="Test MiniRewrite on any text.")
    parser.add_argument("text", nargs="?", help="Text to rewrite")
    parser.add_argument("--interactive", "-i", action="store_true", help="Launch interactive prompt")
    parser.add_argument("--task", default="rewrite", choices=["rewrite", "simplify", "shorten"])
    parser.add_argument("--tone", default="professional", choices=["professional", "friendly", "neutral", "casual"])
    parser.add_argument("--style", default="polite", choices=["polite", "direct", "concise", "plain"])
    parser.add_argument("--weights", type=Path, default=Path("runs/rtx4060/best.pt"))
    parser.add_argument("--tokenizer", type=Path, default=Path("runs/rtx4060"))
    parser.add_argument("--device", default="auto", choices=["auto", "cuda", "cpu"])
    args = parser.parse_args()

    # Determine device
    if args.device == "auto":
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    else:
        device = torch.device(args.device)

    run_dir = Path("runs/rtx4060") if (Path("runs/rtx4060") / "tokenizer.json").exists() else Path(".")
    checkpoint = "best.pt"
    model, tokenizer = load_model(run_dir, checkpoint, device)

    if args.interactive or not args.text:
        run_interactive(model, tokenizer, device)
    else:
        req = [{"source": args.text, "task": args.task, "tone": args.tone, "style": args.style}]
        out = rewrite(model, tokenizer, req, max_source_length=128, max_new_tokens=128)[0]
        print(f"\nOriginal : {args.text}")
        print(f"Settings : task={args.task}, tone={args.tone}, style={args.style}")
        print(f"Rewritten: {out}\n")


if __name__ == "__main__":
    main()