from typing import List, Dict, Any

import pandas as pd
from nltk.metrics.distance import edit_distance
import spacy
from spacy.tokens import Doc
import constants
import json

# Register custom extension for tracking original index
if not Doc.has_extension("original_index"):
    Doc.set_extension("original_index", default=None)

def group_sections(df):
    df = classify_sections_by_title(df, threshold=0.8)
    df, suspected_splits = classify_revisions_by_content_modified(df, threshold=0.99)
    grouped_df = group_titles_using_content(df)
    section_splits_df = validate_section_splits(grouped_df, suspected_splits)
    return grouped_df, section_splits_df

def identify_section_splits(df, length_threshold: int = 100):
    if constants.GROUPED_IDX not in df.columns:
        raise ValueError("The input file must contain 'grouped_idx'")

    df[constants.CHANGED_CONTENT] = df[constants.CHANGED_CONTENT].fillna('').astype(str)

    # Group by 'grouped_idx'
    grouped = df.groupby('grouped_idx', sort=False)
    revision_id_counts = df['Revision ID'].value_counts()

    print("Identified potential section splits:")
    print("Revision ID | Group Index | Length Change")
    print("-" * 50)

    for group_index, group in grouped:
        previous_length = None

        for idx in group.index:  # Ensure iteration is in original order
            current_length = len(df.at[idx, constants.CHANGED_CONTENT])
            revision_id = df.at[idx, 'Revision ID']
            if previous_length is not None:
                length_difference = previous_length - current_length
                # Check if the length difference exceeds the threshold
                if length_difference > length_threshold and revision_id_counts[revision_id] > 1:
                    print(f"{revision_id:<12} | {group_index:<11} | {length_difference}")

            # Update previous length
            previous_length = current_length

def levenshtein_similarity(title_a, title_b):
    """
    Calculate the Levenshtein similarity between two strings.
    """
    title_a, title_b = str(title_a), str(title_b)
    distance = edit_distance(title_a, title_b)
    max_length = max(len(title_a), len(title_b))
    return 1 - (distance / max_length) if max_length > 0 else 0

def classify_sections_by_title(df, threshold=0.8):
    """
    Classify section titles based on Levenshtein similarity.
    """
    section_sets = []
    section_indices = []

    for section in df['Section']:
        max_avg_similarity = 0
        max_set_index = None

        for idx, section_set in enumerate(section_sets):
            similarities = [levenshtein_similarity(section, existing_section) for existing_section in section_set]
            avg_similarity = sum(similarities) / len(similarities)
            if avg_similarity > max_avg_similarity:
                max_avg_similarity = avg_similarity
                max_set_index = idx

        if max_avg_similarity > threshold and max_set_index is not None:
            section_sets[max_set_index].add(section)
            section_indices.append(max_set_index + 1)
        else:
            section_sets.append({section})
            section_indices.append(len(section_sets))

    df[constants.SECTION_IDX_TITLE] = section_indices
    return df

def classify_revisions_by_content_modified(df, threshold=0.99, max_sentence_diff=4):
    """
    Classify revisions based on content similarity using SciSpaCy,
    collect suspected section-splits (sim>threshold AND diff>=max_sentence_diff),
    but defer final split validation until after grouping is complete.
    """
    import spacy
    import pandas as pd

    # Load model
    nlp = spacy.load("en_core_sci_md")
    nlp.max_length = 10000

    # Prepare text
    df[constants.CHANGED_CONTENT] = (
        df[constants.CHANGED_CONTENT]
        .fillna("")
        .astype(str)
    )

    # Build spaCy docs
    content_docs = []
    for text in df[constants.CHANGED_CONTENT]:
        if len(text) > nlp.max_length:
            print(
                f"Warning: Content length {len(text)} exceeds "
                f"nlp.max_length. Truncating to {nlp.max_length} chars."
            )
            text = text[:nlp.max_length]
        content_docs.append(nlp(text))

    section_sets    = []  # lists of Docs
    section_indices = []  # final SECTION_IDX_CONTENT values
    suspected_splits = []  # suspected split-event records (not yet validated)

    # Iterate through each revision‐doc
    for i, doc in enumerate(content_docs):
        doc._.original_index = i
        cur_sents = len(list(doc.sents))

        # Track best small‐diff match
        best_sim, best_idx, best_diff = 0.0, None, None

        # Compare to each existing section's last doc
        for sec_idx, sec in enumerate(section_sets):
            last_doc   = sec[-1]
            last_sents = len(list(last_doc.sents))
            sim        = doc.similarity(last_doc)
            diff       = abs(cur_sents - last_sents)

            # Record a SUSPECTED split event if sim high AND size-change large
            if sim > threshold and diff >= max_sentence_diff:
                suspected_splits.append({
                    "current_index":  i,
                    "previous_index": last_doc._.original_index,
                    "similarity":     sim,
                    "sentence_diff":  diff,
                })

            # Only consider for grouping if size-change is small
            if sim > best_sim and diff < max_sentence_diff:
                best_sim, best_idx, best_diff = sim, sec_idx, diff

        # Group-then-new logic (same as before)
        if best_idx is not None and best_sim > threshold:
            # high similarity + small diff ⇒ same section
            section_sets[best_idx].append(doc)
            section_indices.append(best_idx + 1)
        else:
            # otherwise ⇒ new section
            section_sets.append([doc])
            section_indices.append(len(section_sets))

    # Write section indices back to DataFrame
    df[constants.SECTION_IDX_CONTENT] = section_indices

    return df, suspected_splits

