use std::collections::BTreeMap;
use std::fs::{self, File};
use std::io::{self, BufRead, BufReader, BufWriter, Write};
use std::path::PathBuf;
use std::sync::LazyLock;

use anyhow::{Context, Result};
use clap::Parser;
use serde::Deserialize;
use serde_json::{json, Value};

#[derive(Parser)]
#[command(about = "Repair bracket/case misalignment and mislabeled release groups in the DMHY char-level dataset")]
struct Args {
    #[arg(long, default_value = "datasets\\AnimeName\\dmhy_weak_char.jsonl")]
    input: PathBuf,
    #[arg(long, default_value = "data\\dmhy_weak_char.repaired.jsonl")]
    output: PathBuf,
    #[arg(long, default_value = "data\\repair_report.json")]
    report: PathBuf,
    /// Compute and print statistics without writing the repaired dataset.
    #[arg(long)]
    dry_run: bool,
}

#[derive(Deserialize)]
struct Record {
    filename: String,
    #[serde(default)]
    tokens: Vec<String>,
    labels: Vec<String>,
    #[serde(default)]
    template: Option<String>,
    #[serde(default)]
    tokenizer_variant: Option<String>,
}

/// Release groups that appear un-bracketed as a filename prefix and were found
/// mislabeled as TITLE. Extend this list as more groups are confirmed.
static KNOWN_GROUPS: LazyLock<Vec<String>> = LazyLock::new(|| vec!["teneighty".to_string()]);

#[derive(Clone, Copy, PartialEq)]
enum AlignKind {
    Bracket,
    Case,
}

#[derive(Default)]
struct Stats {
    total: usize,
    already_ok: usize,
    normalized: usize,
    repaired_bracket: usize,
    repaired_case: usize,
    unrepairable: usize,
    group_relabeled: usize,
    examples: BTreeMap<String, Vec<Value>>,
}

