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 services.llamaswap import read_config, write_config
from config import CONFIG_PATH, SPEC_DRAFT_MODEL_PATH

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", {})

    draft_path = SPEC_DRAFT_MODEL_PATH

    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. Ensure prompt caching flags exist
        if "--prompt-cache " not in cmd and not cmd.endswith("--prompt-cache") and not cmd.endswith("--prompt-cache\n"):
            cmd = cmd.strip() + " --prompt-cache"
        if "--prompt-cache-all" not in cmd:
            cmd = cmd.strip() + " --prompt-cache-all"
            
        # 2. Extract aliases/role
        aliases = spec.get("aliases", [])
        role = aliases[0] if aliases else None
        
        # 3. Add parallel and speculative decoding for fast and coder
        if role in ("fast", "coder"):
            if "--parallel" not in cmd:
                cmd = cmd.strip() + " --parallel 2"
            if "--spec-draft-model" not in cmd:
                cmd = cmd.strip() + f" --spec-draft-model {draft_path}"
                
        # 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()