app.py

import gradio as gr
import pandas as pd
import yaml
import io
import random
import os
from collections import deque
import tempfile
from dataclasses import dataclass, field
from typing import List, Dict, Any, Optional
import json  # Добавьте импорт json

# ===============================================================
# 1. DATA CLASSES (полная структура)
# ===============================================================
@dataclass
class MergeChainItemData:
    MergeItemId: str
    RequirementWeight: int
    RewardDifficulty: int

@dataclass
class MergeChain:
    Id: str
    Items: List[MergeChainItemData] = field(default_factory=list)

@dataclass
class GenerationReward:
    Amount: int
    MergeItemId: str = ""
    Type: str = "Energy"
    RewardWeight: int = 100
    ReductionFactor: int = 0

@dataclass
class GenerationRewardWithDifficulty:
    DifficultyScore: int
    Rewards: List[GenerationReward] = field(default_factory=list)

@dataclass
class RequirementWeight:
    Weights: List[int] = field(default_factory=lambda: [70, 30])

@dataclass
class GeneratorSettings:
    Id: str = "DefaultSettings"
    DefaultRequirementWeights: RequirementWeight = field(default_factory=RequirementWeight)
    MaxActiveOrders: int = 4
    ReductionFactor: int = 3
    IncreaseFactor: int = 5

@dataclass
class MergeGeneratorRuleset:
    Id: str = "DefaultRuleset"
    MaxHistoryOrders: int = 5
    IncrementDifficulty: int = 2
    OverallChanceToDropExpeditionEnergy: int = 90
    OverrideWeights: Dict[str, int] = field(default_factory=dict)
    EnergyRewards: List[GenerationRewardWithDifficulty] = field(default_factory=list)
    ItemRewards: List[GenerationRewardWithDifficulty] = field(default_factory=list)
    OverrideMaxRequirementOrders: bool = False
    OverrideMaxOrdersWithWeight: RequirementWeight = field(default_factory=RequirementWeight)

@dataclass
class SimulatedOrder:
    Requirements: list = field(default_factory=list)
    Rewards: list = field(default_factory=list)
    TotalDifficulty: int = 0
    MergeEnergyPrice: int = 0

# ===============================================================
# 2. ADVANCED PARSERS
# ===============================================================
def robust_asset_parser(file_content: str) -> dict:
    try:
        lines = file_content.splitlines()
        mono_behaviour_line_index = next((i for i, line in enumerate(lines) if "MonoBehaviour:" in line), -1)
        if mono_behaviour_line_index == -1: return {}
        data_lines = lines[mono_behaviour_line_index + 1:]
        if not data_lines: return {}
        indentation = len(data_lines[0]) - len(data_lines[0].lstrip(' '))
        dedented_lines = [(line[indentation:] if line.strip() else "") for line in data_lines]
        yaml_string = "\n".join(dedented_lines)
        parsed_data = yaml.safe_load(yaml_string) or {}
        # FIX: Clean binary data fields
        for key, value in parsed_data.items():
            if isinstance(value, dict):
                for sub_key, sub_value in value.items():
                    if isinstance(sub_value, str) and len(sub_value) == 16 and all(c in '0123456789abcdef' for c in sub_value.lower()):
                        parsed_data[key][sub_key] = "70,30"  # Default fallback for binary weight data
        return parsed_data
    except Exception as e:
        print(f"Asset parsing error: {e}")
        return {}

def load_chain_config(file) -> MergeChain:
    with open(file.name, 'r', encoding='utf-8') as f: content = f.read()
    data = robust_asset_parser(content)
    chain = MergeChain(Id=data.get('_id'))
    for item_dict in data.get('_mergeChainItemsData', []):
        chain.Items.append(MergeChainItemData(
            MergeItemId=item_dict.get('MergeItemId'),
            RequirementWeight=int(item_dict.get('RequirementWeight', 0)),
            RewardDifficulty=int(item_dict.get('RewardDifficulty', 0))
        ))
    return chain

