#!/usr/bin/env python3
"""
5e.tools Query Script — search monsters, spells, and items by name.

Usage:
    python3 query.py "Mind Flayer"
    python3 query.py "Fireball" --type spell
    python3 query.py "Bag of Holding" --type item
    python3 query.py "Goblin" --type monster --json
"""

import json
import glob
import sys
import os
import argparse
from difflib import SequenceMatcher

DATA_DIR = os.path.expanduser("~/5etools/data")


def load_all_bestiary():
    """Load all monsters from all bestiary files."""
    monsters = []
    for f in sorted(glob.glob(os.path.join(DATA_DIR, "bestiary", "bestiary-*.json"))):
        try:
            data = json.load(open(f))
            monsters.extend(data.get("monster", []))
        except (json.JSONDecodeError, FileNotFoundError):
            continue
    return monsters


def load_all_spells():
    """Load all spells from all spell files."""
    spells = []
    for f in sorted(glob.glob(os.path.join(DATA_DIR, "spells", "spells-*.json"))):
        try:
            data = json.load(open(f))
            spells.extend(data.get("spell", []))
        except (json.JSONDecodeError, FileNotFoundError):
            continue
    return spells


def load_all_items():
    """Load all items from items.json."""
    try:
        data = json.load(open(os.path.join(DATA_DIR, "items.json")))
        return data.get("item", [])
    except (json.JSONDecodeError, FileNotFoundError):
        return []


def find_by_name(items, name, threshold=0.6):
    """Find items by name using fuzzy matching."""
    name_lower = name.lower()
    exact = []
    fuzzy = []

    for item in items:
        item_name = item.get("name", "").lower()
        if item_name == name_lower:
            exact.append(item)
        elif name_lower in item_name or item_name in name_lower:
            fuzzy.append(item)
        else:
            ratio = SequenceMatcher(None, name_lower, item_name).ratio()
            if ratio >= threshold:
                fuzzy.append((item, ratio))

    if exact:
        return exact
    if fuzzy:
        if isinstance(fuzzy[0], tuple):
            fuzzy.sort(key=lambda x: -x[1])
            return [f[0] for f in fuzzy]
        return fuzzy
    return []


def clean_text(text):
    """Remove 5e.tools formatting tags from text."""
    if not isinstance(text, str):
        return str(text)
    import re
    # Remove {@tag ...} markers including nested braces
    text = re.sub(r'\{@[^}]*\}', '', text)
    # Clean up double spaces
    text = re.sub(r' +', ' ', text)
    return text.strip()


def format_ac(ac):
    """Format armor class for display."""
    if isinstance(ac, list):
        parts = []
        for a in ac:
            if isinstance(a, dict):
                val = a.get("ac", "?")
                src = a.get("from", [])
                src_str = ", ".join(clean_text(s) for s in src) if src else ""
                parts.append(f"{val} ({src_str})" if src_str else str(val))
            else:
                parts.append(str(a))
        return " | ".join(parts)
    return str(ac)


def format_hp(hp):
    """Format hit points for display."""
    if isinstance(hp, dict):
        avg = hp.get("average", "?")
        formula = hp.get("formula", "")
        return f"{avg} ({formula})" if formula else str(avg)
    return str(hp)


def format_stat_block(monster):
    """Format a monster stat block for display."""
    lines = []
    name = monster.get("name", "Unknown")
    size = monster.get("size", "")
    if isinstance(size, list):
        size = ", ".join(size)
    mtype = monster.get("type", "")
    alignment = monster.get("alignment", "")
    if isinstance(alignment, list):
        alignment = ", ".join(alignment)
    lines.append(f"**{name}** — {size} {mtype}, {alignment}")
    lines.append("─" * 50)

    ac = format_ac(monster.get("ac", "?"))
    hp = format_hp(monster.get("hp", "?"))
    speed = monster.get("speed", {})
    if isinstance(speed, dict):
        speed_str = "; ".join(f"{k} {v}" for k, v in speed.items())
    else:
        speed_str = str(speed)
    lines.append(f"AC {ac} | HP {hp} | Speed {speed_str}")

    stats = {s: monster.get(s, "?") for s in ["str", "dex", "con", "int", "wis", "cha"]}
    lines.append(f"STR {stats['str']} | DEX {stats['dex']} | CON {stats['con']} | INT {stats['int']} | WIS {stats['wis']} | CHA {stats['cha']}")

    # Saving throws
    saves = monster.get("save", {})
    if saves:
        save_str = " | ".join(f"{k.upper()} {v}" for k, v in saves.items())
        lines.append(f"Saving Throws: {save_str}")

    # Skills
    skills = monster.get("skill", {})
    if skills:
        skill_str = " | ".join(f"{k.capitalize()} {v}" for k, v in skills.items())
        lines.append(f"Skills: {skill_str}")

    # Senses
    senses = monster.get("senses", "")
    if isinstance(senses, list):
        senses = ", ".join(clean_text(s) for s in senses)
    passive = monster.get("passive", "")
    senses_str = clean_text(senses) if isinstance(senses, str) else str(senses)
    if passive:
        senses_str = f"{senses_str}, Passive Perception {passive}" if senses_str else f"Passive Perception {passive}"
    if senses_str:
        lines.append(f"Senses: {senses_str}")

    # Languages
    languages = monster.get("languages", "")
    if isinstance(languages, list):
        languages = ", ".join(clean_text(l) for l in languages)
    if languages:
        lines.append(f"Languages: {clean_text(languages) if isinstance(languages, str) else languages}")

    # CR
    cr = monster.get("cr", "?")
    xp = monster.get("xp", "")
    cr_str = f"CR {cr}"
    if xp:
        cr_str += f" ({xp} XP)"
    lines.append(cr_str)

    # Traits
    for trait in monster.get("trait", []):
        name_t = clean_text(trait.get("name", ""))
        entries = trait.get("entries", [])
        entry_text = " ".join(clean_text(str(e)) for e in entries) if isinstance(entries, list) else clean_text(str(entries))
        lines.append(f"\n**{name_t}.** {entry_text}")

    # Actions
    for action in monster.get("action", []):
        name_a = clean_text(action.get("name", ""))
        entries = action.get("entries", [])
        entry_text = " ".join(clean_text(str(e)) for e in entries) if isinstance(entries, list) else clean_text(str(entries))
        lines.append(f"\n**{name_a}.** {entry_text}")

    # Legendary Actions
    legendary = monster.get("legendary", [])
    if legendary:
        lines.append("\n**Legendary Actions**")
        for la in legendary:
            name_l = clean_text(la.get("name", ""))
            entries = la.get("entries", [])
            entry_text = " ".join(clean_text(str(e)) for e in entries) if isinstance(entries, list) else clean_text(str(entries))
            lines.append(f"\n**{name_l}.** {entry_text}")

    return "\n".join(lines)