fn main() -> Result<()> {
    let args = Args::parse();

    let input = File::open(&args.input)
        .with_context(|| format!("failed to open {}", args.input.display()))?;
    let reader = BufReader::new(input);

    if !args.dry_run {
        if let Some(parent) = args.output.parent() {
            if !parent.as_os_str().is_empty() {
                fs::create_dir_all(parent)
                    .with_context(|| format!("failed to create {}", parent.display()))?;
            }
        }
    }

    let mut out: Box<dyn Write> = if args.dry_run {
        Box::new(io::sink())
    } else {
        let f = File::create(&args.output)
            .with_context(|| format!("failed to create {}", args.output.display()))?;
        Box::new(BufWriter::new(f))
    };

    let mut stats = Stats::default();

    for (lineno, line) in reader.lines().enumerate() {
        let line = line.with_context(|| format!("failed to read line {}", lineno + 1))?;
        if line.trim().is_empty() {
            continue;
        }
        let rec: Record = serde_json::from_str(&line)
            .with_context(|| format!("failed to parse line {}", lineno + 1))?;
        stats.total += 1;

        let old_joined: String = rec.tokens.concat();
        let fn_chars: Vec<char> = rec.filename.chars().collect();
        let perfectly_aligned = rec.filename == old_joined
            && rec.tokens.len() == fn_chars.len()
            && rec.labels.len() == rec.tokens.len();

        if perfectly_aligned {
            let new_tokens = rec.tokens.clone();
            let mut new_labels = rec.labels.clone();
            if relabel_groups(&new_tokens, &mut new_labels) {
                stats.group_relabeled += 1;
                push_example(
                    &mut stats,
                    "group_relabeled",
                    &rec.filename,
                    &rec.tokens,
                    &rec.labels,
                    &new_tokens,
                    &new_labels,
                );
            }
            stats.already_ok += 1;
            if !args.dry_run {
                write_record(
                    &mut out,
                    &rec.filename,
                    &new_tokens,
                    &new_labels,
                    rec.template.as_deref(),
                    rec.tokenizer_variant.as_deref(),
                )?;
            }
            continue;
        }

        match align_to_filename(&rec.filename, &rec.tokens, &rec.labels) {
            Some((new_tokens, new_labels, kind)) => {
                let mut final_labels = new_labels;
                let grp = relabel_groups(&new_tokens, &mut final_labels);

                if rec.filename == old_joined && rec.tokens.len() != fn_chars.len() {
                    // Content identical but token granularity was not char-level.
                    stats.normalized += 1;
                    push_example(
                        &mut stats,
                        "normalized",
                        &rec.filename,
                        &rec.tokens,
                        &rec.labels,
                        &new_tokens,
                        &final_labels,
                    );
                } else if kind == AlignKind::Case {
                    stats.repaired_case += 1;
                    push_example(
                        &mut stats,
                        "repaired_case",
                        &rec.filename,
                        &rec.tokens,
                        &rec.labels,
                        &new_tokens,
                        &final_labels,
                    );
                } else {
                    stats.repaired_bracket += 1;
                    push_example(
                        &mut stats,
                        "repaired_bracket",
                        &rec.filename,
                        &rec.tokens,
                        &rec.labels,
                        &new_tokens,
                        &final_labels,
                    );
                }

                if grp {
                    stats.group_relabeled += 1;
                    push_example(
                        &mut stats,
                        "group_relabeled",
                        &rec.filename,
                        &rec.tokens,
                        &rec.labels,
                        &new_tokens,
                        &final_labels,
                    );
                }

                if !args.dry_run {
                    write_record(
                        &mut out,
                        &rec.filename,
                        &new_tokens,
                        &final_labels,
                        rec.template.as_deref(),
                        rec.tokenizer_variant.as_deref(),
                    )?;
                }
            }
            None => {
                stats.unrepairable += 1;
                push_example(
                    &mut stats,
                    "unrepairable",
                    &rec.filename,
                    &rec.tokens,
                    &rec.labels,
                    &rec.tokens,
                    &rec.labels,
                );
                if !args.dry_run {
                    write_record(
                        &mut out,
                        &rec.filename,
                        &rec.tokens,
                        &rec.labels,
                        rec.template.as_deref(),
                        rec.tokenizer_variant.as_deref(),
                    )?;
                }
            }
        }
    }
    out.flush()?;

    let summary = json!({
        "total": stats.total,
        "already_ok": stats.already_ok,
        "normalized": stats.normalized,
        "repaired_bracket": stats.repaired_bracket,
        "repaired_case": stats.repaired_case,
        "unrepairable": stats.unrepairable,
        "group_relabeled": stats.group_relabeled,
    });

    if args.dry_run {
        let report = json!({
            "input": args.input.display().to_string(),
            "dry_run": true,
            "summary": summary,
            "examples": stats.examples,
        });
        println!("{}", serde_json::to_string_pretty(&report)?);
    } else {
        let report = json!({
            "input": args.input.display().to_string(),
            "output": args.output.display().to_string(),
            "dry_run": false,
            "summary": summary,
            "examples": stats.examples,
        });
        let rf = File::create(&args.report)
            .with_context(|| format!("failed to create {}", args.report.display()))?;
        serde_json::to_writer_pretty(BufWriter::new(rf), &report)?;
        println!("{}", serde_json::to_string_pretty(&summary)?);
        println!("wrote repaired dataset -> {}", args.output.display());
        println!("wrote report -> {}", args.report.display());
    }

    Ok(())
}

/// Reconstruct character-level tokens/labels from `filename`, projecting the
/// old character-level labels onto the filename characters. Characters present
/// in `filename` but absent from the joined old tokens (e.g. dropped brackets)
/// become `O`; case-only differences are tolerated.
fn align_to_filename(
    filename: &str,
    old_tokens: &[String],
    old_labels: &[String],
) -> Option<(Vec<String>, Vec<String>, AlignKind)> {
    let mut old_chars: Vec<char> = Vec::new();
    for t in old_tokens {
        for c in t.chars() {
            old_chars.push(c);
        }
    }
    if old_chars.len() != old_labels.len() {
        return None;
    }
    let fn_chars: Vec<char> = filename.chars().collect();
    if fn_chars.len() < old_chars.len() {
        // filename is missing characters relative to the labels -> cannot repair
        return None;
    }

    let mut new_labels = Vec::with_capacity(fn_chars.len());
    let mut o_ptr = 0usize;
    let mut had_case = false;

    for &ch in &fn_chars {
        if o_ptr < old_chars.len() && old_chars[o_ptr] == ch {
            new_labels.push(old_labels[o_ptr].clone());
            o_ptr += 1;
        } else if o_ptr < old_chars.len()
            && old_chars[o_ptr].to_lowercase().eq(ch.to_lowercase())
            && old_chars[o_ptr] != ch
        {
            // case-only difference (e.g. filename `s02e01` vs labels `S02E01`)
            new_labels.push(old_labels[o_ptr].clone());
            o_ptr += 1;
            had_case = true;
        } else {
            // character inserted relative to old (typically a bracket)
            new_labels.push("O".to_string());
        }
    }

    if o_ptr != old_chars.len() {
        // some old characters were never matched -> real deletion -> unrepairable
        return None;
    }

    let new_tokens = fn_chars.iter().map(|c| c.to_string()).collect();
    let kind = if had_case {
        AlignKind::Case
    } else {
        AlignKind::Bracket
    };
    Some((new_tokens, new_labels, kind))
}

