"""Patch backend: replace AVAILABLE_MODELS with OpenRouter models + routing."""

import ast
import os
import re

AGENT_FILE = "/app/backend/routes/agent.py"
LLM_PARAMS_FILE = "/app/agent/core/llm_params.py"

# === Step 1: Patch model_ids.py ===
IDS_FILE = "/app/agent/core/model_ids.py"
with open(IDS_FILE, "r") as f:
    content = f.read()

old_models_block = 'CLAUDE_OPUS_48_MODEL_ID = "anthropic/claude-opus-4.8:fal-ai"\nGPT_55_MODEL_ID = "openai/gpt-5.5:fal-ai"\nKIMI_K27_CODE_MODEL_ID = "moonshotai/Kimi-K2.7-Code:novita"\nMINIMAX_M3_MODEL_ID = "MiniMaxAI/MiniMax-M3:novita"\nGLM_52_MODEL_ID = "zai-org/GLM-5.2:novita"\nDEEPSEEK_V4_PRO_MODEL_ID = "deepseek-ai/DeepSeek-V4-Pro:novita"'

new_models_block = 'GPT_55_MODEL_ID = "openai/gpt-5.5:fal-ai"\nKIMI_K27_CODE_MODEL_ID = "deepseek-ai/DeepSeek-V4-Pro"\nMINIMAX_M3_MODEL_ID = "deepseek-ai/DeepSeek-V4-Flash"\nGLM_52_MODEL_ID = "openai/deepseek/deepseek-v4-flash"\nDEEPSEEK_V4_PRO_MODEL_ID = "nvidia/nemotron-3-super-120b-a12b:free"\n\n# === PATCH: OpenRouter models ===\nGEMMA_4_31B_FREE_MODEL_ID = "openai/google/gemma-4-31b-it:free"\nTENCENT_HY3_FREE_MODEL_ID = "openai/tencent/hy3:free"\nLLAMA_3_3_70B_FREE_MODEL_ID = "openai/meta-llama/llama-3.3-70b-instruct:free"\nLLAMA_3_1_8B_MODEL_ID = "openai/meta-llama/llama-3.1-8b-instruct"\nLAGUNA_M1_FREE_MODEL_ID = "openai/poolside/laguna-m.1:free"\nLAGUNA_S21_FREE_MODEL_ID = "openai/poolside/laguna-s-2.1:free"'

content = content.replace(old_models_block, new_models_block)

old_hosted = "HOSTED_MODEL_IDS = {\n    CLAUDE_OPUS_48_MODEL_ID,\n    GPT_55_MODEL_ID,\n    KIMI_K27_CODE_MODEL_ID,\n    MINIMAX_M3_MODEL_ID,\n    GLM_52_MODEL_ID,\n    DEEPSEEK_V4_PRO_MODEL_ID,\n}"
new_hosted = "HOSTED_MODEL_IDS = {\n    GPT_55_MODEL_ID,\n    KIMI_K27_CODE_MODEL_ID,\n    MINIMAX_M3_MODEL_ID,\n    GLM_52_MODEL_ID,\n    DEEPSEEK_V4_PRO_MODEL_ID,\n    GEMMA_4_31B_FREE_MODEL_ID,\n    TENCENT_HY3_FREE_MODEL_ID,\n    LLAMA_3_3_70B_FREE_MODEL_ID,\n    LLAMA_3_1_8B_MODEL_ID,\n    LAGUNA_M1_FREE_MODEL_ID,\n    LAGUNA_S21_FREE_MODEL_ID,\n}"
content = content.replace(old_hosted, new_hosted)

with open(IDS_FILE, "w") as f:
    f.write(content)
print("OK: model_ids.py patched")

# === Step 2: Patch agent.py — replace _available_models() via regex ===
with open(AGENT_FILE) as f:
    content = f.read()

# Make DEFAULT_MODEL_ID and DEFAULT_GPT_MODEL_ID point to our models
content = content.replace(
    "DEFAULT_MODEL_ID = GLM_52_MODEL_ID",
    'DEFAULT_MODEL_ID = "openai/tencent/hy3:free"'
)
content = content.replace(
    "DEFAULT_GPT_MODEL_ID = GPT_55_MODEL_ID",
    'DEFAULT_GPT_MODEL_ID = LAGUNA_S21_FREE_MODEL_ID'
)

# Update imports
old_import = "from agent.core.model_ids import (\n    CLAUDE_OPUS_48_MODEL_ID,\n    DEEPSEEK_V4_PRO_MODEL_ID,\n    GLM_52_MODEL_ID,\n    GPT_55_MODEL_ID,\n    KIMI_K27_CODE_MODEL_ID,\n    MINIMAX_M3_MODEL_ID,\n    strip_huggingface_model_prefix,\n)"
new_import = "from agent.core.model_ids import (\n    DEEPSEEK_V4_PRO_MODEL_ID,\n    GLM_52_MODEL_ID,\n    GPT_55_MODEL_ID,\n    KIMI_K27_CODE_MODEL_ID,\n    MINIMAX_M3_MODEL_ID,\n    GEMMA_4_31B_FREE_MODEL_ID,\n    TENCENT_HY3_FREE_MODEL_ID,\n    LLAMA_3_3_70B_FREE_MODEL_ID,\n    LLAMA_3_1_8B_MODEL_ID,\n    LAGUNA_M1_FREE_MODEL_ID,\n    LAGUNA_S21_FREE_MODEL_ID,\n    strip_huggingface_model_prefix,\n)"
content = content.replace(old_import, new_import)