def parse_rewards_from_data(reward_list_data, reward_type) -> List[GenerationRewardWithDifficulty]:
    rewards_with_difficulty = []
    if not isinstance(reward_list_data, list): return rewards_with_difficulty
    
    for reward_group in reward_list_data:
        # FIX: Safe conversion for DifficultyScore
        try:
            difficulty = int(reward_group.get('DifficultyScore', 0))
        except (ValueError, TypeError):
            difficulty = 0
        
        rewards = []
        
        if reward_type == "Energy":
            if 'Reward' in reward_group and reward_group['Reward']:
                reward_info = reward_group['Reward']
                # FIX: Safe conversion for amount
                try:
                    amount = int(reward_info.get('_amount', 0))
                except (ValueError, TypeError):
                    amount = 0
                rewards.append(GenerationReward(Amount=amount, Type='Energy'))
        elif reward_type == "Item":
            if 'Rewards' in reward_group and isinstance(reward_group['Rewards'], list):
                for weighted_reward in reward_group['Rewards']:
                    if 'Reward' in weighted_reward and weighted_reward['Reward']:
                        reward_info = weighted_reward['Reward']
                        # FIX: Safe conversions for all numeric fields
                        try:
                            amount = int(reward_info.get('_amount', 1))
                        except (ValueError, TypeError):
                            amount = 1
                        try:
                            reward_weight = int(weighted_reward.get('RewardWeight', 100))
                        except (ValueError, TypeError):
                            reward_weight = 100
                        try:
                            reduction_factor = int(weighted_reward.get('ReductionFactor', 0))
                        except (ValueError, TypeError):
                            reduction_factor = 0
                            
                        rewards.append(GenerationReward(
                            Amount=amount,
                            MergeItemId=reward_info.get('_mergeItemId', ''),
                            Type='Item',
                            RewardWeight=reward_weight,
                            ReductionFactor=reduction_factor
                        ))
        if rewards:
            rewards_with_difficulty.append(GenerationRewardWithDifficulty(DifficultyScore=difficulty, Rewards=rewards))
            
    return rewards_with_difficulty

def load_ruleset_config(file) -> MergeGeneratorRuleset:
    # DEBUG_LOG: Asset parsing initiation
    print("--- Loading ruleset config ---")
    with open(file.name, 'r', encoding='utf-8') as f: content = f.read()
    data = robust_asset_parser(content)
    print(f"Parsed raw data from {file.name}")

    # ERROR_FIX: Binary weight data handling + hex validation
    raw_weights = data.get('_overrideMaxOrdersWithWeight', {}).get('_requirementOrderWeights', "70,30")
    if isinstance(raw_weights, str) and ',' in raw_weights:
        try:
            weights_list = [int(w.strip()) for w in raw_weights.split(',')]
        except ValueError:
            weights_list = [70, 30]  # FALLBACK_DEFAULT
    else:
        weights_list = [70, 30]  # FALLBACK_BINARY_DATA

    ruleset = MergeGeneratorRuleset(
        Id=data.get('_id'),
        MaxHistoryOrders=int(data.get('_maxHistoryOrders', 5)),
        IncrementDifficulty=int(data.get('_incrementDifficulty', 2)),
        OverallChanceToDropExpeditionEnergy=int(data.get('_overallChanceToDropExpeditionEnergy', 90)),
        OverrideMaxRequirementOrders=bool(int(data.get('_overrideMaxRequirementOrders', 0))),
        OverrideMaxOrdersWithWeight=RequirementWeight(Weights=weights_list)
    )
    ruleset.OverrideWeights = {ow.get('_mergeItemId'): int(ow.get('_weight', 0)) for ow in data.get('_overrideWeights', [])}
    ruleset.EnergyRewards = parse_rewards_from_data(data.get('_energyRewards', []), "Energy")
    ruleset.ItemRewards = parse_rewards_from_data(data.get('_itemRewards', []), "Item")
    print(f"Loaded {len(ruleset.EnergyRewards)} energy reward groups and {len(ruleset.ItemRewards)} item reward groups.")
    return ruleset