/// Rewrite a leading known release-group span (currently labeled TITLE_*) as
/// GROUP. Returns true if any labels were changed.
fn relabel_groups(tokens: &[String], labels: &mut [String]) -> bool {
    let n = tokens.len();
    if n == 0 {
        return false;
    }
    let start = match (0..n).find(|&i| labels[i] != "O") {
        Some(s) => s,
        None => return false,
    };
    let lab0 = labels[start].clone();
    if !lab0.starts_with("B-") {
        // Not at the start of a BIO span; leave it alone.
        return false;
    }
    let entity = lab0[2..].to_string();
    let continuation = format!("I-{entity}");
    let mut end = start;
    while end + 1 < n && labels[end + 1] == continuation {
        end += 1;
    }
    let span_text: String = tokens[start..=end].concat();
    if !KNOWN_GROUPS.iter().any(|g| g == &span_text) {
        return false;
    }
    labels[start] = "B-GROUP".to_string();
    for i in (start + 1)..=end {
        labels[i] = "I-GROUP".to_string();
    }
    true
}

fn write_record<W: Write>(
    w: &mut W,
    filename: &str,
    tokens: &[String],
    labels: &[String],
    template: Option<&str>,
    variant: Option<&str>,
) -> Result<()> {
    let obj = json!({
        "filename": filename,
        "tokens": tokens,
        "labels": labels,
        "template": template,
        "tokenizer_variant": variant,
    });
    serde_json::to_writer(&mut *w, &obj)?;
    w.write_all(b"\n")?;
    Ok(())
}

fn push_example(
    stats: &mut Stats,
    key: &str,
    filename: &str,
    before_tokens: &[String],
    before_labels: &[String],
    after_tokens: &[String],
    after_labels: &[String],
) {
    let arr = stats.examples.entry(key.to_string()).or_default();
    if arr.len() >= 25 {
        return;
    }
    arr.push(json!({
        "filename": filename,
        "before": spans_text(before_tokens, before_labels),
        "after": spans_text(after_tokens, after_labels),
    }));
}

fn spans_text(tokens: &[String], labels: &[String]) -> String {
    let mut out = String::new();
    let mut i = 0;
    while i < labels.len() {
        let lab = labels[i].clone();
        if lab == "O" {
            let mut j = i;
            while j + 1 < labels.len() && labels[j + 1] == "O" {
                j += 1;
            }
            let text: String = tokens[i..=j].concat();
            out.push_str(&format!(" {text}[O]"));
            i = j + 1;
            continue;
        }
        if !lab.starts_with("B-") {
            // lone I- without a preceding B-: render as its own span
            let mut j = i;
            while j + 1 < labels.len() && labels[j + 1] == lab {
                j += 1;
            }
            let text: String = tokens[i..=j].concat();
            out.push_str(&format!(" {text}[{lab}]"));
            i = j + 1;
            continue;
        }
        let entity = lab[2..].to_string();
        let continuation = format!("I-{entity}");
        let mut j = i;
        while j + 1 < labels.len() && labels[j + 1] == continuation {
            j += 1;
        }
        let text: String = tokens[i..=j].concat();
        out.push_str(&format!(" {text}[{lab}]"));
        i = j + 1;
    }
    out.trim_start().to_string()
}