data_loader.py

# src/tft_analyzer/data_loader.py
import requests
from requests.exceptions import RequestException
import json
from pathlib import Path
from typing import Dict, List, Any

TFT_DATA_URL = "https://raw.communitydragon.org/latest/cdragon/tft/en_us.json"
DATA_DIR = Path(__file__).parent / "data"
TFT_DATA_FILE = DATA_DIR / "tft_data.json"
CONSTANTS_FILE = DATA_DIR / "game_constants.json"

def download_tft_data(url: str = TFT_DATA_URL, output_path: Path = TFT_DATA_FILE):
    """Downloads TFT data from the given URL and saves it to the output path."""
    try:
        print("Downloading latest TFT data...")
        response = requests.get(url, timeout=30)
        response.raise_for_status()
        data = response.json()
        
        DATA_DIR.mkdir(exist_ok=True)
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(data, f, indent=4)
        print(f"Successfully downloaded TFT data to {output_path}")
    except RequestException as e:
        print(f"Error downloading TFT data: {e}. Will try to use existing file.")

def load_tft_data(path: Path = TFT_DATA_FILE) -> Dict[str, Any]:
    """Loads the TFT data from the specified local JSON file."""
    if not path.exists():
        download_tft_data()

    if not path.exists():
        raise FileNotFoundError(
            f"TFT data file not found at '{path}' and could not be downloaded. "
            f"Please check your internet connection and try again, or manually "
            f"place the data file at '{path}'."
        )

    with open(path, 'r', encoding='utf-8') as f:
        return json.load(f)

def get_champion_image_url(asset_path: str) -> str:
    """Converts a CommunityDragon asset path to a CDN image URL."""
    if not asset_path:
        return ""
    base_url = "https://raw.communitydragon.org/latest/game/"
    return base_url + asset_path.lower().replace(".tex", ".png")


def get_champion_data(tft_data: Dict[str, Any]) -> Dict[str, Dict[str, Any]]:
    """Parses the TFT data and returns a dictionary of champion data."""
    champions = {}
    
    if "sets" not in tft_data:
        raise ValueError("TFT data does not contain 'sets' key.")
        
    # Get the latest set
    latest_set = tft_data["sets"][str(max(map(int, tft_data["sets"].keys())))]
    
    if "champions" not in latest_set:
        raise ValueError("Latest set data does not contain 'champions' key.")

    for champ in latest_set["champions"]:
        # Skip dummy/test champions and spawned units
        if champ.get("isSpawn", False) or not champ.get("name"):
            continue

        cost = champ.get("cost")
        name = champ.get("name", "").strip()
        traits = champ.get("traits", [])

        # Ensure traits is a list
        if not isinstance(traits, list):
            traits = []

        # Only include playable champions: cost 1-5 with at least one trait
        if not name or not isinstance(cost, int) or cost < 1 or cost > 5 or not traits:
            continue

        # Sanitize null stat values to 0
        raw_stats = champ.get("stats", {})
        stats = {k: (v if isinstance(v, (int, float)) else 0) for k, v in raw_stats.items()}

        champions[name] = {
            "cost": cost,
            "traits": traits,
            "stats": stats,
            "damageType": champ.get("damageType", 1),
            "roles": champ.get("roles", []),
            "icon": champ.get("icon", ""),
            "squareIcon": champ.get("squareIcon", ""),
            "apiName": champ.get("apiName", ""),
        }
                
    return champions

def get_trait_data(tft_data: Dict[str, Any]) -> Dict[str, Any]:
    """Parses the TFT data and returns a dictionary of trait data."""
    traits = {}
    
    if "sets" not in tft_data:
        raise ValueError("TFT data does not contain 'sets' key.")
        
    # Get the latest set
    latest_set = tft_data["sets"][str(max(map(int, tft_data["sets"].keys())))]
    
    if "traits" not in latest_set:
        raise ValueError("Latest set data does not contain 'traits' key.")

    for trait in latest_set["traits"]:
        name = trait.get("name")
        if name:
            traits[name] = {
                "effects": trait.get("effects", [])
            }
                
    return traits

def get_item_data(tft_data: Dict[str, Any]) -> Dict[str, Any]:
    """Parses the TFT data and returns a dictionary of item data."""
    items = {}
    # Items are usually at the top level
    for item in tft_data.get("items", []):
        name = item.get("name")
        if name:
            items[name] = {
                "id": item.get("id"),
                "desc": item.get("desc"),
                "effects": item.get("effects", {}),
                "icon": item.get("icon")
            }
    return items

def get_champions_by_tier(tft_data: Dict[str, Any]) -> Dict[str, List[str]]:
    """Parses the TFT data and returns a dictionary of champion names by tier."""
    champions_by_tier = {f"{i}-cost": [] for i in range(1, 6)}
    champ_data = get_champion_data(tft_data)
    
    for name, data in champ_data.items():
        cost = data["cost"]
        if 1 <= cost <= 5:
            champions_by_tier[f"{cost}-cost"].append(name)
            
    return champions_by_tier

def save_split_data(tft_data: Dict[str, Any]):
    """Splits and saves the data into separate JSON files."""
    print("Splitting data into categories...")
    
    save_map = {
        "champions.json": get_champion_data(tft_data),
        "traits.json": get_trait_data(tft_data),
        "items.json": get_item_data(tft_data)
    }
    
    for filename, data in save_map.items():
        path = DATA_DIR / filename
        with open(path, 'w', encoding='utf-8') as f:
            json.dump(data, f, indent=4)
        print(f"Saved {path}")

def save_game_constants():
    """
    Saves static game rules (Pool Sizes, Shop Odds) to a JSON file.
    These values are not typically available in the standard en_us.json,
    so we maintain them here as a configurable source of truth.
    """
    constants = {
        "pool_sizes": {
            "1": 29, "2": 22, "3": 18, "4": 12, "5": 10
        },
        "shop_odds": {
            "1":  {"1": 100, "2": 0, "3": 0, "4": 0, "5": 0},
            "2":  {"1": 100, "2": 0, "3": 0, "4": 0, "5": 0},
            "3":  {"1": 75, "2": 25, "3": 0, "4": 0, "5": 0},
            "4":  {"1": 55, "2": 30, "3": 15, "4": 0, "5": 0},
            "5":  {"1": 45, "2": 33, "3": 20, "4": 2, "5": 0},
            "6":  {"1": 25, "2": 40, "3": 30, "4": 5, "5": 0},
            "7":  {"1": 19, "2": 30, "3": 35, "4": 15, "5": 1},
            "8":  {"1": 18, "2": 25, "3": 32, "4": 22, "5": 3},
            "9":  {"1": 10, "2": 20, "3": 25, "4": 35, "5": 10},
            "10": {"1": 5, "2": 10, "3": 20, "4": 40, "5": 25}
        }
    }
    
    if not CONSTANTS_FILE.exists():
        print("Generating default game constants file...")
        with open(CONSTANTS_FILE, 'w', encoding='utf-8') as f:
            json.dump(constants, f, indent=4)
        print(f"Saved {CONSTANTS_FILE}")

if __name__ == "__main__":
    # This makes the script runnable to manually update data
    download_tft_data()
    # Example of loading and parsing
    data = load_tft_data()
    
    # Split the data for the analyzer
    save_split_data(data)
    save_game_constants()