def load_settings_config(file) -> GeneratorSettings:
    with open(file.name, 'r', encoding='utf-8') as f: content = f.read()
    data = robust_asset_parser(content)
    return GeneratorSettings(
        Id=data.get('m_Name', 'MergeGeneratorSettings'),
        MaxActiveOrders=int(data.get('_maxActiveOrders', 4)),
        ReductionFactor=int(data.get('_reductionFactor', 3)),
        IncreaseFactor=int(data.get('_increaseFactor', 5))
    )

# ===============================================================
# 3. SIMULATION LOGIC
# ===============================================================
def get_requirement_count(ruleset, settings):
    weights_source = settings.DefaultRequirementWeights if not ruleset.OverrideMaxRequirementOrders else ruleset.OverrideMaxOrdersWithWeight
    weights = weights_source.Weights
    if len(weights) < 2: weights.extend([0] * (2 - len(weights)))
    return random.choices([1, 2], weights=weights[:2], k=1)[0]

def generate_rewards(order, ruleset):
    if random.randint(1, 100) <= ruleset.OverallChanceToDropExpeditionEnergy:
        # FIX: Handle NaN values in DifficultyScore
        suitable_reward_groups = [rg for rg in ruleset.EnergyRewards if pd.notna(rg.DifficultyScore) and order.TotalDifficulty >= int(rg.DifficultyScore)]
        if suitable_reward_groups:
            chosen_group = max(suitable_reward_groups, key=lambda rg: int(rg.DifficultyScore))
            if chosen_group.Rewards: order.Rewards.append(chosen_group.Rewards[0])
    else:
        # FIX: Handle NaN values in DifficultyScore
        suitable_reward_groups = [rg for rg in ruleset.ItemRewards if pd.notna(rg.DifficultyScore) and order.TotalDifficulty >= int(rg.DifficultyScore)]
        if suitable_reward_groups:
            chosen_group = max(suitable_reward_groups, key=lambda rg: int(rg.DifficultyScore))
            if chosen_group.Rewards:
                rewards, weights = chosen_group.Rewards, [r.RewardWeight for r in chosen_group.Rewards]
                if sum(weights) > 0: order.Rewards.append(random.choices(rewards, weights=weights, k=1)[0])

def run_simulation_logic(chains, ruleset, settings, iteration_count, initial_energy):
    order_history = deque(maxlen=ruleset.MaxHistoryOrders)
    chain_unlock_levels = {chain.Id: 1 for chain in chains}
    simulation_results = []
    current_energy = initial_energy

    for i in range(iteration_count):
        all_available_items = [item for chain in chains for level, item in enumerate(chain.Items, 1) if item.RequirementWeight > 0 and level <= chain_unlock_levels.get(chain.Id, 1)]
        recently_used_ids = {req.MergeItemId for order in order_history for req in order.Requirements}
        final_items = [item for item in all_available_items if item.MergeItemId not in recently_used_ids] or all_available_items
        if not final_items: continue

        req_count = get_requirement_count(ruleset, settings)
        order = SimulatedOrder()
        used_items_in_order = []

        for _ in range(req_count):
            selectable_items = [item for item in final_items if item not in used_items_in_order]
            if not selectable_items: 
                break
            
            weights = [ruleset.OverrideWeights.get(item.MergeItemId, item.RequirementWeight) for item in selectable_items]
            if sum(weights) == 0: 
                continue
            
            selected_item = random.choices(selectable_items, weights=weights, k=1)[0]
            order.Requirements.append(selected_item)
            order.TotalDifficulty += selected_item.RewardDifficulty
            used_items_in_order.append(selected_item)
        
        if not order.Requirements: 
            continue

        total_cost = sum(2**(chain.Items.index(req)) for req in order.Requirements 
                        if (chain := next((c for c in chains if req in c.Items), None)))
        order.MergeEnergyPrice = total_cost
        current_energy -= total_cost
        
        generate_rewards(order, ruleset)
        order_history.append(order)

        for req in order.Requirements:
            chain_of_item = next((c for c in chains if req in c.Items), None)
            if chain_of_item and (current_level := chain_of_item.Items.index(req) + 1) == chain_unlock_levels.get(chain_of_item.Id, 1):
                chain_unlock_levels[chain_of_item.Id] += 1
        
        row = {
            "Order": i + 1, 
            "Total_Difficulty": order.TotalDifficulty, 
            "MergeEnergyPrice": order.MergeEnergyPrice, 
            "MEnergy_Amount": current_energy
        }
        
        for j, req in enumerate(order.Requirements):
            chain_of_item = next((c for c in chains if req in c.Items), None)
            row[f'Requirement_{j+1}'] = req.MergeItemId
            row[f'Weight_{j+1}'] = ruleset.OverrideWeights.get(req.MergeItemId, req.RequirementWeight)
            row[f'ChainId_{j+1}'] = chain_of_item.Id if chain_of_item else "N/A"
            row[f'Level_{j+1}'] = chain_of_item.Items.index(req) + 1 if chain_of_item else 0
            row[f'RewardDifficulty_{j+1}'] = req.RewardDifficulty
        
        row['ExpeditionEnergyReward'] = next((r.Amount for r in order.Rewards if r.Type == 'Energy'), 0)
        row['MergeItemReward'] = next((r.MergeItemId for r in order.Rewards if r.Type == 'Item'), "")
        simulation_results.append(row)

    df = pd.DataFrame(simulation_results).fillna(0)
    full_column_list = [
        'Order', 'MergeEnergyPrice', 'MEnergy_Amount', 'Total_Difficulty', 'ExpeditionEnergyReward', 'MergeItemReward',
        'Requirement_1', 'Weight_1', 'ChainId_1', 'Level_1', 'RewardDifficulty_1',
        'Requirement_2', 'Weight_2', 'ChainId_2', 'Level_2', 'RewardDifficulty_2'
    ]
    
    for col in full_column_list:
        if col not in df.columns: 
            df[col] = 0
    
    stats_report = f"SIMULATION_RESULT: ORDERS={len(df)} AVG_DIFFICULTY={df['Total_Difficulty'].mean():.2f} FINAL_ENERGY={current_energy}"
    return df[full_column_list], stats_report