def validate_section_splits(grouped_df, suspected_splits):
    """
    Validate suspected section splits by checking if the items have different grouped indices.
    Prevent duplicate splits by ensuring:
    1. Each previous_idx is only used once (allows cascading X→Y→Z)
    2. Each current_idx (revision) only appears once in the final splits (prevents A→X, B→X duplicates)
    
    Args:
        grouped_df: DataFrame with grouped_idx column
        suspected_splits: List of suspected split records from classify_revisions_by_content_modified
    
    Returns:
        section_splits_df: DataFrame containing validated section splits
    """
    import pandas as pd
    
    if not suspected_splits:
        return pd.DataFrame()
    
    validated_splits = []
    used_previous_indices = set()  # Track previous_idx that have been used in confirmed splits
    used_current_indices = set()   # Track current_idx that have been used in confirmed splits
    
    for split in suspected_splits:
        current_idx = split["current_index"]
        previous_idx = split["previous_index"]
        
        # Skip if this previous_idx has already been used in a confirmed split
        if previous_idx in used_previous_indices:
            continue
            
        # Skip if this current_idx has already been used in a confirmed split
        # This prevents the same revision from appearing multiple times in splits
        if current_idx in used_current_indices:
            continue
            
        # Get grouped indices for both items
        current_grouped_idx = grouped_df.iloc[current_idx][constants.GROUPED_IDX]
        previous_grouped_idx = grouped_df.iloc[previous_idx][constants.GROUPED_IDX]
        
        # Validate split: different grouped indices = true section split
        if current_grouped_idx != previous_grouped_idx:
            validated_splits.append(split)
            # Mark both indices as used
            used_previous_indices.add(previous_idx)
            used_current_indices.add(current_idx)
    
    # Build enriched splits DataFrame (same format as original)
    if validated_splits:
        raw = pd.DataFrame(validated_splits)
        out = grouped_df.iloc[raw["current_index"]].copy()
        out["Previous Section"] = grouped_df.iloc[raw["previous_index"], grouped_df.columns.get_loc("Section")].values
        out["Previous Revision ID"] = grouped_df.iloc[raw["previous_index"], grouped_df.columns.get_loc("Revision ID")].values
        out["Previous Changed Content"] = grouped_df.iloc[raw["previous_index"], grouped_df.columns.get_loc(constants.CHANGED_CONTENT)].values
        out["Similarity"] = raw["similarity"].values
        out["Sentence Diff"] = raw["sentence_diff"].values
        
        # Add grouped index information for debugging/analysis
        out["Current Grouped Index"] = grouped_df.iloc[raw["current_index"], grouped_df.columns.get_loc(constants.GROUPED_IDX)].values
        out["Previous Grouped Index"] = grouped_df.iloc[raw["previous_index"], grouped_df.columns.get_loc(constants.GROUPED_IDX)].values
        
        section_splits_df = out
    else:
        section_splits_df = pd.DataFrame()
    
    return section_splits_df

