import os
import sys
import argparse
import json
from collections import defaultdict, deque
import pandas as pd
from openpyxl.styles import Font

# -------- optional command-line arguments --------

parser = argparse.ArgumentParser()
parser.add_argument("--proband", default=None, help="Full name of the proband (web use)")
parser.add_argument("--proband-id", default=None, help="GEDCOM ID of the proband, for example I0123 (web use)")
parser.add_argument("--outdir", default=None, help="Output directory (web use)")
parser.add_argument("--gedfile", default=None, help="Path to the GEDCOM file (web use)")
args, _ = parser.parse_known_args()

RED = "\033[31m"
GREEN = "\033[32m"
YELLOW = "\033[33m"
BLUE = "\033[34m"
RESET = "\033[0m"

SCRIPT_VERSION = "2026-08-04 output-folder check"

# -------- locate the script directory --------

script_dir = os.path.dirname(os.path.abspath(__file__))
os.chdir(script_dir)

print("Relationship script version:", SCRIPT_VERSION)
print("Working directory:", script_dir)


# -------- locate the relationship file --------

english_relations_file = "relationship_names.xlsx"
danish_relations_file = "Slægtsnavne.xlsx"

english_relations_exists = os.path.exists(english_relations_file)
danish_relations_exists = os.path.exists(danish_relations_file)

if english_relations_exists:
    relations_file = english_relations_file

    if danish_relations_exists:
        print(
            f"\n{YELLOW}Warning: Both the English and Danish relationship tables "
            f"were found. The English relationship table will be used.{RESET}\n"
        )
elif danish_relations_exists:
    relations_file = danish_relations_file
else:
    print(
        f"\n{RED}No relationship table was found. "
        f"Add either the English or Danish relationship table to the script folder.\n{RESET}"
    )
    sys.exit(1)


relations_df = pd.read_excel(relations_file)

relations_df.columns = relations_df.columns.str.strip().str.lower()

# Load the selected relationship table
# Map (generations up, generations down) to relationship text

relation_table = {}
affinal_table = {}

for _, row in relations_df.iterrows():

    up = row.get("up")
    down = row.get("down")
    path = row.get("path")

    entry = {
        "default": row.get("relation"),
        "m": row.get("relation_m"),
        "f": row.get("relation_f")
    }

    # blood relationships

    if pd.notna(up) and pd.notna(down):

        relation_table[(int(up), int(down))] = entry


    # relationships by marriage

    if pd.notna(path):

        affinal_table[path.strip().lower()] = entry


# -------- locate the GEDCOM file --------

parent_dir = os.path.dirname(script_dir)

if args.gedfile:

    ged_file = args.gedfile

else:

    ged_file = None

    # Search the script directory first
    for f in os.listdir(script_dir):

        if f.lower().endswith((".ged", ".ged2")):

            ged_file = os.path.join(script_dir, f)

            break

    # Fall back to the parent directory if no file was found
    if not ged_file:

        for f in os.listdir(parent_dir):

            if f.lower().endswith((".ged", ".ged2")):

                ged_file = os.path.join(parent_dir, f)

                break

if not ged_file:

    print(f"{RED}No GEDCOM or GED2 file was found.{RESET}")

    sys.exit(1)

# -------- read the requested proband --------

argument_mode = bool(args.proband or args.proband_id)
proband_name = args.proband.strip() if args.proband else None
proband_id_argument = args.proband_id.strip() if args.proband_id else None

if not argument_mode:

    proband_name = input("Enter the proband's full name: ").strip()

    if not proband_name:
        print(f"{RED}No proband name was entered.{RESET}")
        sys.exit(1)

# -------- parse the GEDCOM file --------

individuals = {}
families = {}

