import json
import copy
from typing import Dict, List, Tuple, Optional

class ArtifactOptimizer:
    def __init__(self):
        # Special numbers for scoring (from the post)
        self.stat_weights = {
            'hp_': 5.0,      # HP%
            'hp': 700.0,     # Flat HP
            'atk_': 5.0,     # ATK%
            'atk': 45.0,     # Flat ATK
            'def_': 6.0,     # DEF%
            'def': 50.0,     # Flat DEF
            'eleMas': 20.0,  # Elemental Mastery
            'enerRech_': 5.5, # Energy Recharge%
            'critRate_': 3.0, # CRIT Rate%
            'critDMG_': 6.0   # CRIT DMG%
        }

        # Slot mappings
        self.slot_names = {
            'flower': 'Flower',
            'plume': 'Feather',
            'sands': 'Sands',
            'goblet': 'Goblet',
            'circlet': 'Circlet'
        }

    def load_data(self, artifacts_file: str, characters_file: str):
        """Load artifact and character data from JSON files"""
        with open(artifacts_file, 'r') as f:
            self.artifact_data = json.load(f)

        with open(characters_file, 'r') as f:
            self.character_data = json.load(f)

        self.artifacts = self.artifact_data['artifacts']

    def get_character_relevant_stats(self, character: str) -> List[str]:
        """Get the stats that are relevant for a character based on their substat priority"""
        char_data = self.character_data.get(character, {})
        substat_priority = char_data.get('substat_priority', [])

        # Convert display names to internal stat keys
        stat_mapping = {
            'ER%': 'enerRech_',
            'CRIT Rate': 'critRate_',
            'CRIT DMG': 'critDMG_',
            'HP%': 'hp_',
            'Flat HP': 'hp',
            'ATK%': 'atk_',
            'Flat ATK': 'atk',
            'DEF%': 'def_',
            'Flat DEF': 'def',
            'Elemental Mastery': 'eleMas'
        }

        relevant_stats = []
        for stat in substat_priority:
            if stat in stat_mapping:
                relevant_stats.append(stat_mapping[stat])

        return relevant_stats

    def calculate_artifact_score(self, artifact: Dict, character: str) -> float:
        """Calculate artifact score for a specific character"""
        relevant_stats = self.get_character_relevant_stats(character)

        total_value = 0.0

        # Score substats
        for substat in artifact.get('substats', []):
            stat_key = substat['key']
            stat_value = substat['value']

            if stat_key in relevant_stats and stat_key in self.stat_weights:
                total_value += stat_value / self.stat_weights[stat_key]

        # Convert to percentage score (max 9 substat rolls possible)
        score = (total_value / 9.0) * 100.0
        return score

    def is_set_suitable(self, artifact: Dict, character: str) -> bool:
        """Check if artifact's set is suitable for the character"""
        char_data = self.character_data.get(character, {})
        required_sets = char_data.get('required_sets', [])

        # If no required sets specified, accept any set
        if not required_sets:
            return True

        return artifact['setKey'] in required_sets

    def is_main_stat_suitable(self, artifact: Dict, character: str) -> bool:
        """Check if artifact's main stat is suitable for the character"""
        char_data = self.character_data.get(character, {})
        ideal_main_stats = char_data.get('ideal_main_stats', {})

        slot = artifact['slotKey']
        main_stat = artifact['mainStatKey']

        # Convert internal stat keys to display names for comparison
        stat_display_mapping = {
            'hp': 'HP',
            'atk': 'ATK',
            'hp_': 'HP%',
            'atk_': 'ATK%',
            'def_': 'DEF%',
            'enerRech_': 'ER%',
            'eleMas': 'Elemental Mastery',
            'critRate_': 'CRIT Rate',
            'critDMG_': 'CRIT DMG',
            'heal_': 'Healing Bonus',
            'physical_dmg_': 'Physical DMG Bonus',
            'anemo_dmg_': 'Anemo DMG Bonus',
            'geo_dmg_': 'Geo DMG Bonus',
            'electro_dmg_': 'Electro DMG Bonus',
            'dendro_dmg_': 'Dendro DMG Bonus',
            'hydro_dmg_': 'Hydro DMG Bonus',
            'pyro_dmg_': 'Pyro DMG Bonus',
            'cryo_dmg_': 'Cryo DMG Bonus'
        }

        main_stat_display = stat_display_mapping.get(main_stat, main_stat)

        # Flower and Feather have fixed main stats
        if slot == 'flower':
            return main_stat == 'hp'
        elif slot == 'plume':
            return main_stat == 'atk'

        # Check if main stat is in the ideal list for this slot
        slot_ideals = ideal_main_stats.get(slot, [])

        # If no ideal main stats specified for this slot, accept any reasonable main stat
        if not slot_ideals:
            # Accept common main stats for each slot
            if slot == 'sands':
                return main_stat in ['atk_', 'hp_', 'def_', 'enerRech_', 'eleMas']
            elif slot == 'goblet':
                return main_stat in ['atk_', 'hp_', 'def_', 'eleMas'] or main_stat.endswith('_dmg_')
            elif slot == 'circlet':
                return main_stat in ['atk_', 'hp_', 'def_', 'critRate_', 'critDMG_', 'heal_', 'eleMas']

        return main_stat_display in slot_ideals

    def check_er_requirement(self, artifacts: List[Dict], character: str) -> Dict:
        """Check if the artifact set meets ER requirements"""
        char_data = self.character_data.get(character, {})
        er_req = char_data.get('er_requirement', {})

        # Calculate total ER from artifacts
        total_er = 100.0  # Base ER

        for artifact in artifacts:
            # Main stat ER
            if artifact['mainStatKey'] == 'enerRech_':
                if artifact['level'] == 20:
                    total_er += 51.8  # Max ER% main stat at level 20
                else:
                    # Approximate ER based on level
                    total_er += 51.8 * (artifact['level'] / 20.0)

            # Substat ER
            for substat in artifact.get('substats', []):
                if substat['key'] == 'enerRech_':
                    total_er += substat['value']

        # Determine requirement type
        if isinstance(er_req, dict):
            if 'min' in er_req:
                # Simple min/max requirement
                req_min = er_req['min']
                req_max = er_req.get('max', req_min)
                meets_req = req_min <= total_er <= req_max + 20  # Allow some flexibility
                req_type = f"{req_min}-{req_max}%"
            else:
                # Multiple scenarios (like Fischl)
                meets_req = False
                req_type = "Various scenarios"
                for scenario, req_data in er_req.items():
                    if req_data['min'] <= total_er <= req_data['max'] + 20:
                        meets_req = True
                        req_type = f"{scenario}: {req_data['min']}-{req_data['max']}%"
                        break
        else:
            meets_req = True
            req_type = "No specific requirement"

        return {
            'total_er': total_er,
            'meets_requirement': meets_req,
            'requirement_type': req_type
        }

    def get_available_artifacts_by_slot(self, slot: str, exclude_ids: set, character: str = None, min_rarity: int = 5) -> List[Dict]:
        """Get all available artifacts for a specific slot, optionally filtered by character requirements"""
        available = []
        for artifact in self.artifacts:
            if (artifact['slotKey'] == slot and
                artifact['id'] not in exclude_ids and
                artifact['rarity'] >= min_rarity):

                # If character specified, check set requirements
                if character and not self.is_set_suitable(artifact, character):
                    continue

                available.append(artifact)
        return available

    def find_best_build_for_character(self, character: str, exclude_ids: set) -> Tuple[List[Dict], float, Dict, Dict]:
        """Find the best artifact build for a character"""
        slots = ['flower', 'plume', 'sands', 'goblet', 'circlet']
        best_build = None
        best_score = -1
        best_er_info = None

        # Get available artifacts for each slot and collect debug info
        slot_artifacts = {}
        debug_info = {}

        # First try with 5-star artifacts only
        for slot in slots:
            available_5star = self.get_available_artifacts_by_slot(slot, exclude_ids, character, min_rarity=5)
            suitable_5star = [art for art in available_5star if self.is_main_stat_suitable(art, character)]

            # If no suitable 5-star artifacts, include 4-star as fallback
            if len(suitable_5star) == 0:
                available_all = self.get_available_artifacts_by_slot(slot, exclude_ids, character, min_rarity=4)
                suitable_all = [art for art in available_all if self.is_main_stat_suitable(art, character)]
                slot_artifacts[slot] = suitable_all
                used_fallback = len(suitable_all) > 0
            else:
                slot_artifacts[slot] = suitable_5star
                available_all = available_5star
                used_fallback = False

            # Collect debug information
            char_data = self.character_data.get(character, {})
            required_sets = char_data.get('required_sets', [])

            main_stats_found = list(set(art['mainStatKey'] for art in available_all))
            suitable_main_stats_found = list(set(art['mainStatKey'] for art in slot_artifacts[slot]))
            sets_found = list(set(art['setKey'] for art in available_all))

            debug_info[slot] = {
                'total_available': len(available_all),
                'suitable_main_stats': len(slot_artifacts[slot]),
                'main_stats_found': main_stats_found,
                'suitable_main_stats_found': suitable_main_stats_found,
                'sets_found': sets_found,
                'required_sets': required_sets,
                'used_4star_fallback': used_fallback
            }

        # Try all combinations (this is computationally expensive but thorough)
        from itertools import product

        artifact_combinations = list(product(*[slot_artifacts[slot] for slot in slots]))

        for combination in artifact_combinations:
            if len(set(art['id'] for art in combination)) != 5:
                continue  # Skip if any artifacts are duplicated

            # Calculate total score
            total_score = sum(self.calculate_artifact_score(art, character) for art in combination)
            avg_score = total_score / 5

            # Check ER requirement
            er_info = self.check_er_requirement(list(combination), character)

            # Prioritize builds that meet ER requirements
            adjusted_score = avg_score
            if er_info['meets_requirement']:
                adjusted_score += 10  # Bonus for meeting ER req

            if adjusted_score > best_score:
                best_score = adjusted_score
                best_build = list(combination)
                best_er_info = er_info

        return best_build, best_score - (10 if best_er_info and best_er_info['meets_requirement'] else 0), best_er_info, debug_info

    def optimize_builds(self, character_priority: List[str]) -> Dict:
        """Optimize artifact builds for characters in priority order"""
        results = {}
        used_artifact_ids = set()

        for character in character_priority:
            if character not in self.character_data:
                print(f"Warning: Character '{character}' not found in character data")
                continue

            print(f"Finding best build for {character}...")

            best_build, score, er_info, debug_info = self.find_best_build_for_character(character, used_artifact_ids)

            if best_build:
                # Mark these artifacts as used
                for artifact in best_build:
                    used_artifact_ids.add(artifact['id'])

                # Calculate individual artifact scores
                artifact_scores = []
                for artifact in best_build:
                    art_score = self.calculate_artifact_score(artifact, character)
                    artifact_scores.append({
                        'artifact': artifact,
                        'score': art_score
                    })

                results[character] = {
                    'build': best_build,
                    'average_score': score,
                    'artifact_scores': artifact_scores,
                    'er_info': er_info,
                    'debug_info': debug_info
                }
            else:
                results[character] = {
                    'build': None,
                    'average_score': 0,
                    'artifact_scores': [],
                    'er_info': {'total_er': 100, 'meets_requirement': False, 'requirement_type': 'No artifacts available'},
                    'debug_info': debug_info
                }

        return results

    def format_results(self, results: Dict) -> str:
        """Format the optimization results for display"""
        output = []
        output.append("=" * 80)
        output.append("ARTIFACT OPTIMIZATION RESULTS")
        output.append("=" * 80)

        for character, data in results.items():
            output.append(f"\n{'='*20} {character.upper()} {'='*20}")

            if not data['build']:
                output.append("❌ No suitable build found")

                # Add debug information
                debug_info = data.get('debug_info', {})
                output.append("\n🔍 DEBUG INFO:")
                for slot, info in debug_info.items():
                    fallback_note = " (using 4⭐ fallback)" if info.get('used_4star_fallback') else ""
                    output.append(f"  {slot}: {info['suitable_main_stats']}/{info['total_available']} suitable artifacts{fallback_note}")
                    if info.get('required_sets'):
                        output.append(f"    Required sets: {', '.join(info['required_sets'])}")
                    if info.get('sets_found'):
                        output.append(f"    Available sets: {', '.join(info['sets_found'])}")
                    if info['main_stats_found']:
                        output.append(f"    Available main stats: {', '.join(info['main_stats_found'])}")
                    if info['suitable_main_stats_found']:
                        output.append(f"    Suitable main stats: {', '.join(info['suitable_main_stats_found'])}")

                continue

            output.append(f"📊 Average Artifact Score: {data['average_score']:.2f}/100")

            # ER Information
            er_info = data['er_info']
            er_status = "✅" if er_info['meets_requirement'] else "⚠️"
            output.append(f"{er_status} Energy Recharge: {er_info['total_er']:.1f}% ({er_info['requirement_type']})")

            output.append("\n📋 ARTIFACT BUILD:")

            for i, art_data in enumerate(data['artifact_scores']):
                artifact = art_data['artifact']
                score = art_data['score']

                slot_name = self.slot_names[artifact['slotKey']]
                set_name = artifact['setKey']
                main_stat = artifact['mainStatKey']
                level = artifact['level']
                rarity = artifact['rarity']
                rarity_stars = '⭐' * rarity

                output.append(f"\n{slot_name} ({set_name}) - Level {level} {rarity_stars}")
                output.append(f"  Main Stat: {main_stat}")
                output.append(f"  Score: {score:.2f}/100")
                output.append(f"  Substats:")

                for substat in artifact['substats']:
                    stat_key = substat['key']
                    stat_value = substat['value']

                    # Format the value appropriately
                    if stat_key in ['hp', 'atk', 'def', 'eleMas']:
                        formatted_value = f"{stat_value:.0f}"
                    else:
                        formatted_value = f"{stat_value:.1f}%"

                    output.append(f"    • {stat_key}: {formatted_value}")

        # Summary
        output.append(f"\n{'='*20} SUMMARY {'='*20}")
        total_chars = len(results)
        successful_builds = sum(1 for data in results.values() if data['build'] is not None)

        output.append(f"Characters processed: {total_chars}")
        output.append(f"Successful builds: {successful_builds}")
        output.append(f"Failed builds: {total_chars - successful_builds}")

        if successful_builds > 0:
            avg_score = sum(data['average_score'] for data in results.values() if data['build']) / successful_builds
            output.append(f"Average build quality: {avg_score:.2f}/100")

        return "\n".join(output)

def main():
    # Initialize optimizer
    optimizer = ArtifactOptimizer()

    # Load data
    optimizer.load_data('data.json', 'characters.json')

    # Define character priority (example)
    character_priority = ['Furina', 'Escoffier', 'Fischl', 'Chiori']

    print("Starting artifact optimization...")
    print(f"Character priority: {' -> '.join(character_priority)}")
    print(f"Total artifacts available: {len(optimizer.artifacts)}")

    # Optimize builds
    results = optimizer.optimize_builds(character_priority)

    # Display results
    formatted_output = optimizer.format_results(results)
    print(formatted_output)

    # Save results to file
    with open('optimization_results.txt', 'w') as f:
        f.write(formatted_output)

    print(f"\nResults saved to 'optimization_results.txt'")

if __name__ == "__main__":
    main()