def classify_revisions_by_content(df, threshold=0.99, max_sentence_diff=4):
    """
    Original function - kept for backward compatibility.
    Classify revisions based on content similarity using SciSpaCy,
    record any section-splits (sim>threshold AND diff>=max_sentence_diff),
    but always prefer grouping into an existing section when diff<max_sentence_diff.
    """
    import spacy
    import pandas as pd

    # Load model
    nlp = spacy.load("en_core_sci_md")
    nlp.max_length = 10000

    # Prepare text
    df[constants.CHANGED_CONTENT] = (
        df[constants.CHANGED_CONTENT]
        .fillna("")
        .astype(str)
    )

    # Build spaCy docs
    content_docs = []
    for text in df[constants.CHANGED_CONTENT]:
        if len(text) > nlp.max_length:
            print(
                f"Warning: Content length {len(text)} exceeds "
                f"nlp.max_length. Truncating to {nlp.max_length} chars."
            )
            text = text[:nlp.max_length]
        content_docs.append(nlp(text))

    section_sets    = []  # lists of Docs
    section_indices = []  # final SECTION_IDX_CONTENT values
    splits          = []  # raw split-event records

    # Iterate through each revision‐doc
    for i, doc in enumerate(content_docs):
        doc._.original_index = i
        cur_sents = len(list(doc.sents))

        # Track best small‐diff match
        best_sim, best_idx, best_diff = 0.0, None, None

        # Compare to each existing section's last doc
        for sec_idx, sec in enumerate(section_sets):
            last_doc   = sec[-1]
            last_sents = len(list(last_doc.sents))
            sim        = doc.similarity(last_doc)
            diff       = abs(cur_sents - last_sents)

            # Record a split event if sim high AND size-change large
            if sim > threshold and diff >= max_sentence_diff:
                splits.append({
                    "current_index":  i,
                    "previous_index": last_doc._.original_index,
                    "similarity":     sim,
                    "sentence_diff":  diff,
                })

            # Only consider for grouping if size-change is small
            if sim > best_sim and diff < max_sentence_diff:
                best_sim, best_idx, best_diff = sim, sec_idx, diff

        # Group-then-new logic
        if best_idx is not None and best_sim > threshold:
            # high similarity + small diff ⇒ same section
            section_sets[best_idx].append(doc)
            section_indices.append(best_idx + 1)
        else:
            # otherwise ⇒ new section
            section_sets.append([doc])
            section_indices.append(len(section_sets))

    # Write section indices back to DataFrame
    df[constants.SECTION_IDX_CONTENT] = section_indices

    # Build enriched splits DataFrame
    if splits:
        raw = pd.DataFrame(splits)
        out = df.iloc[raw["current_index"]].copy()
        out["Previous Section"]        = df.iloc[raw["previous_index"],    df.columns.get_loc("Section")].values
        out["Previous Revision ID"]    = df.iloc[raw["previous_index"],    df.columns.get_loc("Revision ID")].values
        out["Previous Changed Content"] = df.iloc[raw["previous_index"],   df.columns.get_loc(constants.CHANGED_CONTENT)].values
        out["Similarity"]              = raw["similarity"].values
        out["Sentence Diff"]           = raw["sentence_diff"].values
        section_splits_df = out
    else:
        section_splits_df = pd.DataFrame()

    return df, section_splits_df

def group_titles_using_content(df):
    filtered_df = df[df['Section'] != "(Top)"]
    idx_dict = create_map_of_content_idx_to_relative_title_indexes(filtered_df)
    grouped_indices = unite_indices(idx_dict)
    title_to_group = {
        title_index: group_index
        for group_index, group_set in enumerate(grouped_indices, start=1)  # Groups are 1-indexed
        for title_index in group_set
    }

    # Step 5: Add the 'grouped_idx' column to the DataFrame
    df[constants.GROUPED_IDX] = df[constants.SECTION_IDX_TITLE].map(title_to_group).fillna(0).astype(int)
    return df

def create_map_of_content_idx_to_relative_title_indexes(filtered_df):
    """
    Create a mapping from content index to title indexes.
    :param filename_combined: str
        Path to the input CSV file.
    :return: dict
        Dictionary where content index maps to a list of title indexes.
    """
    return (
        filtered_df.groupby(constants.SECTION_IDX_CONTENT)[constants.SECTION_IDX_TITLE]
        .apply(list)
        .to_dict()
    )