with open(ged_file,"r",encoding="utf-8",errors="ignore") as f:

    cid=None
    ctype=None
    in_occu = False

    for line in f:

        parts=line.strip().split(" ",2)

        if len(parts)<2:
            continue

        level=parts[0]

        if len(parts)==2:
            tag,data=parts[1],""
        else:
            if parts[1].startswith("@"):
                data,tag=parts[1],parts[2]
            else:
                tag,data=parts[1],parts[2]

        if level=="0":

            if data.startswith("@") and tag in ("INDI","FAM"):

                cid=data
                ctype=tag

                if tag=="INDI":
                    individuals[cid]={"name":None,"sex":None,"occupation":""}

                else:
                    families[cid]={"husb":None,"wife":None,"chil":[],"pedi":{}}

            else:
                cid=None
                ctype=None
                in_occu = False

            continue

        if ctype=="INDI":

            if tag=="NAME" and individuals[cid]["name"] is None:
                individuals[cid]["name"]=data.replace("/","").strip()

            elif tag=="SEX":
                individuals[cid]["sex"]=data

            elif tag=="FAMC":
                individuals[cid].setdefault("famc", []).append({
                "fam": data,
                "pedi": "birth"
                })

            elif tag=="PEDI":
                if "famc" in individuals[cid]:
                    individuals[cid]["famc"][-1]["pedi"] = data

            elif tag == "OCCU":
                in_occu = True

                if not isinstance(individuals[cid].get("occupation"), list):
                    individuals[cid]["occupation"] = []

                # Flush the previous pending value: without PLAC, OCCU contains the occupation
                prev = individuals[cid].get("_occu_pending")
                if prev:
                    individuals[cid]["occupation"].append(prev)

                # Store the new OCCU value temporarily and ignore it if PLAC follows
                # (Gramps stores the description in OCCU and the occupation itself in PLAC)
                individuals[cid]["_occu_pending"] = data if data else None

            elif in_occu and tag == "PLAC":

                if not isinstance(individuals[cid].get("occupation"), list):
                    individuals[cid]["occupation"] = []

                if data:
                    individuals[cid]["occupation"].append(data)

                # The OCCU value was a description, so ignore it
                individuals[cid]["_occu_pending"] = None
                in_occu = False

        elif ctype=="FAM":

            if tag=="HUSB":
                families[cid]["husb"]=data

            elif tag=="WIFE":
                families[cid]["wife"]=data

            elif tag=="CHIL":
                families[cid]["chil"].append(data)
                
            elif tag=="PEDI":
                if families[cid]["chil"]:
                    last_child = families[cid]["chil"][-1]
                    families[cid]["pedi"][last_child] = data

# -------- flush any remaining OCCU value (without PLAC, it is the occupation) --------

for iid in individuals:
    pending = individuals[iid].pop("_occu_pending", None)
    if pending:
        if not isinstance(individuals[iid].get("occupation"), list):
            individuals[iid]["occupation"] = []
        individuals[iid]["occupation"].append(pending)

# -------- build the relationship graph --------
edges = defaultdict(list)

def add(a, b, label):
    edges[a].append((b, label))

for fam_id, fam in families.items():
    h = fam["husb"]
    w = fam["wife"]

    if h and w:
        add(h, w, "spouse")
        add(w, h, "spouse")

    for c in fam["chil"]:
        # Check whether the child belongs to this family as a foster child
        is_foster = False
        if c in individuals and "famc" in individuals[c]:
            for f_link in individuals[c]["famc"]:
                if f_link["fam"] == fam_id and f_link["pedi"] == "foster":
                    is_foster = True
                    break
        
        # Skip foster-child relationships
        if is_foster:
            continue

        if h:
            add(h, c, "child")
            add(c, h, "parent")
        if w:
            add(w, c, "child")
            add(c, w, "parent")

# -------- find the proband in the GEDCOM data --------

def normalize_gedcom_id(person_id):
    """Return a GEDCOM ID in the internal @I123@ format."""
    clean_id = person_id.strip().strip("@").upper()
    return f"@{clean_id}@"


def display_gedcom_id(person_id):
    """Return a GEDCOM ID without the surrounding @ characters."""
    return person_id.strip("@")


proband = None

