package cn.edu.bjut.textmining.chapter5.support;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.UncheckedIOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.text.Normalizer;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Set;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/** 读取官方 JSONL；类别仅用于选择样本和事后评价。 */
public final class TNewsCorpusLoader {
    private static final Pattern FIELD =
            Pattern.compile("\\\"(label_desc|sentence)\\\"\\s*:\\s*\\\"((?:\\\\.|[^\\\"])*)\\\"");

    private TNewsCorpusLoader() {}

    public static final class NewsRecord {
        private final String label;
        private final String title;
        private final int sourceLine;

        public NewsRecord(String label, String title) {
            this(label, title, 0);
        }

        public NewsRecord(String label, String title, int sourceLine) {
            this.label = label;
            this.title = title;
            this.sourceLine = sourceLine;
        }

        public String getLabel() {
            return label;
        }

        public String getTitle() {
            return title;
        }

        public int getSourceLine() {
            return sourceLine;
        }
    }

    public static List<NewsRecord> readAll(Path file) {
        List<NewsRecord> records = new ArrayList<>();
        try (BufferedReader reader = Files.newBufferedReader(file, StandardCharsets.UTF_8)) {
            String line;
            int number = 0;
            while ((line = reader.readLine()) != null) {
                number++;
                Matcher matcher = FIELD.matcher(line);
                Map<String, String> values = new HashMap<>();
                while (matcher.find()) {
                    values.put(matcher.group(1), unescape(matcher.group(2)));
                }
                if (!values.containsKey("sentence") || !values.containsKey("label_desc")) {
                    throw new IllegalArgumentException("JSONL缺少字段，行号：" + number);
                }
                records.add(
                        new NewsRecord(values.get("label_desc"), values.get("sentence"), number));
            }
        } catch (IOException e) {
            throw new UncheckedIOException("读取新闻失败：" + file, e);
        }
        return records;
    }

    public static String normalizedKey(String title) {
        return Normalizer.normalize(title, Normalizer.Form.NFKC)
                .replaceAll("[\\s\\p{Z}]+", "")
                .toLowerCase(Locale.ROOT);
    }

    public static List<NewsRecord> readBalanced(Path file, List<String> categories, int limit) {
        if (limit < 1) {
            throw new IllegalArgumentException("每类样本数必须为正数");
        }
        Map<String, NewsRecord> unique = new LinkedHashMap<>();
        Set<String> conflicts = new HashSet<>();
        for (NewsRecord record : readAll(file)) {
            String key = normalizedKey(record.title);
            NewsRecord old = unique.get(key);
            if (old != null && !old.label.equals(record.label)) {
                conflicts.add(key);
            }
            if (old == null && !key.isEmpty()) {
                unique.put(key, record);
            }
        }
        Map<String, Integer> counts = new LinkedHashMap<>();
        for (String category : categories) {
            counts.put(category, 0);
        }
        List<NewsRecord> result = new ArrayList<>();
        for (Map.Entry<String, NewsRecord> item : unique.entrySet()) {
            NewsRecord record = item.getValue();
            Integer count = counts.get(record.label);
            if (!conflicts.contains(item.getKey()) && count != null && count < limit) {
                result.add(record);
                counts.put(record.label, count + 1);
            }
        }
        for (String category : categories) {
            if (counts.get(category) != limit) {
                throw new IllegalStateException("样本不足：" + category);
            }
        }
        return result;
    }

    public static List<NewsRecord> readCuratedSample(Path file, int perCategory) {
        if (perCategory < 1 || perCategory > 3) {
            throw new IllegalArgumentException("每类选择1至3条");
        }
        int[] lines = {208, 1241, 1431, 625, 676, 824, 274, 308, 566};
        List<NewsRecord> all = readAll(file);
        List<NewsRecord> result = new ArrayList<>();
        for (int group = 0; group < 3; group++) {
            for (int j = 0; j < perCategory; j++) {
                result.add(all.get(lines[group * 3 + j] - 1));
            }
        }
        return result;
    }

    public static List<String> titles(List<NewsRecord> records) {
        List<String> titles = new ArrayList<>();
        for (NewsRecord record : records) {
            titles.add(record.title);
        }
        return titles;
    }

    private static String unescape(String text) {
        StringBuilder result = new StringBuilder();
        for (int i = 0; i < text.length(); i++) {
            char c = text.charAt(i);
            if (c != '\\') {
                result.append(c);
                continue;
            }
            if (++i >= text.length()) {
                throw new IllegalArgumentException("JSON转义不完整");
            }
            c = text.charAt(i);
            switch (c) {
                case 'u':
                    result.append((char) Integer.parseInt(text.substring(i + 1, i + 5), 16));
                    i += 4;
                    break;
                case 'n':
                    result.append('\n');
                    break;
                case 'r':
                    result.append('\r');
                    break;
                case 't':
                    result.append('\t');
                    break;
                case 'b':
                    result.append('\b');
                    break;
                case 'f':
                    result.append('\f');
                    break;
                case '"':
                case '\\':
                case '/':
                    result.append(c);
                    break;
                default:
                    throw new IllegalArgumentException("未知JSON转义：" + c);
            }
        }
        return result.toString();
    }
}