import gradio as gr
import json
from collections import defaultdict
import re

def load_families_from_jsonl():
    people = {}  # name -> record
    # 建立夫妻映射表，用于补齐缺失的父母
    spouse_map = {} 

    try:
        if not os.path.exists('family_data.jsonl'):
            return {"__ERROR__": "文件 family_data.jsonl 不存在"}
            
        with open('family_data.jsonl', 'r', encoding='utf-8') as f:
            for line in f:
                line = line.strip()
                if not line: continue
                try:
                    rec = json.loads(line)
                    name = rec.get('n', '').strip()
                    if name:
                        people[name] = rec
                        # 记录配偶关系 (处理字符串或列表)
                        sp = rec.get('sp', "无")
                        if sp and sp != "无":
                            sp_list = sp if isinstance(sp, list) else [sp]
                            for s in sp_list:
                                spouse_map[name] = s
                                spouse_map[s] = name
                except: continue
    except Exception as e:
        return {"__ERROR__": f"读取失败: {e}"}

    # 第二步：按父母分组子女
    children_by_parents = defaultdict(list)
    
    for name, rec in people.items():
        f = rec.get('f', '').strip()
        m = rec.get('m', '').strip()
        
        # 核心逻辑：补齐缺失的父母
        # 如果只有父亲，尝试找母亲
        if f and (not m or m == "无"):
            m = spouse_map.get(f, "不详")
        # 如果只有母亲，尝试找父亲
        elif m and (not f or f == "无"):
            f = spouse_map.get(m, "不详")
        
        # 只要有一方存在，就归入家庭
        if (f and f != "无") or (m and m != "无"):
            # 统一占位符
            f = f if f else "不详"
            m = m if m else "不详"
            children_by_parents[(f, m)].append((name, rec.get('info', '')))

    # 第三步：构建家庭列表
    families = []
    
    def extract_year(info):
        match = re.search(r'(\d{4})', str(info))
        return int(match.group(1)) if match else 9999

    for (father, mother), children in children_by_parents.items():
        # 获取信息
        father_info = people.get(father, {}).get('info', '资料待补充')
        mother_info = people.get(mother, {}).get('info', '资料待补充')

        # 子女去重并按年份排序
        unique_children = sorted(list(set(children)), key=lambda x: extract_year(x[1]))

        families.append({
            "father": father,
            "father_info": father_info,
            "mother": mother,
            "mother_info": mother_info,
            "children": unique_children
        })

    # 按父亲/母亲名字排序
    families.sort(key=lambda x: x["father"] if x["father"] != "不详" else x["mother"])
    return families

def format_family_output():
    families = load_families_from_jsonl()
    
    if isinstance(families, dict) and "__ERROR__" in families:
        return families["__ERROR__"]
    
    if not families:
        return "未找到任何家庭数据。"

    lines = ["# 冯氏家族关系网 · 家庭成员表\n"]
    
    for fam in families:
        f_name = fam["father"]
        m_name = fam["mother"]
        f_info = fam["father_info"]
        m_info = fam["mother_info"]
        children = fam["children"]
        
        lines.append(f"## {f_name} & {m_name} 家庭")
        lines.append(f"- **父亲**：{f_name} {'（'+f_info+'）' if f_name != '不详' else ''}")
        lines.append(f"- **母亲**：{m_name} {'（'+m_info+'）' if m_name != '不详' else ''}")
        lines.append("- **子女**：")
        if children:
            for child_name, child_info in children:
                lines.append(f"  - **{child_name}** — {child_info}")
        else:
            lines.append("  - （无子女记录）")
        lines.append("\n---") 

    return "\n".join(lines)

# Gradio 保持不变
import os
demo = gr.Interface(
    fn=format_family_output,
    inputs=None,
    outputs=gr.Markdown(),
    title="冯氏家族关系网 · 家庭成员表",
    description="自动关联单亲记录并补齐配偶信息。"
)

if __name__ == "__main__":
    demo.launch()