def format_spell(spell):
    """Format a spell for display."""
    lines = []
    name = spell.get("name", "Unknown")
    level = spell.get("level", "?")
    school = spell.get("school", "")
    lines.append(f"**{name}** — Level {level} {school}")
    lines.append("─" * 50)

    time = spell.get("time", [])
    if time:
        t = time[0] if isinstance(time, list) else time
        if isinstance(t, dict):
            lines.append(f"Casting Time: {t.get('number', '')} {t.get('unit', '')}")
        else:
            lines.append(f"Casting Time: {t}")

    lines.append(f"Range: {spell.get('range', '?')}")
    components = spell.get("components", {})
    comp_str = ""
    if components.get("v"): comp_str += "V "
    if components.get("s"): comp_str += "S "
    if components.get("m"): comp_str += f"M ({components['m']})"
    if comp_str:
        lines.append(f"Components: {comp_str.strip()}")

    duration = spell.get("duration", [])
    if duration:
        d = duration[0] if isinstance(duration, list) else duration
        if isinstance(d, dict):
            dur_str = f"{d.get('type', '')}"
            if d.get("concentration"):
                dur_str = f"Concentration, {dur_str}"
            lines.append(f"Duration: {dur_str}")
        else:
            lines.append(f"Duration: {d}")

    classes = spell.get("classes", {})
    if classes:
        class_list = []
        for cls_type, cls_names in classes.items():
            if isinstance(cls_names, list):
                class_list.extend(cls_names)
        if class_list:
            lines.append(f"Classes: {', '.join(class_list)}")

    entries = spell.get("entries", [])
    for entry in entries:
        if isinstance(entry, str):
            lines.append(f"\n{entry}")
        elif isinstance(entry, dict):
            if "name" in entry:
                lines.append(f"\n**{entry['name']}.** {' '.join(str(e) for e in entry.get('entries', []))}")
            else:
                lines.append(f"\n{' '.join(str(e) for e in entry.get('entries', []))}")

    return "\n".join(lines)


def format_item(item):
    """Format a magic item for display."""
    lines = []
    name = item.get("name", "Unknown")
    rarity = item.get("rarity", "")
    item_type = item.get("type", "")
    lines.append(f"**{name}** — {rarity} {item_type}")
    lines.append("─" * 50)

    value = item.get("value")
    weight = item.get("weight")
    if value or weight:
        details = []
        if value:
            details.append(f"Value: {value}")
        if weight:
            details.append(f"Weight: {weight} lb")
        lines.append(" | ".join(details))

    entries = item.get("entries", [])
    for entry in entries:
        if isinstance(entry, str):
            lines.append(f"\n{entry}")
        elif isinstance(entry, dict):
            if "name" in entry:
                lines.append(f"\n**{entry['name']}.** {' '.join(str(e) for e in entry.get('entries', []))}")
            else:
                lines.append(f"\n{' '.join(str(e) for e in entry.get('entries', []))}")

    return "\n".join(lines)


def main():
    parser = argparse.ArgumentParser(description="Query 5e.tools data")
    parser.add_argument("name", help="Name to search for")
    parser.add_argument("--type", choices=["monster", "spell", "item"], default="monster",
                        help="Type of data to search")
    parser.add_argument("--json", action="store_true", help="Output raw JSON")
    args = parser.parse_args()

    if args.type == "monster":
        items = load_all_bestiary()
        formatter = format_stat_block
    elif args.type == "spell":
        items = load_all_spells()
        formatter = format_spell
    else:
        items = load_all_items()
        formatter = format_item

    results = find_by_name(items, args.name)

    if not results:
        print(f"No {args.type}s found matching '{args.name}'")
        sys.exit(1)

    if args.json:
        print(json.dumps(results, indent=2))
        return

    for i, result in enumerate(results):
        if i > 0:
            print("\n" + "=" * 50 + "\n")
        print(formatter(result))


if __name__ == "__main__":
    main()