if proband_id_argument:

    requested_id = normalize_gedcom_id(proband_id_argument)

    # GEDCOM IDs are treated case-insensitively for command-line convenience.
    proband = next(
        (person_id for person_id in individuals if person_id.upper() == requested_id),
        None,
    )

    if not proband:
        print(
            f"{RED}No person with GEDCOM ID "
            f"{display_gedcom_id(requested_id)} was found.{RESET}"
        )
        sys.exit(1)

    selected_name = individuals[proband].get("name") or "Unknown name"

    if proband_name and proband_name.casefold() != selected_name.casefold():
        print(
            f"{RED}The supplied name and GEDCOM ID do not identify the same person.{RESET}"
        )
        print(f"Name supplied: {proband_name}")
        print(
            f"ID supplied: {display_gedcom_id(proband)} "
            f"({selected_name})"
        )
        sys.exit(1)

    proband_name = selected_name
    print(
        f"{GREEN}\nProband received as an argument: "
        f"{proband_name} — ID {display_gedcom_id(proband)}{RESET}"
    )
    print()

else:

    matches = [
        person_id
        for person_id, person_data in individuals.items()
        if person_data.get("name")
        and proband_name.casefold() == person_data["name"].casefold()
    ]

    if not matches:
        print(f"{RED}Proband not found.{RESET}")
        if not argument_mode:
            input("Press Enter to close.")
        sys.exit(1)

    if len(matches) == 1:
        proband = matches[0]

    elif argument_mode:
        print(
            f"{RED}More than one person was found with the name "
            f"{proband_name}.{RESET}"
        )
        print("\nAvailable people:")

        for person_id in matches:
            print(
                f"  {individuals[person_id]['name']} — "
                f"ID {display_gedcom_id(person_id)}"
            )

        print("\nRun the script again with --proband-id followed by the desired ID.")
        sys.exit(1)

    else:
        print(
            f"{YELLOW}\nMore than one person was found with the name "
            f"{proband_name}.{RESET}"
        )
        print()

        for number, person_id in enumerate(matches, start=1):
            print(
                f"{number}. {individuals[person_id]['name']} — "
                f"ID {display_gedcom_id(person_id)}"
            )

        while True:
            choice = input(f"\nSelect the person [1-{len(matches)}]: ").strip()

            if choice.isdigit() and 1 <= int(choice) <= len(matches):
                proband = matches[int(choice) - 1]
                break

            print(
                f"{RED}Enter a number from 1 to {len(matches)}.{RESET}"
            )

        print(
            f"{GREEN}\nSelected proband: {proband_name} — "
            f"ID {display_gedcom_id(proband)}{RESET}"
        )
        print()

# -------- BFS --------

prev={proband:(None,None)}

q=deque([proband])

while q:

    u=q.popleft()

    for v,l in edges[u]:

        if v not in prev:

            prev[v]=(u,l)

            q.append(v)

# -------- reconstruct the relationship path --------

def reconstruct(i):

    nodes=[]
    labels=[]

    n=i

    while n:

        p,l=prev.get(n,(None,None))

        nodes.append(n)

        if l:
            labels.append(l)

        n=p

    nodes.reverse()
    labels.reverse()

    return nodes,labels

def sex(i):
    return individuals.get(i,{}).get("sex","")

# -------- build the readable relationship path --------

def chain(nodes,labels):

    if not labels:
        return "self"

    words=[]

    for i,l in enumerate(labels):

        sid=nodes[i+1]
        s=sex(sid)

        if l=="parent":
            if s=="M":
                w="father"
            elif s=="F":
                w="mother"
            else:
                w="parent (unknown)"

        elif l=="child":
            if s=="M":
                w="son"
            elif s=="F":
                w="daughter"
            else:
                w="child (unknown)"

        elif l=="spouse":
            w="spouse"

        else:
            w="relative"

        if i<len(labels)-1:
            words.append(w+"'s")
        else:
            words.append(w)

    return " ".join(words)


# -------- build a simplified relationship path --------

def _possessive(phrase):
    """Return an English possessive phrase."""
    if phrase.endswith("s"):
        return phrase + "'"
    return phrase + "'s"


def _join_relationship_parts(parts):
    """Join relationship parts as an English possessive chain."""
    return " ".join(
        _possessive(part) if index < len(parts) - 1 else part
        for index, part in enumerate(parts)
    )


def _spouse_word(person_id):
    """Return husband, wife, or spouse from the person's recorded sex."""
    person_sex = sex(person_id)
    if person_sex == "M":
        return "husband"
    if person_sex == "F":
        return "wife"
    return "spouse"


