//! GLiNER2 ONNX inference example using the `ort` crate.
//!
//! Cargo.toml dependencies:
//! ```toml
//! [dependencies]
//! ort = "2.0.0-rc.12"
//! tokenizers = { version = "0.21", default-features = false, features = ["fancy-regex"] }
//! regex = "1"
//! ```
//!
//! Usage:
//!     cargo run --release

use ort::session::Session;
use ort::value::Tensor;
use regex::Regex;
use tokenizers::Tokenizer;

const MODEL_PATH: &str = "model.onnx";
const TOKENIZER_PATH: &str = "tokenizer.json";
const THRESHOLD: f32 = 0.5;
const MAX_WIDTH: usize = 8;

const TEXT: &str = "Steve Jobs founded Apple Inc. in Cupertino, California on April 1, 1976.";
const LABELS: &[&str] = &["person", "company", "city", "date"];

/// A word with its character span in the original text.
struct Word {
    text: String,
    start: usize,
    end: usize,
}

/// An extracted entity.
struct Entity {
    label: String,
    text: String,
    start: usize,
    end: usize,
    score: f32,
}

/// Split text into words using the WhitespaceTokenSplitter regex.
/// Words are lowercased to match GLiNER2's preprocessing.
fn split_words(text: &str) -> Vec<Word> {
    let re = Regex::new(
        r"(?i)(?:https?://[^\s]+|www\.[^\s]+)|[a-z0-9._%+-]+@[a-z0-9.-]+\.[a-z]{2,}|@[a-z0-9_]+|\w+(?:[-_]\w+)*|\S"
    ).unwrap();

    let lowered = text.to_lowercase();
    re.find_iter(&lowered)
        .map(|m| Word {
            text: m.as_str().to_string(),
            start: m.start(),
            end: m.end(),
        })
        .collect()
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Load tokenizer and model
    let tokenizer = Tokenizer::from_file(TOKENIZER_PATH)
        .map_err(|e| format!("Failed to load tokenizer: {e}"))?;
    let mut session = Session::builder()?.commit_from_file(MODEL_PATH)?;

    // Split text into words
    let words = split_words(TEXT);
    let word_strings: Vec<&str> = words.iter().map(|w| w.text.as_str()).collect();

    // Build schema tokens: ( [P] entities ( [E] label1 [E] label2 ... ) )
    let mut schema_tokens: Vec<String> = vec![
        "(".into(),
        "[P]".into(),
        "entities".into(),
        "(".into(),
    ];
    for label in LABELS {
        schema_tokens.push("[E]".into());
        for part in label.split_whitespace() {
            schema_tokens.push(part.into());
        }
    }
    schema_tokens.push(")".into());
    schema_tokens.push(")".into());

    let num_schema_words = schema_tokens.len() + 1; // +1 for [SEP_TEXT]

    // Full pre-tokenized sequence: schema + [SEP_TEXT] + words
    let mut full_sequence: Vec<&str> = schema_tokens.iter().map(|s| s.as_str()).collect();
    full_sequence.push("[SEP_TEXT]");
    full_sequence.extend_from_slice(&word_strings);

    // Tokenize with is_pretokenized (Vec<&str> → InputSequence::PreTokenized)
    let encoding = tokenizer
        .encode(full_sequence, false)
        .map_err(|e| format!("Tokenization failed: {e}"))?;

    let token_ids: Vec<i64> = encoding.get_ids().iter().map(|&id| id as i64).collect();
    let raw_word_ids = encoding.get_word_ids();
    let seq_len = token_ids.len();
    let num_words = words.len();

    // input_ids: (1, seq_len)
    let input_ids = Tensor::from_array(
        (vec![1i64, seq_len as i64], token_ids),
    )?;

    // attention_mask: (1, seq_len)
    let attention_mask = Tensor::from_array(
        (vec![1i64, seq_len as i64], vec![1i64; seq_len]),
    )?;

    // text_positions: (num_words,)
    let mut text_positions_vec = Vec::with_capacity(num_words);
    for wi in 0..num_words {
        let full_word_idx = (num_schema_words + wi) as u32;
        let pos = raw_word_ids
            .iter()
            .position(|&wid| wid == Some(full_word_idx))
            .ok_or_else(|| format!("Word {wi} ('{}') not found in token mapping", words[wi].text))?;
        text_positions_vec.push(pos as i64);
    }
    let text_positions = Tensor::from_array(
        (vec![num_words as i64], text_positions_vec),
    )?;

    // schema_positions: [P] then each [E]
    let mut schema_positions_vec = Vec::new();
    for (i, tok) in schema_tokens.iter().enumerate() {
        if tok == "[P]" || tok == "[E]" {
            let idx = i as u32;
            let pos = raw_word_ids
                .iter()
                .position(|&wid| wid == Some(idx))
                .ok_or_else(|| format!("Schema token '{tok}' at index {i} not found"))?;
            schema_positions_vec.push(pos as i64);
        }
    }
    let num_schema_pos = schema_positions_vec.len();
    let schema_positions = Tensor::from_array(
        (vec![num_schema_pos as i64], schema_positions_vec),
    )?;

    // span_idx: (1, num_words * max_width, 2)
    let mut spans = Vec::with_capacity(num_words * MAX_WIDTH * 2);
    for start in 0..num_words {
        for width in 1..=MAX_WIDTH {
            let end = start + width;
            if end <= num_words {
                spans.push(start as i64);
                spans.push((end - 1) as i64);
            } else {
                spans.push(0i64);
                spans.push(0i64);
            }
        }
    }
    let span_idx = Tensor::from_array(
        (vec![1i64, (num_words * MAX_WIDTH) as i64, 2i64], spans),
    )?;

    // Run inference
    let outputs = session.run(ort::inputs![
        "input_ids" => input_ids,
        "attention_mask" => attention_mask,
        "text_positions" => text_positions,
        "schema_positions" => schema_positions,
        "span_idx" => span_idx,
    ])?;

    // Extract span_scores: (1, num_fields, num_words, max_width)
    let (shape, scores_data) = outputs["span_scores"]
        .try_extract_tensor::<f32>()?;
    let num_fields = shape[1] as usize;

    // Decode entities
    let mut entities: Vec<Entity> = Vec::new();
    for fi in 0..num_fields {
        for start in 0..num_words {
            for wi in 0..MAX_WIDTH {
                let idx = fi * num_words * MAX_WIDTH + start * MAX_WIDTH + wi;
                let score = scores_data[idx];
                if score >= THRESHOLD {
                    let end = start + wi; // inclusive word index
                    if end >= words.len() {
                        continue;
                    }
                    let char_start = words[start].start;
                    let char_end = words[end].end;
                    entities.push(Entity {
                        label: LABELS[fi].to_string(),
                        text: TEXT[char_start..char_end].to_string(),
                        start: char_start,
                        end: char_end,
                        score,
                    });
                }
            }
        }
    }

    entities.sort_by(|a, b| a.start.cmp(&b.start).then(a.end.cmp(&b.end)));

    // Print results
    println!("Text: {TEXT}");
    println!("Labels: {LABELS:?}");
    println!("Threshold: {THRESHOLD}");
    println!();
    if entities.is_empty() {
        println!("  (no entities found)");
    } else {
        for ent in &entities {
            println!(
                "  {:>10}  {:.3}  [{:3}:{:3}]  {}",
                ent.label, ent.score, ent.start, ent.end, ent.text
            );
        }
    }

    Ok(())
}