"""Prediction entry point for COMBO."""

import argparse
import re
from pathlib import Path

from combo.config import load_config
from combo.data import CoNLLUReader
from combo.data.conllu_writer import CoNLLUWriter
from combo.inference import Predictor
from combo.utils import get_device


def tokenize(text: str) -> list[str]:
    """Simple tokenizer that separates punctuation from words."""
    text = re.sub(r'([.,!?;:\(\)\[\]"\'„"«»–—])', r' \1 ', text)
    return [token for token in text.split() if token]


def main():
    parser = argparse.ArgumentParser(description="Run predictions with COMBO")

    # Model source (mutually exclusive)
    model_source = parser.add_mutually_exclusive_group(required=True)
    model_source.add_argument(
        "--model-dir",
        type=Path,
        help="Path to exported model directory (contains config.json, pytorch_model.bin, encoders/)",
    )
    model_source.add_argument(
        "--task-config",
        type=Path,
        help="Path to task-specific configuration override (requires --checkpoint)",
    )

    parser.add_argument(
        "--config",
        type=Path,
        default=Path("config/base.yaml"),
        help="Path to base configuration file (used with --task-config)",
    )
    parser.add_argument(
        "--checkpoint",
        type=Path,
        default=None,
        help="Path to model checkpoint (required with --task-config)",
    )
    parser.add_argument(
        "--input",
        type=Path,
        default=None,
        help="Input CoNLL-U file (if not provided, uses interactive mode)",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=None,
        help="Output CoNLL-U file for predictions",
    )
    parser.add_argument(
        "--text",
        type=str,
        default=None,
        help="Text to parse (space-separated tokens)",
    )
    args = parser.parse_args()

    if args.model_dir:
        # === Load from exported model directory ===
        model_dir = args.model_dir
        print(f"Loading model from exported directory: {model_dir}")

        config = load_config(model_dir / "config.json")
        device = get_device(config.device)
        max_length = config.max_length

        predictor = Predictor.from_checkpoint(
            checkpoint_path=model_dir / "pytorch_model.bin",
            encoders_dir=model_dir / "encoders",
            config=config.model,
            device=device,
            max_length=max_length,
        )
    else:
        # === Load from task-config + checkpoint ===
        if not args.checkpoint:
            parser.error("--checkpoint is required when using --task-config")

        print(f"Loading configuration from {args.config}")
        config = load_config(args.config, args.task_config)
        device = get_device(config.device)

        predictor = Predictor.from_checkpoint(
            checkpoint_path=args.checkpoint,
            encoders_dir=config.language_output_dir / "encoders",
            config=config.model,
            device=device,
            max_length=config.data.max_length,
        )

    print("Model loaded successfully!")

    # Process input
    if args.input:
        # Process CoNLL-U file
        print(f"\nProcessing {args.input}...")
        doc = CoNLLUReader.read(args.input, treebank_id="input", language="unknown")

        predictions = []
        for i, sentence in enumerate(doc.sentences):
            words = [t.form for t in sentence.words]
            pred = predictor.predict(words, sent_id=sentence.sent_id)
            pred.text = sentence.text
            predictions.append(pred)

            if (i + 1) % 100 == 0:
                print(f"  Processed {i + 1}/{len(doc.sentences)} sentences")

        print(f"  Processed {len(predictions)} sentences")

        # Save output
        if args.output:
            CoNLLUWriter.write(predictions, args.output)
            print(f"\nPredictions saved to {args.output}")
        else:
            # Print to stdout
            for pred in predictions:
                print(pred.to_conllu())

    elif args.text:
        # Process single text
        words = tokenize(args.text)
        print(f"\nParsing: {args.text}")

        pred = predictor.predict(words)
        print("\n" + pred.to_conllu())

        if args.output:
            CoNLLUWriter.write([pred], args.output)
            print(f"Saved to {args.output}")

    else:
        # Interactive mode
        print("\nInteractive mode - enter text or 'quit' to exit")
        print("-" * 50)

        sent_id = 0
        while True:
            try:
                text = input("\nEnter text: ").strip()
            except EOFError:
                break

            if text.lower() in ("quit", "exit", "q"):
                break

            if not text:
                continue

            words = tokenize(text)
            pred = predictor.predict(words, sent_id=f"interactive_{sent_id}")
            sent_id += 1

            print("\n" + pred.to_conllu())


if __name__ == "__main__":
    main()