def run_simulation_interface(
    chain_df, energy_rewards_df, item_rewards_df,
    max_hist, inc_diff, energy_chance, req_weights_str,
    red_factor, inc_factor, iteration_count, initial_energy
):
    """
    INTERFACE_WRAPPER: SIMULATION_EXECUTION
    """
    if chain_df is None or chain_df.empty: 
        raise gr.Error("CHAIN_DATA: EMPTY → LOAD_REQUIRED")
    
    chains = [MergeChain(Id=chain_id, Items=[
        MergeChainItemData(row['MergeItemId'], int(row['RequirementWeight']), int(row['RewardDifficulty'])) 
        for _, row in group.iterrows()
    ]) for chain_id, group in chain_df.groupby('ChainId')]
    
    settings = GeneratorSettings(
        ReductionFactor=red_factor, 
        IncreaseFactor=inc_factor, 
        DefaultRequirementWeights=RequirementWeight(Weights=[int(w.strip()) for w in req_weights_str.split(',')])
    )
    
    ruleset = MergeGeneratorRuleset(
        MaxHistoryOrders=max_hist, 
        IncrementDifficulty=inc_diff, 
        OverallChanceToDropExpeditionEnergy=energy_chance
    )
    
    if energy_rewards_df is not None and not energy_rewards_df.empty:
        df_copy = energy_rewards_df.dropna().copy()
        df_copy['DifficultyScore'] = pd.to_numeric(df_copy['DifficultyScore'])
        df_copy['Amount'] = pd.to_numeric(df_copy['Amount'])
        ruleset.EnergyRewards = [GenerationRewardWithDifficulty(r['DifficultyScore'], [GenerationReward(Amount=r['Amount'], Type='Energy')]) 
                               for i, r in df_copy.iterrows()]
        
    if item_rewards_df is not None and not item_rewards_df.empty:
        df_copy = item_rewards_df.dropna().copy()
        df_copy['DifficultyScore'] = pd.to_numeric(df_copy['DifficultyScore'])
        df_copy['Amount'] = pd.to_numeric(df_copy['Amount'])
        df_copy['RewardWeight'] = pd.to_numeric(df_copy['RewardWeight'])
        df_copy['ReductionFactor'] = pd.to_numeric(df_copy['ReductionFactor'])
        
        for score, group in df_copy.groupby('DifficultyScore'):
            rewards = [GenerationReward(
                Amount=r['Amount'], 
                MergeItemId=r['MergeItemId'], 
                Type='Item', 
                RewardWeight=r['RewardWeight'], 
                ReductionFactor=r['ReductionFactor']
            ) for i, r in group.iterrows()]
            ruleset.ItemRewards.append(GenerationRewardWithDifficulty(int(score), rewards))

    df, stats_report = run_simulation_logic(chains, ruleset, settings, iteration_count, initial_energy)
    
    with tempfile.NamedTemporaryFile(delete=False, mode='w', suffix='.csv', encoding='utf-8', newline='') as tmp_csv:
        df.to_csv(tmp_csv.name, index=False)
        csv_path = tmp_csv.name
        
    with tempfile.NamedTemporaryFile(delete=False, mode='w', suffix='.txt', encoding='utf-8') as tmp_txt:
        tmp_txt.write(stats_report)
        report_path = tmp_txt.name

    return df, stats_report, gr.update(value=csv_path, visible=True), gr.update(value=report_path, visible=True)

