import os, re

path = 'client/cocos/assets/scripts/combat/CombatController.ts'
with open(path, 'r', encoding='utf-8') as f:
    content = f.read()

# Replace skills record with loadouts
content = re.sub(
    r'  // Skills\n  public skills: Record<string, SkillSlot> = \{.*?\n  \};\n',
    '''  // Skills Loadouts (Chân Võ Đan + Tà Ấn)
  public loadouts: Record<string, SkillLoadout> = {
    'Q': { baseSkillId: 1001, supportIds: [2001, 2002], currentCooldown: 0 },
    'W': { baseSkillId: 1007, supportIds: [2005], currentCooldown: 0 },
    'E': { baseSkillId: 1011, supportIds: [2004], currentCooldown: 0 }
  };\n''',
    content, flags=re.DOTALL
)

# Fix Object.keys(this.skills) -> loadouts
content = content.replace('Object.keys(this.skills)', 'Object.keys(this.loadouts)')
content = content.replace('const s = this.skills[key];', 'const s = this.loadouts[key];')

new_method = '''
  public resolveSkillContext(loadout: SkillLoadout): SkillExecutionCtx {
    const baseSkill = CANONICAL_SKILLS[loadout.baseSkillId];
    if (!baseSkill) throw new Error("Skill not found");

    const ctx: SkillExecutionCtx = {
      skillId: loadout.baseSkillId,
      projectileCount: 1, // Default or pull from baseSkill if it has one
      animationSpeedMulti: 1.0,
      damageMulti: 1.0,
      aoeRadiusMulti: 1.0,
      costMulti: 1.0,
      vfxOverrideColor: null
    };

    // Apply Support Skills (Tà Ấn) Modifiers
    for (const sId of loadout.supportIds) {
      const support = SUPPORT_SKILLS[sId];
      if (!support) continue;

      switch (support.augmentType) {
        case 'multi_projectile':
          ctx.projectileCount += 3;
          ctx.damageMulti *= 0.65;
          break;
        case 'aoe_expansion':
          ctx.aoeRadiusMulti *= 1.45;
          ctx.damageMulti *= 0.85;
          break;
        case 'rapid_cast':
          ctx.animationSpeedMulti *= 1.45;
          ctx.damageMulti *= 0.75;
          break;
        case 'blood_magic':
          ctx.vfxOverrideColor = '#ef4444';
          break;
      }
    }

    return ctx;
  }
'''

content = re.sub(
    r'  public castSkill.*?return true;\n  }',
    new_method + '''
  public castSkill(key: 'Q' | 'W' | 'E'): boolean {
    const loadout = this.loadouts[key];
    if (!loadout || loadout.currentCooldown > 0 || this.isIFrame) {
      return false;
    }

    const ctx = this.resolveSkillContext(loadout);
    const baseSkill = CANONICAL_SKILLS[loadout.baseSkillId];

    loadout.currentCooldown = baseSkill?.cooldownSec || 3.0;
    const baseDamage = 110;
    const damage = baseDamage * ctx.damageMulti;
    const isCrit = Math.random() < 0.35;

    if (this.animFSM) {
      // Pass animation speed multiplier to animation engine
      this.animFSM.setState('attack', true, { playbackRate: ctx.animationSpeedMulti });
    }
    CombatFeel.triggerImpactShake(this.camera, true, true);
    CombatFeel.triggerHitStop(this.animFSM, 50);

    EventBus.emit('skillCasted', { skillId: ctx.skillId, damage, isCrit, ctx });
    return true;
  }''',
    content, flags=re.DOTALL
)

with open(path, 'w', encoding='utf-8') as f:
    f.write(content)

print('Updated CombatController with Loadouts and resolveSkillContext.')