#!/usr/bin/env python3
"""Add newly-hunted keys (FreeModel / DeepSeek / Kimi) to NewAPI.

Reads each provider's non-401 key file, probes both endpoints to find
which one actually accepts the key, creates a channel for every
working combination, tests it, and auto-disables channels that fail
(so they don't break routing, but can be re-enabled later).
"""

import json
import sys
import time
from pathlib import Path

from newapi_client import (
    add_channel, list_channels, update_channel, test_channel,
    TYPE_OPENAI, TYPE_ANTHROPIC,
)

HERE = Path(__file__).parent
HUNTER = HERE / "llm-key-hunter" / "results"

# Each entry: provider label, key file, list of endpoint specs
# Endpoint spec: (channel_name_suffix, type, base_url, models, header_template)
PROVIDERS = [
    {
        "name": "FreeModel",
        "file": HUNTER / "freemodel" / "all_non_401.txt",
        "parse": lambda line: line.split("|", 2)[1],  # USABLE|key|src
        "endpoints": [
            {
                "suffix": "Anthropic",
                "type": TYPE_ANTHROPIC,
                "base_url": "https://cc.freemodel.dev",
                "models": (
                    "claude-sonnet-4-6,claude-opus-4-6,claude-opus-5,"
                    "claude-opus-4-8,claude-opus-4-7,claude-fable-5,"
                    "claude-haiku-4-5,claude-sonnet-4-20250514,"
                    "claude-opus-4-20250514,claude-3-5-haiku-20241022,"
                    "claude-3-7-sonnet-20250219"
                ),
            },
            {
                "suffix": "OpenAI",
                "type": TYPE_OPENAI,
                "base_url": "https://api.freemodel.dev",
                "models": (
                    "gpt-5.6-luna,gpt-5.6-sol,gpt-5.6-terra,"
                    "claude-sonnet-4-20250514"
                ),
            },
        ],
    },
    {
        "name": "DeepSeek",
        "file": HUNTER / "deepseek" / "all_non_401.txt",
        "parse": lambda line: line.split("|", 2)[1],
        "endpoints": [
            {
                "suffix": "",
                "type": TYPE_OPENAI,
                "base_url": "https://api.deepseek.com",
                "models": "deepseek-chat,deepseek-reasoner",
            },
        ],
    },
    {
        "name": "Kimi",
        "file": HUNTER / "kimi" / "all_non_401.txt",
        "parse": lambda line: line.split("|", 2)[1],
        "endpoints": [
            {
                "suffix": "",
                "type": TYPE_OPENAI,
                "base_url": "https://api.moonshot.cn",
                "models": (
                    "moonshot-v1-8k,moonshot-v1-32k,moonshot-v1-128k,"
                    "kimi-k2-0905-preview,kimi-k2-turbo-preview,"
                    "kimi-latest"
                ),
            },
        ],
    },
]


def existing_channel_keys():
    """Return set of keys already present in NewAPI (from full channel details)."""
    keys = set()
    for ch in list_channels():
        cid = ch["id"]
        try:
            from newapi_client import get_channel
            d = get_channel(cid)
            k = (d or {}).get("key", "")
            if k:
                keys.add(k)
        except Exception:
            pass
    return keys


def main():
    print("Loading existing channels to avoid duplicates...")
    existing_keys = existing_channel_keys()
    print(f"  {len(existing_keys)} unique keys already configured.\n")

    added = []
    skipped = 0
    failed_test = 0

    for prov in PROVIDERS:
        path = prov["file"]
        if not path.exists():
            print(f"[{prov['name']}] no file at {path}, skipping")
            continue

        raw_keys = []
        seen = set()
        for line in path.read_text().splitlines():
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            try:
                k = prov["parse"](line)
            except IndexError:
                continue
            if k and k not in seen:
                seen.add(k)
                raw_keys.append(k)

        print(f"[{prov['name']}] {len(raw_keys)} keys from {path.name}")

        for idx, key in enumerate(raw_keys, 1):
            if key in existing_keys:
                print(f"  [{idx}/{len(raw_keys)}] {key[:18]}... already in NewAPI, skip")
                skipped += 1
                continue

            for ep in prov["endpoints"]:
                suffix = ep["suffix"]
                name = f"{prov['name']}_{key[-6:]}"
                if suffix:
                    name = f"{name}-{suffix}"

                channel = {
                    "name": name,
                    "type": ep["type"],
                    "key": key,
                    "base_url": ep["base_url"],
                    "models": ep["models"],
                    "groups": ["default"],
                    "model_mapping": "",
                    "priority": 0,
                    "weight": 0,
                    "auto_ban": 1,
                }
                success, msg = add_channel(channel)
                if not success:
                    print(f"  [{idx}] {name} ADD FAIL: {msg}")
                    continue

                # find new channel id
                new_id = None
                for ch in list_channels():
                    if ch["name"] == name:
                        new_id = ch["id"]
                        break
                if new_id is None:
                    print(f"  [{idx}] {name} added but couldn't find id")
                    continue

                # test it
                t_ok, t_msg, _ = test_channel(new_id)
                if t_ok:
                    print(f"  [{idx}] {name} ID={new_id} ✅ {t_msg[:80]}")
                    added.append((new_id, name, key[:18], t_msg[:80]))
                else:
                    # auto-disable so it doesn't burn retries
                    update_channel(new_id, status=2)
                    print(f"  [{idx}] {name} ID={new_id} ❌ DISABLED ({t_msg[:80]})")
                    failed_test += 1
                time.sleep(0.3)

        print()

    print("=== Summary ===")
    print(f"  Added & test-passed: {len(added)}")
    print(f"  Added but auto-disabled: {failed_test}")
    print(f"  Skipped (already present): {skipped}")
    if added:
        print("\nLive channels:")
        for cid, name, k, msg in added:
            print(f"  ID={cid:3d} {name:35s} key={k}...  {msg}")


if __name__ == "__main__":
    main()