def unite_indices(data):
    """
    Unites overlapping indices from dictionary values into distinct sets.
    :param data: dict
        A dictionary where keys map to lists of indices.
    :return: list
        A list of sets with unified indices.
    """
    # Convert the dictionary values into sets
    sets = [set(value) for value in data.values()]

    # Iteratively merge sets with overlaps
    merged = []
    while sets:
        current = sets.pop(0)
        overlap_found = False

        for other_set in merged:
            if current & other_set:  # Check for overlap
                other_set.update(current)  # Merge overlapping sets
                overlap_found = True
                break

        if not overlap_found:
            merged.append(current)  # Add as a new set

    return merged

def check_section_name_changes(filename: str, threshold: float = 0.8) -> str:
    """
    Parse the output CSV by section index and check for significant section name changes.
    Log events where Levenshtein similarity is below the threshold.
    """
    df = pd.read_csv(filename)
    if 'Section from similarity' not in df.columns or 'Section' not in df.columns:
        raise ValueError("The input file must contain 'Section' and 'Section from similarity' columns.")

    df['Levenshtein Similarity from Previous'] = None
    grouped = df.groupby('Section from similarity', sort=False)

    for section_index, group in grouped:
        previous_name = None

        for idx in group.index:  # Use original indices to ensure order
            current_name = df.at[idx, 'Section']

            if previous_name is not None:
                distance = levenshtein_similarity(previous_name, current_name)
                df.at[idx, 'Levenshtein Distance from Previous'] = distance
            previous_name = current_name

    output_filename = filename.replace(".csv", "_with_levenshtein_distance.csv")
    df.to_csv(output_filename, index=False)
    print(f"Levenshtein distance calculation completed. Results saved to {output_filename}.")
    return output_filename

def generate_diff_json(df: pd.DataFrame, article_name: str, output_path: str = "diff_urls.json") -> List[Dict[str, Any]]:
    """
    Generate a JSON file listing Wikipedia diff URLs and associated section title changes.

    Each entry includes:
        - 'url': diff link
        - 'section_before': section title in previous revision
        - 'section_after': section title in current revision

    Grouping is done on both GROUPED_IDX and SECTION_IDX_CONTENT.

    Args:
        df (pd.DataFrame): DataFrame with revision data.
        article_name (str): Wikipedia article title.
        output_path (str): Path to save the output JSON file (default: diff_urls.json)
    """
    required_cols = [
        constants.GROUPED_IDX,
        constants.SECTION_IDX_CONTENT,
        constants.SECTION_IDX_TITLE,
        'Revision ID',
        'Section'
    ]
    if not all(col in df.columns for col in required_cols):
        raise ValueError(f"Missing one or more required columns: {required_cols}")

    title_slug = article_name.replace(" ", "_")
    diffs = []

    # Sort and group by both GROUPED_IDX and SECTION_IDX_CONTENT
    df_sorted = df.sort_values(by=[constants.GROUPED_IDX, constants.SECTION_IDX_CONTENT, 'Revision ID'])
    grouped = df_sorted.groupby([constants.GROUPED_IDX, constants.SECTION_IDX_CONTENT], sort=False)

    for (_, _), group_df in grouped:
        group_df = group_df.reset_index(drop=True)

        for i in range(1, len(group_df)):
            current_title_idx = group_df.loc[i, constants.SECTION_IDX_TITLE]
            prev_title_idx = group_df.loc[i - 1, constants.SECTION_IDX_TITLE]

            if current_title_idx != prev_title_idx:
                current_rev_id = group_df.loc[i, 'Revision ID']
                prev_rev_id = group_df.loc[i - 1, 'Revision ID']

                url = (
                    f"https://en.wikipedia.org/w/index.php"
                    f"?title={title_slug}"
                    f"&diff={current_rev_id}"
                    f"&oldid={prev_rev_id}"
                )

                diffs.append({
                    "url": url,
                    "section_before": group_df.loc[i - 1, 'Section'],
                    "section_after": group_df.loc[i, 'Section']
                })

    # Save JSON
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(diffs, f, indent=2, ensure_ascii=False)

    print(f"Saved {len(diffs)} diff entries to {output_path}")
    return diffs