strategy.py

from dataclasses import dataclass, field
from typing import List, Dict, Optional, Any, Set

@dataclass
class Champion:
    name: str
    cost: int
    traits: List[str]
    stats: Dict[str, Any]
    damage_type_id: int
    explicit_tags: List[str] = field(default_factory=list)

    # --- Configurable trait classification sets (Set 14) ---
    # Update these when the set rotates. They drive both role and damage
    # inference so item suggestions stay accurate even when the API
    # returns damageType: 1 for every champion.

    # Traits whose holders deal primarily magic / ability damage
    MAGIC_TRAITS: Set[str] = frozenset({
        "Arcanist", "Invoker", "Sorcerer",       # caster archetypes
        "Dark Child", "Rune Mage", "Star Forger", # unique AP carries
        "Chronokeeper",                            # AP-oriented unique
    })

    # Traits whose holders deal primarily physical / auto-attack damage
    ATTACK_TRAITS: Set[str] = frozenset({
        "Slayer", "Gunslinger", "Quickstriker",   # AD melee / ranged
        "Longshot", "Vanquisher",                  # AD ranged / crit
        "Chosen Wolves", "Soulbound",              # AD-oriented uniques
        "Huntress",                                # AD unique
    })

    # Traits that indicate a frontline / tank role
    TANK_TRAITS: Set[str] = frozenset({
        "Bruiser", "Defender", "Juggernaut", "Warden",  # Set 14 tank traits
        "Bastion", "Sentinel", "Guardian", "Vanguard",  # legacy / future-proof
    })

    @property
    def damage_source(self) -> str:
        """Returns 'Attack', 'Magic', or 'Flexible'.

        The CommunityDragon API returns damageType: 1 for ALL champions,
        so we only trust explicit non-1 values (2 = Magic, 3+ = Flexible).
        When damageType is 1 we infer the real damage type from traits
        and, as a last resort, from stats.
        """
        # Trust the API value when it is explicitly Magic or Flexible
        if self.damage_type_id in [2, "Magic"]:
            return "Magic"
        if self.damage_type_id not in [1, "Physical"]:
            return "Flexible"

        # damageType == 1 -- unreliable, infer from traits
        has_magic = any(t in self.MAGIC_TRAITS for t in self.traits)
        has_attack = any(t in self.ATTACK_TRAITS for t in self.traits)

        if has_magic and not has_attack:
            return "Magic"
        if has_attack and not has_magic:
            return "Attack"
        if has_magic and has_attack:
            # Champion has both magic and physical traits -- use stats as
            # tie-breaker: high range + low base AD suggests a caster
            dmg_val = self.stats.get('damage', 50)
            rng_val = self.stats.get('range', 1)
            if not isinstance(dmg_val, (int, float)):
                dmg_val = 50
            if not isinstance(rng_val, (int, float)):
                rng_val = 1
            if rng_val >= 3 and dmg_val < 50:
                return "Magic"
            return "Attack"

        # No recognised damage trait -- stat-based fallback
        rng_val = self.stats.get('range', 1)
        dmg_val = self.stats.get('damage', 50)
        if not isinstance(rng_val, (int, float)):
            rng_val = 1
        if not isinstance(dmg_val, (int, float)):
            dmg_val = 50
        if rng_val >= 4 and dmg_val < 40:
            return "Magic"
        return "Attack"

    @property
    def role(self) -> str:
        """Dynamically determines role based on traits and stats."""
        # 1. Check Traits (The most reliable method)
        if any(t in self.TANK_TRAITS for t in self.traits):
            return "Tank"

        # 2. Check Stats
        # Sanitize stats to ensure they are numbers (handle potential bad data like lists or None)
        rng_val = self.stats.get('range', 1)
        hp_val = self.stats.get('hp', 0)

        if not isinstance(rng_val, (int, float)):
            rng_val = 1

        if not isinstance(hp_val, (int, float)):
            hp_val = 0

        if rng_val >= 3:
            return "Caster" # Backline
        elif hp_val > 800:
            return "Tank" # Frontline
        else:
            return "Fighter" # Melee Carry

    @property
    def granular_role(self) -> str:
        """
        Combines Damage Source + Base Role for itemization logic.
        Examples: "Magic Tank", "Attack Marksman"
        """
        base = self.role
        dmg = self.damage_source

        # Normalize damage source for the string
        if dmg == "Flexible":
            dmg = "Magic" # Default to magic for flexible units in heuristics for now
            
        return f"{dmg} {base}"


