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

import argparse
import sys
from pathlib import Path

from combo.data.conllu_writer import CoNLLUWriter


def main():
    parser = argparse.ArgumentParser(
        description="Run predictions with COMBO",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Language name — auto-resolves model + LAMBO segmenter
  combo-nlp-predict --model Polish --text "Ala ma kota."

  # HuggingFace model ID or local directory
  combo-nlp-predict --model clarin-pl/combo-nlp-xlm-roberta-base-polish-pdb-ud2.17 --input data.conllu --output out.conllu

  # Raw text file (not CoNLL-U)
  combo-nlp-predict --model Polish --input raw.txt --noconllu-format --output out.conllu

  # Training checkpoint (development workflow)
  combo-nlp-predict --task-config config/tasks/pl.yaml --checkpoint best.pt --input data.conllu

  # Turn-level segmentation (e.g. for dialogue/speech)
  combo-nlp-predict --model Polish --input dialogue.txt --noconllu-format --split-level TURN
""",
    )

    # Model source
    model_source = parser.add_mutually_exclusive_group(required=True)
    model_source.add_argument(
        "--model",
        type=str,
        help="Model source: language name (e.g. 'Polish'), HuggingFace model ID "
             "(e.g. 'clarin-pl/combo-nlp-polish'), or local directory path.",
    )
    model_source.add_argument(
        "--task-config",
        type=Path,
        help="Path to task-specific configuration override (requires --checkpoint). "
             "Used for development/training workflow.",
    )

    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)",
    )

    # Input/output
    parser.add_argument(
        "--input",
        type=str,
        default=None,
        help="Input file path. CoNLL-U by default, or raw text with --noconllu-format. "
             "Use '-' to read from stdin.",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=None,
        help="Output CoNLL-U file for predictions (default: stdout)",
    )
    parser.add_argument(
        "--text",
        type=str,
        default=None,
        help="Text to parse directly (raw text, segmented by LAMBO if available)",
    )
    parser.add_argument(
        "--noconllu-format",
        action="store_true",
        help="Input file is raw text, not CoNLL-U. Requires a segmenter "
             "(use --model with a language name or a model that has LAMBO configured).",
    )

    # Segmentation options
    parser.add_argument(
        "--split-level",
        type=str,
        default=None,
        choices=["SENTENCE", "TURN"],
        help="Segmentation level: SENTENCE (default, split into sentences) or "
             "TURN (treat each turn/paragraph as one sentence). "
             "Turns are separated by double newlines in the input.",
    )
    parser.add_argument(
        "--split-multiwords",
        action=argparse.BooleanOptionalAction,
        default=None,
        help="Split multi-word tokens into subwords (default: True). "
             "Use --no-split-multiwords to keep MWTs as single tokens.",
    )

    # Other
    parser.add_argument(
        "--batch-size",
        type=int,
        default=1,
        help="Number of sentences to process in parallel (default: 1)",
    )

    args = parser.parse_args()

    # === Load model ===
    if args.model:
        nlp = _load_from_model(args.model, args.batch_size)
    else:
        nlp = _load_from_checkpoint(args, parser)

    print("Model loaded successfully!")

    # Segmenter kwargs for raw text processing
    seg_kwargs = {}
    if args.split_level is not None:
        seg_kwargs["split_level"] = args.split_level
    if args.split_multiwords is not None:
        seg_kwargs["split_multiwords"] = args.split_multiwords

    # === Process input ===
    if args.text:
        _process_text(nlp, args.text, args.output, seg_kwargs)
    elif args.input:
        if args.noconllu_format:
            _process_raw_text_file(nlp, args.input, args.output, seg_kwargs)
        else:
            _process_conllu_file(nlp, args.input, args.output)
    else:
        _interactive_mode(nlp, seg_kwargs)


def _load_from_model(model_spec: str, batch_size: int):
    """Load COMBO from a language name, HF model ID, or local path."""
    from combo.api import COMBO

    # Try as language name first
    try:
        from combo.languages import resolve_language
        resolve_language(model_spec)
        # It's a valid language name
        print(f"Loading model for language: {model_spec}")
        return COMBO(model_spec, batch_size=batch_size)
    except Exception:
        pass

    # Try as local path or HF model ID
    print(f"Loading model: {model_spec}")
    return COMBO.from_pretrained(model_spec, batch_size=batch_size)


def _load_from_checkpoint(args, parser):
    """Load COMBO from a task-config + checkpoint (development workflow)."""
    from combo.api import COMBO
    from combo.config import load_config
    from combo.inference import Predictor
    from combo.utils import get_device

    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,
    )

    return COMBO(predictor, batch_size=args.batch_size)


def _process_text(nlp, text: str, output, seg_kwargs: dict):
    """Process a single text string."""
    print(f"\nParsing: {text}")
    result = nlp(text, **seg_kwargs)

    if output:
        CoNLLUWriter.write(result.sentences, output)
        print(f"Saved to {output}")
    else:
        for sent in result.sentences:
            print(sent.to_conllu())


def _process_raw_text_file(nlp, input_path: str, output, seg_kwargs: dict):
    """Process a raw text file (not CoNLL-U)."""
    if nlp._segmenter is None:
        print("Error: --noconllu-format requires a segmenter. "
              "Use --model with a language name (e.g. --model Polish).",
              file=sys.stderr)
        sys.exit(1)

    if input_path == "-":
        text = sys.stdin.read()
        print("Reading from stdin...")
    else:
        path = Path(input_path)
        print(f"\nProcessing raw text file: {path}")
        text = path.read_text(encoding="utf-8")

    result = nlp(text, **seg_kwargs)
    print(f"  {len(result.sentences)} sentences")

    if output:
        CoNLLUWriter.write(result.sentences, output)
        print(f"Predictions saved to {output}")
    else:
        for sent in result.sentences:
            print(sent.to_conllu())


def _process_conllu_file(nlp, input_path: str, output):
    """Process a CoNLL-U file (gold tokenization, predict annotations)."""
    from combo.data import CoNLLUReader

    if input_path == "-":
        print("Error: stdin not supported for CoNLL-U input. "
              "Use a file path or --noconllu-format for raw text.",
              file=sys.stderr)
        sys.exit(1)

    path = Path(input_path)
    print(f"\nProcessing CoNLL-U file: {path}")
    doc = CoNLLUReader.read(path, treebank_id="input", language="unknown")

    # Use pre-tokenized input — pass word lists directly
    all_words = [[t.form for t in s.tokens] for s in doc.sentences]

    predictions = nlp(all_words, tokenized=True).sentences

    # Restore original sentence text
    for pred, orig in zip(predictions, doc.sentences):
        pred.text = orig.text

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

    if output:
        CoNLLUWriter.write(predictions, path=output)
        print(f"\nPredictions saved to {output}")
    else:
        for pred in predictions:
            print(pred.to_conllu())


def _interactive_mode(nlp, seg_kwargs: dict):
    """Interactive mode — read text from user, parse and print."""
    print("\nInteractive mode - enter text or 'quit' to exit")
    print("-" * 50)

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

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

        if not text:
            continue

        result = nlp(text, **seg_kwargs)
        for sent in result.sentences:
            print(sent.to_conllu())


if __name__ == "__main__":
    main()
