import sys
from pathlib import Path

# Add backend directory to sys.path so we can import services
sys.path.append(str(Path(__file__).resolve().parent))

from config import CONFIG_PATH
from services.llamaswap import _PATH_RE, read_config, spec_draft_flags, write_config


def migrate():
    print(f"Reading config from {CONFIG_PATH}...")
    if not CONFIG_PATH.exists():
        print(f"Config path {CONFIG_PATH} does not exist. Skipping.")
        return
        
    cfg = read_config()
    models = cfg.get("models", {})

    for name, spec in models.items():
        if not isinstance(spec, dict):
            continue
        cmd = spec.get("cmd", "")
        if not cmd:
            continue
            
        print(f"Migrating model: {name}")

        # 1. Defektes --prompt-cache/--prompt-cache-all entfernen (llama-CLI-Flags,
        #    die llama-server ablehnt → Start scheitert). Caching macht llama-server
        #    automatisch pro Slot.
        cmd = cmd.replace(" --prompt-cache-all", "").replace(" --prompt-cache", "")

        # 2. Extract aliases/role
        aliases = spec.get("aliases", [])
        role = aliases[0] if aliases else None
        
        # 3. Add parallel + (nur vocab-kompatibles) Speculative Decoding für fast/coder.
        #    spec_draft_flags() prüft die Vocab-Kompatibilität und hängt --spec-type an;
        #    ein inkompatibler Draft (z.B. qwen2.5 ↔ Qwen3.6) wird NICHT gesetzt.
        if role in ("fast", "coder"):
            if "--parallel" not in cmd:
                cmd = cmd.strip() + " --parallel 2"
            if "--spec-draft-model" not in cmd:
                target = mt.group(1) if (mt := _PATH_RE.search(cmd)) else ""
                cmd = cmd.strip() + spec_draft_flags(target)
                
        # Update cmd
        from ruamel.yaml.scalarstring import LiteralScalarString
        spec["cmd"] = LiteralScalarString(cmd.strip() + "\n")
        
    print(f"Writing updated config back to {CONFIG_PATH}...")
    write_config(cfg)
    print("Migration completed successfully!")

if __name__ == "__main__":
    migrate()