# ===============================================================
# 4. UI UPDATE FUNCTIONS
# ===============================================================
def update_ui_from_files(chain_files, ruleset_file, settings_file):
    """
    ERROR_FIX: ENHANCED_FILE_PROCESSING + REWARD_EXTRACTION
    """
    chain_data = []
    if chain_files:
        for file in chain_files:
            chain = load_chain_config(file)
            for item in chain.Items:
                chain_data.append([chain.Id, item.MergeItemId, item.RequirementWeight, item.RewardDifficulty])
    
    chain_df = pd.DataFrame(chain_data, columns=['ChainId', 'MergeItemId', 'RequirementWeight', 'RewardDifficulty'])

    max_hist, inc_diff, energy_chance, req_weights = 5, 2, 90, "70,30"
    energy_df_data, item_df_data = [], []
    
    if ruleset_file:
        # ERROR_FIX: Enhanced ruleset parsing with validation
        ruleset = load_ruleset_config(ruleset_file)
        max_hist, inc_diff, energy_chance = ruleset.MaxHistoryOrders, ruleset.IncrementDifficulty, ruleset.OverallChanceToDropExpeditionEnergy
        req_weights = ",".join(map(str, ruleset.OverrideMaxOrdersWithWeight.Weights))
        
        # ERROR_FIX: Extract rewards data with NaN validation
        for rg in ruleset.EnergyRewards:
            if pd.notna(rg.DifficultyScore):
                for r in rg.Rewards:
                    energy_df_data.append([rg.DifficultyScore, r.Amount])
        for rg in ruleset.ItemRewards:
            if pd.notna(rg.DifficultyScore):
                for r in rg.Rewards:
                    item_df_data.append([rg.DifficultyScore, r.Amount, r.MergeItemId, r.RewardWeight, r.ReductionFactor])
        
        # DEBUG_LOG: Reward extraction statistics
        print(f"LOADED: {len(energy_df_data)} energy_rewards, {len(item_df_data)} item_rewards")
        
    energy_df = pd.DataFrame(energy_df_data, columns=['DifficultyScore', 'Amount'])
    item_df = pd.DataFrame(item_df_data, columns=['DifficultyScore', 'Amount', 'MergeItemId', 'RewardWeight', 'ReductionFactor'])
        
    red_factor, inc_factor = 3, 5
    if settings_file:
        settings = load_settings_config(settings_file)
        red_factor, inc_factor = settings.ReductionFactor, settings.IncreaseFactor
        if not ruleset_file:
            req_weights = ",".join(map(str, settings.DefaultRequirementWeights.Weights))

    return chain_df, max_hist, inc_diff, energy_chance, req_weights, red_factor, inc_factor, energy_df, item_df