def simplified_path(nodes, labels):
    """
    Build a shorter, everyday English relationship path.

    Examples:
        father's daughter's daughter -> sister's daughter
        father's daughter's spouse -> sister's husband
        spouse's father -> wife's father or husband's father

    Unknown sex is handled with neutral words such as sibling, child,
    parent, and spouse.
    """
    full_path = chain(nodes, labels)

    if not labels:
        return full_path

    # Handle relationships by marriage recursively. This allows one rule to
    # cover paths such as wife's father, sister's husband, and more distant
    # combinations containing several spouse links.
    if "spouse" in labels:
        spouse_index = labels.index("spouse")
        parts = []

        # Blood relationship from the proband to the person immediately before
        # the spouse link.
        if spouse_index > 0:
            left_nodes = nodes[: spouse_index + 1]
            left_labels = labels[:spouse_index]
            parts.append(simplified_path(left_nodes, left_labels))

        spouse_person = nodes[spouse_index + 1]
        parts.append(_spouse_word(spouse_person))

        # Continue from that spouse to the final person. The suffix can itself
        # contain another spouse link and is therefore handled recursively.
        if spouse_index + 1 < len(labels):
            right_nodes = nodes[spouse_index + 1 :]
            right_labels = labels[spouse_index + 1 :]
            parts.append(simplified_path(right_nodes, right_labels))

        return _join_relationship_parts(parts)

    first_down = next(
        (index for index, label in enumerate(labels) if label == "child"),
        None,
    )

    # Direct ancestors or direct descendants do not need shortening.
    if first_down is None or first_down == 0:
        return full_path

    generations_up = first_down
    branch_person = nodes[first_down + 1]
    branch_sex = sex(branch_person)

    if generations_up == 1:
        if branch_sex == "M":
            first_word = "brother"
        elif branch_sex == "F":
            first_word = "sister"
        else:
            first_word = "sibling"

    elif generations_up == 2:
        if branch_sex == "M":
            first_word = "uncle"
        elif branch_sex == "F":
            first_word = "aunt"
        else:
            first_word = "uncle or aunt"

    elif generations_up == 3:
        if branch_sex == "M":
            first_word = "granduncle"
        elif branch_sex == "F":
            first_word = "grandaunt"
        else:
            first_word = "granduncle or grandaunt"

    else:
        prefix = "great-" * (generations_up - 3)
        if branch_sex == "M":
            first_word = f"{prefix}granduncle"
        elif branch_sex == "F":
            first_word = f"{prefix}grandaunt"
        else:
            first_word = f"{prefix}granduncle or {prefix}grandaunt"

    words = [first_word]

    # Add the descendants from the branch person to the focus person.
    for index in range(first_down + 1, len(labels)):
        node_id = nodes[index + 1]
        node_sex = sex(node_id)

        if node_sex == "M":
            word = "son"
        elif node_sex == "F":
            word = "daughter"
        else:
            word = "child"

        words.append(word)

    return _join_relationship_parts(words)

# -------- determine the relationship --------

def detect_relation(labels, sex):

    # ---------- check relationships by marriage first ----------

    if "spouse" in labels:

        path_key = "-".join(labels)

        entry = affinal_table.get(path_key)

        if entry:
        
            relation_m = entry["m"]
            relation_f = entry["f"]


            if sex == "M" and relation_m == relation_m:
                return relation_m

            if sex == "F" and relation_f == relation_f:
                return relation_f


            return entry["default"]
 
        return ""
    
    # ---------- count generations up and down ----------

    up = 0
    down = 0

    direction = "up"

    for l in labels:

        if direction == "up":

            if l == "parent":
                up += 1
                continue

            direction = "down"

        if direction == "down":

            if l == "child":
                down += 1

    # ---------- look up the blood relationship in Excel ----------

    entry = relation_table.get((up, down))

    if entry:
        
        relation_m = entry["m"]
        relation_f = entry["f"]


        if sex == "M" and relation_m == relation_m:
            return relation_m

        if sex == "F" and relation_f == relation_f:
            return relation_f


        return entry["default"]
 
    return ""
    
# -------- calculate the generation --------

def generation(labels):

    g=0

    for l in labels:

        if l=="parent":
            g-=1

        elif l=="child":
            g+=1

    return g

