import sys
import os
import pandas as pd
import spacy

BATCH_SIZE = 10

TARGET_ENTITY_LABELS = {
    "TIME", "DATE", "EVENT", "LOC", "ORG", "PRODUCT", "GPE"
}

EXCLUDED_VERB_LEMMAS = {"be", "have", "do"}
EXCLUDED_VERB_DEPS = {"aux", "auxpass"}


def load_resume_point(output_file):
    if not os.path.exists(output_file):
        return 0
    try:
        existing = pd.read_csv(output_file)
        if "row_id" in existing.columns and not existing.empty:
            return int(existing["row_id"].max())
    except Exception:
        pass
    return 0


def append_csv(rows, output_file):
    if not rows:
        return
    write_header = not os.path.exists(output_file)
    pd.DataFrame(rows).to_csv(
        output_file,
        mode="a",
        index=False,
        header=write_header
    )


def main():
    if len(sys.argv) < 3:
        print("Usage: python extract_contexts.py <csv_file> <text_column>")
        sys.exit(1)

    csv_file = sys.argv[1]
    text_column = sys.argv[2]

    base_name = os.path.splitext(os.path.basename(csv_file))[0]

    entity_file = f"{base_name}_entities.csv"
    verb_file = f"{base_name}_verbs.csv"
    verb_object_file = f"{base_name}_verb_objects.csv"

    try:
        df = pd.read_csv(csv_file)
    except Exception as e:
        print(f"Error reading CSV: {e}")
        sys.exit(1)

    if text_column not in df.columns:
        print(f"CSV must contain column '{text_column}'")
        sys.exit(1)

    print(f"📄 Processing column: '{text_column}'")
    print("🔍 Loading spaCy model...")
    nlp = spacy.load("en_core_web_md")

    processed_until = max(
        load_resume_point(entity_file),
        load_resume_point(verb_file),
        load_resume_point(verb_object_file),
    )

    if processed_until > 0:
        print(f"▶️ Resuming from row {processed_until + 1}")

    entity_buffer = []
    verb_buffer = []
    verb_object_buffer = []

    for i, text in enumerate(df[text_column].dropna(), start=1):
        if i <= processed_until:
            continue

        doc = nlp(str(text))

        # -------- ENTITIES --------
        for ent in doc.ents:
            if ent.label_ in TARGET_ENTITY_LABELS:
                entity_buffer.append({
                    "row_id": i,
                    "source_column": text_column,
                    "text": text,
                    "entity_text": ent.text,          # keep original case
                    "entity_label": ent.label_
                })

        # -------- VERBS + VERB–OBJECTS --------
        for token in doc:
            if (
                token.pos_ == "VERB"
                and token.dep_ not in EXCLUDED_VERB_DEPS
                and token.lemma_ not in EXCLUDED_VERB_LEMMAS
            ):
                verb_buffer.append({
                    "row_id": i,
                    "source_column": text_column,
                    "text": text,
                    "verb_text": token.text.lower(),
                    "verb_lemma": token.lemma_.lower(),
                    "verb_dep": token.dep_
                })

                for child in token.children:
                    if child.dep_ == "dobj":
                        verb_object_buffer.append({
                            "row_id": i,
                            "source_column": text_column,
                            "text": text,
                            "verb_lemma": token.lemma_.lower(),
                            "object_text": child.text.lower(),
                            "object_lemma": child.lemma_.lower(),
                            "object_dep": child.dep_
                        })

        # -------- BATCH WRITE --------
        if i % BATCH_SIZE == 0 or i == len(df):
            append_csv(entity_buffer, entity_file)
            append_csv(verb_buffer, verb_file)
            append_csv(verb_object_buffer, verb_object_file)

            print(
                f"💾 Saved up to row {i} | "
                f"entities={len(entity_buffer)}, "
                f"verbs={len(verb_buffer)}, "
                f"verb_objects={len(verb_object_buffer)}"
            )

            entity_buffer.clear()
            verb_buffer.clear()
            verb_object_buffer.clear()

    print("✅ Done.")
    print(f"📄 Entities: {entity_file}")
    print(f"📄 Verbs: {verb_file}")
    print(f"📄 Verb–Objects: {verb_object_file}")


if __name__ == "__main__":
    main()