# ===============================================================
# 5. GRADIO UI
# ===============================================================
with gr.Blocks(theme=gr.themes.Soft(), title="Инструмент Балансировки Merge-2") as demo:
    gr.Markdown("# Инструмент для Балансировки Игр Merge-2")
    with gr.Tabs():
        with gr.TabItem("Симуляция генератора заказов"):
            with gr.Row():
                with gr.Column(scale=2):
                    gr.Markdown("### 1. Загрузите файлы (опционально)")
                    chain_files_upload = gr.File(label="Файлы цепочек (.asset)", file_count="multiple")
                    ruleset_upload = gr.File(label="Ruleset (.asset)")
                    settings_upload = gr.File(label="Settings (.asset)")
                    
                    with gr.Accordion("Редактор Мердж-цепочек", open=False):
                        chain_editor_df = gr.DataFrame(headers=['ChainId', 'MergeItemId', 'RequirementWeight', 'RewardDifficulty'], datatype=['str', 'str', 'number', 'number'], label="Состав цепочек", interactive=True, row_count=(10, "dynamic"))

                    with gr.Accordion("Настройки генератора", open=True):
                        with gr.Row():
                            max_history_input = gr.Slider(1, 20, value=5, step=1, label="Макс. заказов в истории")
                            increment_diff_input = gr.Slider(0, 10, value=2, step=1, label="Инкремент сложности")
                        with gr.Row():
                            energy_chance_input = gr.Slider(0, 100, value=90, step=5, label="Шанс награды-энергии (%)")
                            req_weights_input = gr.Textbox(label="Веса требований (1, 2)", value="70, 30")
                        with gr.Row():
                             reduction_factor_input = gr.Number(label="Reduction Factor", value=3)
                             increase_factor_input = gr.Number(label="Increase Factor", value=5)
                    
                    with gr.Accordion("Редактор Наград", open=False):
                        with gr.Row():
                            energy_rewards_df = gr.DataFrame(headers=['DifficultyScore', 'Amount'], datatype=['number', 'number'], label="Энергетические награды", interactive=True, col_count=(2, "fixed"), row_count=(5, "dynamic"))
                            item_rewards_df = gr.DataFrame(headers=['DifficultyScore', 'Amount', 'MergeItemId', 'RewardWeight', 'ReductionFactor'], datatype=['number', 'number', 'str', 'number', 'number'], label="Предметные награды", interactive=True, col_count=(5, "fixed"), row_count=(5, "dynamic"))

                    gr.Markdown("### 2. Запустите симуляцию")
                    sim_iterations = gr.Slider(10, 1000, value=100, step=10, label="Количество итераций")
                    sim_initial_energy_input = gr.Number(value=10000, label="Начальное количество энергии")
                    sim_run_button = gr.Button("Запустить симуляцию", variant="primary")

                with gr.Column(scale=3):
                    gr.Markdown("### Результаты симуляции")
                    sim_results_df = gr.DataFrame(label="Данные по заказам", wrap=True)
                    gr.Markdown("### Сводный отчет")
                    sim_stats_report = gr.Textbox(label="Статистика", lines=10)
                    with gr.Row():
                         sim_download_csv = gr.File(label="Скачать CSV", visible=False, interactive=False)
                         sim_download_report = gr.File(label="Скачать отчет", visible=False, interactive=False)

    # --- Event Handlers ---
    files_to_update_ui = [chain_files_upload, ruleset_upload, settings_upload]
    ui_outputs_to_update = [chain_editor_df, max_history_input, increment_diff_input, energy_chance_input, req_weights_input, reduction_factor_input, increase_factor_input, energy_rewards_df, item_rewards_df]
    
    for file_input in files_to_update_ui:
        file_input.upload(update_ui_from_files, files_to_update_ui, ui_outputs_to_update)

    sim_run_button.click(
        fn=run_simulation_interface,
        inputs=[chain_editor_df, energy_rewards_df, item_rewards_df, max_history_input, increment_diff_input, energy_chance_input, req_weights_input, reduction_factor_input, increase_factor_input, sim_iterations, sim_initial_energy_input],
        outputs=[sim_results_df, sim_stats_report, sim_download_csv, sim_download_report]
    )

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