# -------- build the internal path data --------

def build_sti(nodes, labels):
    """
    Build an explicit two-row ancestry path for use by the person-page script.

    Row 1 follows the proband upward to the turning ancestor.
    Row 2 follows the focus person's line downward from that ancestor.

    The path is created only for blood relatives. Spouses and relatives by
    marriage receive an empty Path, although a turning ancestor's spouse may
    still be included in another person's path with Spouse=True.
    """
    # No ancestry path for relationships by marriage
    if "spouse" in labels:
        return []

    if not labels:
        return []   # the proband

    # Find the first child step: the turning point
    first_down = next(
        (i for i, l in enumerate(labels) if l == "child"), None
    )

    sti = []

    # The proband is always on row 1
    sti.append({"ID": nodes[0], "Row": 1})

    # Use row 2 only when the path first goes up and then down
    # Direct descendants and direct ancestors remain on row 1
    # because their paths do not branch
    use_two_rows = first_down is not None and first_down > 0

    # Intermediate people and the final person
    for idx, label in enumerate(labels):
        node_id = nodes[idx + 1]
        linje = 1 if (not use_two_rows or idx < first_down) else 2
        sti.append({"ID": node_id, "Row": linje})

    # Spouse at the turning point
    # Only when at least one parent step comes first
    if first_down is not None and first_down > 0:
        turning_node = nodes[first_down]       # the highest shared ancestor
        child_node   = nodes[first_down + 1]   # the first child on the downward path

        # Find the child's parents and select the one who is not the turning point
        child_data = individuals.get(child_node, {})
        for famc_entry in child_data.get("famc", []):
            fam_id = famc_entry.get("fam") if isinstance(famc_entry, dict) else famc_entry
            fam = families.get(fam_id, {})
            husb = fam.get("husb")
            wife = fam.get("wife")
            other = None
            if husb == turning_node and wife:
                other = wife
            elif wife == turning_node and husb:
                other = husb
            if other:
                sti.append({"ID": other, "Row": 2, "Spouse": True})
                break

    return sti

# -------- build the report rows --------

rows=[]

for i,d in individuals.items():

    name=d["name"]

    if i not in prev:
        continue

    nodes,labels=reconstruct(i)

    ch=chain(nodes,labels)
    simple_ch=simplified_path(nodes,labels)
    
    gen=generation(labels)


    relation_text = detect_relation(labels, sex(i))

 
    rows.append({

        "ID": i,
        "Name": name,
        "Relationship": relation_text,
        "Simplified path": simple_ch,
        "Relationship path": ch,
        "Generation": gen,
        "Path": build_sti(nodes, labels)

    })

df=pd.DataFrame(rows)

df["Generation"] = df["Generation"].fillna(0).astype(int)

# place the proband first
df_proband = df[df["ID"] == proband]
df_rest = df[df["ID"] != proband]

df = pd.concat([df_proband, df_rest])

# remove ID and Path again (used only in relationships.json)
df = df.drop(columns=["ID", "Path"])

# place the proband first
df_proband = df[df["Generation"] == 0]
df_rest = df[df["Generation"] != 0]
df = pd.concat([df_proband, df_rest])

safe_name = proband_name.replace(" ", "_")

if args.outdir:
    # Web mode: use the output directory supplied by the website.
    out_dir = os.path.abspath(args.outdir)
else:
    # Local mode: identify the selected proband in the directory name.
    proband_directory = f"{safe_name}_{display_gedcom_id(proband)}"
    out_dir = os.path.join(script_dir, proband_directory)

print("Output directory:", out_dir)
os.makedirs(out_dir, exist_ok=True)

out = os.path.join(out_dir, f"{safe_name}.xlsx")

with pd.ExcelWriter(out, engine="openpyxl") as writer:

    # write the title at the top
    title = f"Relationship report for {proband_name}"

    df.to_excel(
        writer,
        index=False,
        sheet_name="Relationships",
        startrow=2
    )

    ws = writer.sheets["Relationships"]

    # title
    ws["A1"] = title

    # make the title larger and bold
    ws["A1"].font = Font(size=14, bold=True)
    # make the column headings bold
    for cell in ws[3]:
        cell.font = Font(bold=True)

    # enable filtering
    ws.auto_filter.ref = f"A3:E{ws.max_row}"

    # freeze the top rows
    ws.freeze_panes = "A3"
    
    # adjust column widths
    for col in ws.columns:
        max_length = 0
        column = col[0].column_letter

        for cell in col:
            if cell.value:
                max_length = max(max_length, len(str(cell.value)))

        ws.column_dimensions[column].width = max_length + 2