class StrategyEngine:
    def __init__(self, champions_data: Dict[str, Dict]):
        self.champions: Dict[str, Champion] = {}
        self._initialize_champions(champions_data)

    def _initialize_champions(self, data: Dict[str, Dict]):
        for name, props in data.items():
            self.champions[name] = Champion(
                name=name,
                cost=props.get('cost', 1),
                traits=props.get('traits', []),
                stats=props.get('stats', {}),
                damage_type_id=props.get('damageType', 1),
                explicit_tags=props.get('roles', [])
            )

    def get_champion(self, name: str) -> Optional[Champion]:
        return self.champions.get(name)

    def analyze_composition(self, champion_names: List[str]) -> Dict[str, Any]:
        """
        Analyzes a list of champions for team balance (Frontline vs Backline, Damage Types).
        """
        team = [self.champions[name] for name in champion_names if name in self.champions]
        
        roles = {"Tank": 0, "Fighter": 0, "Caster": 0}
        damage = {"Attack": 0, "Magic": 0, "Flexible": 0}
        
        for champ in team:
            roles[champ.role] = roles.get(champ.role, 0) + 1
            damage[champ.damage_source] = damage.get(champ.damage_source, 0) + 1
            
        # Basic Heuristics for "Optimal Structure"
        frontline = roles["Tank"] + roles["Fighter"]
        backline = roles["Caster"]
        
        score = 100
        suggestions = []
        
        if frontline == 0:
            score -= 50
            suggestions.append("⚠️ Critical: No frontline units. Team will melt.")
        elif frontline < 2 and len(team) >= 4:
            score -= 20
            suggestions.append("⚠️ Weak frontline. Consider adding a Tank.")
            
        return {
            "score": score,
            "roles": roles,
            "damage_profile": damage,
            "suggestions": suggestions
        }

    def get_suggested_items(self, champion_name: str) -> List[str]:
        """Returns a list of heuristic items based on the champion's granular role."""
        champ = self.get_champion(champion_name)
        if not champ:
            return []

        # Heuristic Item Map
        # In a real app, this could be loaded from external meta data
        item_map = {
            "Attack Tank": ["Warmog's Armor", "Gargoyle Stoneplate", "Steadfast Heart"],
            "Magic Tank": ["Ionic Spark", "Sunfire Cape", "Redemption"], # Magic tanks often apply shred/burn
            "Flexible Tank": ["Warmog's Armor", "Gargoyle Stoneplate", "Sunfire Cape"],
            
            "Attack Fighter": ["Bloodthirster", "Titan's Resolve", "Sterak's Gage"],
            "Magic Fighter": ["Ionic Spark", "Jeweled Gauntlet", "Hand of Justice"],
            "Flexible Fighter": ["Hand of Justice", "Titan's Resolve", "Edge of Night"],
            
            "Attack Caster": ["Infinity Edge", "Last Whisper", "Guinsoo's Rageblade"],
            "Magic Caster": ["Blue Buff", "Jeweled Gauntlet", "Rabadon's Deathcap"],
            "Flexible Caster": ["Guinsoo's Rageblade", "Archangel's Staff", "Hand of Justice"],
        }

        # Fallback for specific roles not in map (e.g. "Attack Caster" might map to generic AD)
        role = champ.granular_role
        return item_map.get(role, ["Thief's Gloves"]) # Default to Thief's Gloves if unsure