# Replace _available_models() using regex
func_pattern = re.compile(
    r'def _available_models\(\) -> list\[dict\[str, Any\]\]:\s*\n'
    r'\s+models\s*=\s*\[.*?\]\s*\n\s+return models',
    re.DOTALL
)

new_func = '''def _available_models() -> list[dict[str, Any]]:
    models = [
        {
            "id": TENCENT_HY3_FREE_MODEL_ID,
            "label": "Tencent HY3:free",
            "recommended": True,
        },
        {
            "id": GEMMA_4_31B_FREE_MODEL_ID,
            "label": "Gemma 4 31B:free",
            "recommended": True,
        },
        {
            "id": LLAMA_3_3_70B_FREE_MODEL_ID,
            "label": "Llama 3.3 70B:free",
            "recommended": True,
        },
        {
            "id": LAGUNA_M1_FREE_MODEL_ID,
            "label": "Laguna M.1:free",
            "recommended": True,
        },
        {
            "id": DEFAULT_GPT_MODEL_ID,
            "label": "Laguna S 2.1:free",
            "recommended": True,
        },
        {
            "id": LLAMA_3_1_8B_MODEL_ID,
            "label": "Llama 3.1 8B",
        },
        {
            "id": KIMI_K27_CODE_MODEL_ID,
            "label": "DeepSeek V4 Pro",
        },
        {
            "id": MINIMAX_M3_MODEL_ID,
            "label": "DeepSeek V4 Flash",
        },
        {
            "id": DEFAULT_MODEL_ID,
            "label": "DeepSeek V4 Flash",
        },
        {
            "id": DEEPSEEK_V4_PRO_MODEL_ID,
            "label": "Nemotron 3 Super 120B",
        },
    ]
    return models'''

func_match = func_pattern.search(content)
if func_match:
    content = content[:func_match.start()] + new_func + content[func_match.end():]
    print("OK: Replaced _available_models() via regex")
else:
    print("WARN: Regex failed, trying string replace...")
    old_available = 'def _available_models() -> list[dict[str, Any]]:\n    models = [\n        {\n            "id": CLAUDE_OPUS_48_MODEL_ID,\n            "label": "Claude Opus 4.8",\n        },\n        {\n            "id": DEFAULT_GPT_MODEL_ID,\n            "label": "GPT-5.5",\n        },\n        {\n            "id": KIMI_K27_CODE_MODEL_ID,\n            "label": "Kimi K2.7 Code",\n        },\n        {\n            "id": MINIMAX_M3_MODEL_ID,\n            "label": "MiniMax M3",\n        },\n        {\n            "id": DEFAULT_MODEL_ID,\n            "label": "GLM 5.2",\n            "recommended": True,\n        },\n        {\n            "id": DEEPSEEK_V4_PRO_MODEL_ID,\n            "label": "DeepSeek V4 Pro",\n        },\n    ]\n    return models'
    if old_available in content:
        content = content.replace(old_available, new_func)
        print("OK: String replace worked")
    else:
        print("FAIL: Cannot find _available_models()!")
        import sys
        sys.exit(1)

# Update title generation
old_title = '"openai/gpt-oss-120b:cerebras",'
new_title = '"huggingface/deepseek-ai/DeepSeek-V4-Pro",'
content = content.replace(old_title, new_title)

try:
    ast.parse(content)
    print("OK: agent.py syntax OK")
except SyntaxError as e:
    print(f"FAIL: agent.py syntax error: {e}")
    raise

with open(AGENT_FILE, "w") as f:
    f.write(content)
print("OK: agent.py patched")

# === Step 3: Patch _resolve_llm_params for OpenRouter routing ===
with open(LLM_PARAMS_FILE) as f:
    llm_content = f.read()

if "normalized_model.startswith" not in llm_content:
    api_key_find = "api_key = _resolve_hf_router_token(session_hf_token)"
    if api_key_find in llm_content:
        line_end = llm_content.find('\n', llm_content.find(api_key_find) + len(api_key_find)) + 1
        routing_insert = '''    # === PATCH: Route ALL openai/-prefixed models to OpenRouter ===
    if normalized_model.startswith("openai/"):
        return {
            "model": normalized_model,
            "api_base": "https://openrouter.ai/api/v1",
            "api_key": os.environ.get("OPENROUTER_API_KEY") or api_key or "",
        }
    # === END PATCH ===

'''
        llm_content = llm_content[:line_end] + routing_insert + llm_content[line_end:]
        print("OK: Patched _resolve_llm_params (catch-all openai/ -> OR)")
    else:
        print("WARN: Could not find api_key line")
else:
    print("OK: llm_params.py already patched")

try:
    ast.parse(llm_content)
    print("OK: llm_params.py syntax OK")
except SyntaxError as e:
    print(f"FAIL: llm_params.py syntax error: {e}")
    raise

with open(LLM_PARAMS_FILE, "w") as f:
    f.write(llm_content)

print("OK: ALL PATCHES APPLIED")