print(f"{GREEN}Excel file created: {RESET}{out}")

# JSON file
json_file = os.path.join(out_dir, "relationships.json")

json_rows = []

for row in rows:

    person_id = row["ID"]
    clean_id = person_id.replace("@", "")

    profession = individuals.get(person_id, {}).get("occupation", "")

    json_rows.append({
        "ID": clean_id,
        "Name": row["Name"],
        "Relationship": row["Relationship"],
        "SimplifiedPath": row["Simplified path"],
        "RelationshipPath": row["Relationship path"],
        "Generation": int(row["Generation"]) if row["Generation"] not in (None, "") else 0,
        "Occupation": profession,
        "Roles": [],
        "Path": [
            {k: (v.replace("@", "") if k == "ID" else v)
             for k, v in entry.items()}
            for entry in row.get("Path", [])
        ]
    })

if not args.outdir:

    with open(json_file, "w", encoding="utf-8") as f:
        json.dump(json_rows, f, ensure_ascii=False, indent=2)

    print(f"{GREEN}JSON file created: {RESET}{json_file}")

else:

    print(f"{YELLOW}Web mode: relationships.json is not created.{RESET}")

html_file = os.path.join(out_dir, f"{safe_name}.html")

title = f"Relationship report for {proband_name}"
count = len(df)

html = f"""
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<title>{title}</title>

<link rel="stylesheet" href="https://cdn.datatables.net/1.13.6/css/jquery.dataTables.min.css">

<script src="https://code.jquery.com/jquery-3.7.0.min.js"></script>
<script src="https://cdn.datatables.net/1.13.6/js/jquery.dataTables.min.js"></script>

<style>
body {{
    font-family: Arial, sans-serif;
    margin: 30px;
    background:#f7f7f5;
}}

h1 {{
    margin-bottom: 5px;
}}

table {{
    width: 100%;
}}

#relationships_filter input {{
    font-size: 18px;
    width: 177px;
    padding: 6px;
    margin-left: 8px;
}}

.btn {{
    display:inline-flex;
    align-items:center;
    gap:6px;
    padding:8px 12px;
    border:1px solid #e5e7eb;
    border-radius:10px;
    background:#ffffff;
    text-decoration:none;
    color:#111827;
    font-weight:600;
    white-space:nowrap;
}}

.topbar {{
    display:flex;
    align-items:center;
    justify-content:space-between;
    margin:12px 0 20px 0;
}}

.dataTables_length select {{
    font-size:18px;
    padding:8px 10px;
}}

td:nth-child(5), th:nth-child(5) {{
    text-align:right;
}}
</style>

</head>

<body>

<h1>{title}</h1>
<p>{count} people in the report</p>

<div class="topbar">
<a class="btn" href="javascript:history.back()">← Back</a>
<p class="btn">Use &#9652; and &#9662; to sort the columns</p>
</div>

{df.to_html(index=False, table_id="relationships")}

<script>
$(document).ready(function() {{
    $('#relationships').DataTable({{
        pageLength: 50,
        
        columnDefs: [
            {{
                targets: [2, 3],
                searchable: false
            }}
        ],
        
        language: {{
            search: "Search name, relationship, and generation: <B>Here &#8594;</B>",
            lengthMenu: "Show _MENU_ people",
            info: "Showing _START_ to _END_ of _TOTAL_ people",
            infoEmpty: "No people found",
            infoFiltered: "(filtered from _MAX_ people)",
            zeroRecords: "No matching people found",
            paginate: {{
                first: "First",
                last: "Last",
                next: "Next",
                previous: "Previous"
            }}
        }}
    }});
}});
</script>

</body>
</html>
"""

with open(html_file, "w", encoding="utf-8") as f:
    f.write(html)

print(f"{GREEN}HTML file created: {RESET}{html_file}")