import os
import requests

OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")


TASK_MODELS = {
    "agent": [
        "anthropic/claude-sonnet-4.5",
        "openai/gpt-4o",
        "deepseek/deepseek-chat"
    ],
    "code_hard": [
        "anthropic/claude-sonnet-4.5",
        "openai/gpt-4.1",
        "qwen/qwen3-coder"
    ],
    "code_cheap": [
        "qwen/qwen3-coder",
        "deepseek/deepseek-chat"
    ],
    "chinese": [
        "deepseek/deepseek-chat",
        "qwen/qwen-2.5-72b-instruct"
    ],
    "cheap": [
        "google/gemini-flash-1.5",
        "mistralai/mistral-small"
    ]
}

def classify_task(prompt: str) -> str:
    p = prompt.lower()

    if any(x in p for x in ["debug", "error", "stack trace", "fix", "exception"]):
        return "code_hard"

    if any(x in p for x in ["code", "python", "node", "swift", "typescript"]):
        return "code_cheap"

    if any(x in p for x in ["计划", "总结", "中文", "翻译", "写一段"]):
        return "chinese"

    if any(x in p for x in ["agent", "tool", "automation", "workflow", "n8n"]):
        return "agent"

    return "cheap"


def choose_models(prompt: str):
    task = classify_task(prompt)
    return TASK_MODELS[task]


def call_router(prompt: str):
    models = choose_models(prompt)

    payload = {
        "model": models[0],
        "models": models,
        "messages": [
            {"role": "user", "content": prompt}
        ]
    }

    r = requests.post(
        "https://openrouter.ai/api/v1/chat/completions",
        headers={
            "Authorization": f"Bearer {OPENROUTER_API_KEY}",
            "Content-Type": "application/json",
            "HTTP-Referer": "http://localhost",
            "X-Title": "Hermes Auto Router"
        },
        json=payload,
        timeout=120
    )

    r.raise_for_status()
    return r.json()