import sys

sys.dont_write_bytecode = True

import argparse
import json
import os
import re
from urllib.parse import urlparse

import requests

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from keycheck_common import (
    combined_provider_routing_hint,
    commit_status_transaction,
    default_input_file,
    default_proxy_file,
    ensure_output_files,
    iter_findings,
    keycheck_input_mode,
    load_checked_statuses,
    load_known_keys,
    load_proxies,
    mask_secret,
    provider_routing_database_failed,
    read_plain_keys,
    record_validation_result,
    recover_status_transaction,
    require_provider_authority,
    service_output_dir,
    should_skip_key,
    write_keycheck_event,
)
from keycheckers.provider_resolution import resolve_provider_key


SERVICE = "kimi"
DETECTOR = "KimiMoonshot"
OUTPUT_DIR = os.getenv("KEYCHECK_OUTPUT_DIR") or service_output_dir(SERVICE)
INPUT_FILE = os.getenv("KEYCHECK_INPUT_FILE") or default_input_file()
PROXY_FILE = os.getenv("KEYCHECK_PROXY_FILE") or default_proxy_file()

CHECKED_FILE = os.path.join(OUTPUT_DIR, "kimiChecked.txt")
RESULTS_FILE = os.path.join(OUTPUT_DIR, "kimiResults.jsonl")
STATUS_FILES = {
    "VALID": os.path.join(OUTPUT_DIR, "kimiAlive.txt"),
    "NO_BALANCE": os.path.join(OUTPUT_DIR, "kimiNoBalance.txt"),
    "DEAD": os.path.join(OUTPUT_DIR, "kimiDead.txt"),
    "RESTRICTED": os.path.join(OUTPUT_DIR, "kimiRestricted.txt"),
    "LIMITED": os.path.join(OUTPUT_DIR, "kimiLimited.txt"),
    "NETWORK": os.path.join(OUTPUT_DIR, "kimiNetwork.txt"),
    "UNKNOWN": os.path.join(OUTPUT_DIR, "kimiUnknown.txt"),
}

KIMI_DETECTOR_NAMES = {"kimimoonshot", "moonshotai", "moonshot", "kimi"}
KIMI_EXPLICIT_DETECTOR_NAMES = {"kimimoonshot", "moonshotai"}
QWEN_EXPLICIT_DETECTOR_NAMES = {"qwendashscope", "qwen_dashscope"}
DEEPSEEK_EXPLICIT_DETECTOR_NAMES = {"deepseekapikey", "deepseek_api_key"}
GENERIC_SK_PROVIDERS = {"qwen", "deepseek", "kimi", "zai"}
AMBIGUOUS_QWEN_DEEPSEEK_HINT = "ambiguous_qwen_deepseek"
AMBIGUOUS_GENERIC_SK_HINT = "ambiguous_generic_sk"
EXPLICIT_ASSIGNMENT_HINT_SOURCE = "explicit_assignment"
CANDIDATE_PROVIDER_ROUTE_FIELD = "_keycheck_candidate_provider_route"
KIMI_KEY_MAX_BYTES = 512
KIMI_KEY_REGEX = re.compile(
    r"(?<![A-Za-z0-9_-])sk-[A-Za-z0-9][A-Za-z0-9_-]{20,505}(?![A-Za-z0-9_-])"
)
FOREIGN_KEY_PREFIXES = ("sk-ant-", "sk-or-", "sk-proj-", "sk-svcacct-", "sk-admin-")
DEFAULT_BASE_URLS = (
    "https://api.moonshot.ai/v1",
    "https://api.moonshot.cn/v1",
)
QWEN_CONTEXT_REGEX = re.compile(
    r"(?:DASHSCOPE_API_KEY|QWEN_API_KEY|dashscope|qwen|model[_-]?studio|bailian)",
    re.IGNORECASE,
)
DEEPSEEK_CONTEXT_REGEX = re.compile(
    r"(?:DEEPSEEK_API_KEY|deepseek|api\.deepseek\.com)", re.IGNORECASE,
)
KIMI_CONTEXT_REGEX = re.compile(
    r"(?:MOONSHOT_API_KEY|KIMI_API_KEY|api\.moonshot\.(?:ai|cn)|platform\.kimi\.(?:ai|com))",
    re.IGNORECASE,
)


def normalize_base_url(value):
    return str(value or "").strip().rstrip("/")


def split_csv(value):
    if not value:
        return []
    if isinstance(value, str):
        return [item.strip() for item in value.split(",") if item.strip()]
    return [str(item).strip() for item in value if str(item).strip()]


def unique_ordered(values):
    output = []
    seen = set()
    for value in values:
        normalized = normalize_base_url(value)
        if normalized and normalized not in seen:
            seen.add(normalized)
            output.append(normalized)
    return output


def endpoint_label(base_url):
    parsed = urlparse(base_url)
    return parsed.netloc or base_url


def finding_detector_names(finding):
    if not isinstance(finding, dict):
        return set()
    extra = finding.get("ExtraData") if isinstance(finding.get("ExtraData"), dict) else {}
    names = {
        str(finding.get("DetectorName") or finding.get("detector") or "").strip().lower(),
        str(extra.get("name") or "").strip().lower(),
    }
    return {name for name in names if name}


def finding_has_detector(finding, detector_names):
    return bool(finding_detector_names(finding) & set(detector_names))


def key_from_text(*values):
    for value in values:
        for match in KIMI_KEY_REGEX.finditer(str(value or "")):
            key = match.group(0)
            if not key.startswith(FOREIGN_KEY_PREFIXES):
                return key
    return ""


def key_rejection_reason(key):
    value = str(key or "")
    try:
        key_bytes = len(value.encode("utf-8", errors="strict"))
    except UnicodeEncodeError:
        return "candidate is not valid UTF-8"
    if key_bytes > KIMI_KEY_MAX_BYTES:
        return f"candidate exceeds the {KIMI_KEY_MAX_BYTES}-byte key limit"
    if value.count("sk-") != 1:
        return "candidate contains multiple key prefixes"
    if not KIMI_KEY_REGEX.fullmatch(value) or value.startswith(FOREIGN_KEY_PREFIXES):
        return "candidate does not match the bounded Kimi/Moonshot key format"
    return ""


def finding_provider_routing_hint(finding):
    if not isinstance(finding, dict):
        return ""
    context = finding.get("ScannerContext") if isinstance(finding.get("ScannerContext"), dict) else {}
    persisted_hint = context.get("provider_hint")
    if context.get("provider_hint_source") == EXPLICIT_ASSIGNMENT_HINT_SOURCE and persisted_hint in (
        *GENERIC_SK_PROVIDERS, AMBIGUOUS_QWEN_DEEPSEEK_HINT, AMBIGUOUS_GENERIC_SK_HINT,
    ):
        return persisted_hint

    parts = [str(context.get(key) or "") for key in ("nearby", "file")]
    metadata = finding.get("SourceMetadata") if isinstance(finding.get("SourceMetadata"), dict) else {}
    data = metadata.get("Data") if isinstance(metadata.get("Data"), dict) else {}
    for details in data.values():
        if isinstance(details, dict):
            parts.extend(str(details.get(key) or "") for key in ("file", "repository", "repo", "link", "image"))

    text = "\n".join(parts)
    evidence = set()
    if QWEN_CONTEXT_REGEX.search(text) or finding_has_detector(finding, QWEN_EXPLICIT_DETECTOR_NAMES):
        evidence.add("qwen")
    if DEEPSEEK_CONTEXT_REGEX.search(text) or finding_has_detector(finding, DEEPSEEK_EXPLICIT_DETECTOR_NAMES):
        evidence.add("deepseek")
    if KIMI_CONTEXT_REGEX.search(text) or finding_has_detector(finding, KIMI_EXPLICIT_DETECTOR_NAMES):
        evidence.add("kimi")
    if persisted_hint == AMBIGUOUS_QWEN_DEEPSEEK_HINT:
        evidence.update(("qwen", "deepseek"))
    elif persisted_hint == AMBIGUOUS_GENERIC_SK_HINT:
        evidence.update(GENERIC_SK_PROVIDERS)
    elif persisted_hint in GENERIC_SK_PROVIDERS:
        evidence.add(persisted_hint)
    if len(evidence) > 1:
        return (
            AMBIGUOUS_QWEN_DEEPSEEK_HINT
            if evidence == {"qwen", "deepseek"}
            else AMBIGUOUS_GENERIC_SK_HINT
        )
    return next(iter(evidence)) if evidence else ""


def iter_candidate_keys(input_file, plain_files, trusted_retry_files=None):
    detector_names = [
        "KimiMoonshot", "MoonshotAI", "Moonshot", "Kimi",
        "kimimoonshot", "moonshotai", "moonshot", "kimi", "CustomRegex",
    ]
    routing_decisions = {}
    seen = set()
    for item in iter_findings(input_file, detector_names):
        finding = dict(item.get("finding") or {})
        finding.pop(CANDIDATE_PROVIDER_ROUTE_FIELD, None)
        candidate_metadata = item.get("candidate_metadata")
        persisted_route = ""
        if keycheck_input_mode() == "postgres" and isinstance(candidate_metadata, dict):
            persisted_route = str(candidate_metadata.get("provider_hint") or "").lower()
        if persisted_route == SERVICE:
            finding[CANDIDATE_PROVIDER_ROUTE_FIELD] = SERVICE
        if persisted_route != SERVICE and not finding_has_detector(finding, KIMI_DETECTOR_NAMES):
            continue
        key = key_from_text(item.get("raw"), item.get("raw_v2"), finding.get("Raw"), finding.get("RawV2"))
        if not key or key_rejection_reason(key):
            continue
        if persisted_route == SERVICE:
            hint = SERVICE
        else:
            local_hint = finding_provider_routing_hint(finding)
            if key not in routing_decisions:
                routing_decisions[key] = combined_provider_routing_hint(key, local_hint)
            hint = routing_decisions[key]
            if provider_routing_database_failed():
                raise RuntimeError("provider routing evidence lookup failed closed")
        if hint == "kimi" or (
            keycheck_input_mode() == "postgres"
            and hint in (AMBIGUOUS_QWEN_DEEPSEEK_HINT, AMBIGUOUS_GENERIC_SK_HINT)
        ):
            seen.add(key)
            yield key, item.get("source") or input_file, finding

    owned_retry_paths = {
        os.path.normcase(os.path.abspath(path)) for path in STATUS_FILES.values()
    }
    retry_files = [
        path for path in (trusted_retry_files or [])
        if os.path.normcase(os.path.abspath(path)) in owned_retry_paths
    ]
    for item in read_plain_keys(retry_files, KIMI_KEY_REGEX):
        key = item["key"]
        if key_rejection_reason(key) or key in seen:
            continue
        hint = combined_provider_routing_hint(key, "kimi")
        if provider_routing_database_failed():
            raise RuntimeError("provider routing evidence lookup failed closed")
        if hint == "kimi":
            seen.add(key)
            yield key, item["source"], {}


def redact_text(value, key):
    text = str(value or "")[:1000]
    if key:
        text = text.replace(key, "***REDACTED***")
    return KIMI_KEY_REGEX.sub("***REDACTED***", text)


def parse_error(response, key):
    try:
        payload = response.json()
    except ValueError:
        payload = {}
    error = payload.get("error") if isinstance(payload, dict) else {}
    if not isinstance(error, dict):
        error = {}
    message = error.get("message") or response.text[:500]
    return {
        "http_status": response.status_code,
        "code": error.get("code") or error.get("type") or "",
        "message": redact_text(message, key),
    }


def classify_error(error):
    http_status = int(error.get("http_status") or 0)
    code = str(error.get("code") or "").lower()
    message = str(error.get("message") or "").lower()
    if http_status == 401 or any(marker in code for marker in ("invalid_authentication", "invalid_api_key")):
        return "DEAD"
    if http_status == 403:
        return "RESTRICTED"
    if http_status == 429:
        if any(marker in code + " " + message for marker in ("quota", "balance", "payment")):
            return "NO_BALANCE"
        return "LIMITED"
    if 500 <= http_status <= 599:
        return "NETWORK"
    return "UNKNOWN"


def check_base_url(key, base_url, proxy, timeout, debug=False):
    url = f"{normalize_base_url(base_url)}/users/me/balance"
    headers = {"Authorization": f"Bearer {key}", "Accept": "application/json"}
    try:
        response = requests.get(url, headers=headers, proxies=proxy, timeout=timeout)
    except requests.RequestException as exc:
        return {
            "status": "NETWORK", "base_url": base_url,
            "region": endpoint_label(base_url), "message": redact_text(exc, key),
        }
    if debug:
        print(
            f"  DEBUG {endpoint_label(base_url)} balance: HTTP {response.status_code}: "
            f"{redact_text(response.text[:500], key)}"
        )
    if response.status_code == 200:
        try:
            payload = response.json()
            data = payload.get("data") if isinstance(payload, dict) else None
            if not isinstance(data, dict) or "available_balance" not in data:
                raise ValueError("missing data.available_balance")
            available = float(data.get("available_balance"))
            voucher = float(data.get("voucher_balance", 0) or 0)
            cash = float(data.get("cash_balance", 0) or 0)
        except (TypeError, ValueError, json.JSONDecodeError) as exc:
            return {
                "status": "UNKNOWN", "base_url": base_url,
                "region": endpoint_label(base_url),
                "message": f"invalid balance response: {exc}",
            }
        status = "VALID" if available > 0 else "NO_BALANCE"
        return {
            "status": status,
            "authenticated": True,
            "base_url": base_url,
            "region": endpoint_label(base_url),
            "balance_usd": round(available, 6),
            "voucher_balance_usd": round(voucher, 6),
            "cash_balance_usd": round(cash, 6),
            "message": f"available_balance=${available:.6f}",
        }
    error = parse_error(response, key)
    return {
        "status": classify_error(error),
        "base_url": base_url,
        "region": endpoint_label(base_url),
        "http_status": response.status_code,
        "error": error,
        "message": error.get("message") or "",
    }


def choose_final_status(attempts):
    statuses = [attempt.get("status") for attempt in attempts]
    for status in ("VALID", "NO_BALANCE", "LIMITED", "RESTRICTED", "UNKNOWN", "NETWORK"):
        if status in statuses:
            return status
    return "DEAD"


def check_key(key, base_urls, proxy, timeout, debug=False):
    rejection = key_rejection_reason(key)
    if rejection:
        return {"status": "UNKNOWN", "message": rejection, "candidate_rejected": True}
    attempts = []
    for base_url in base_urls:
        result = check_base_url(key, base_url, proxy, timeout, debug)
        attempts.append(result)
        if result.get("status") in ("VALID", "NO_BALANCE"):
            return {**result, "attempts": attempts}
    status = choose_final_status(attempts)
    selected = next((attempt for attempt in attempts if attempt.get("status") == status), {})
    return {**selected, "status": status, "attempts": attempts}


def ensure_files():
    ensure_output_files([CHECKED_FILE, RESULTS_FILE, *STATUS_FILES.values()])
    recover_status_transaction(CHECKED_FILE, STATUS_FILES)


def write_result(key, result, source, finding):
    status = result.get("status") or "UNKNOWN"
    write_keycheck_event(SERVICE, RESULTS_FILE, key, result, source, finding, DETECTOR)
    commit_status_transaction(
        CHECKED_FILE, STATUS_FILES, key, status,
        result.get("message", ""), result.get("region") or source,
    )
    record_validation_result(SERVICE, key, result, source, finding, DETECTOR)


def retry_statuses_from_args(args):
    statuses = set()
    if args.retry_network:
        statuses.add("NETWORK")
    if args.retry_limited:
        statuses.add("LIMITED")
    if args.retry_unknown:
        statuses.add("UNKNOWN")
    if args.retry_restricted:
        statuses.add("RESTRICTED")
    if args.retry_no_balance:
        statuses.add("NO_BALANCE")
    if args.retry_valid:
        statuses.add("VALID")
    return statuses


def retry_input_files_from_args(args):
    if args.recheck_all:
        statuses = list(STATUS_FILES)
    else:
        statuses = [
            status for flag, status in (
                (args.retry_network, "NETWORK"),
                (args.retry_limited, "LIMITED"),
                (args.retry_unknown, "UNKNOWN"),
                (args.retry_restricted, "RESTRICTED"),
                (args.retry_no_balance, "NO_BALANCE"),
                (args.retry_valid, "VALID"),
            ) if flag
        ]
    return [STATUS_FILES[status] for status in statuses]


def base_urls_from_args(args):
    custom = []
    for value in args.base_url:
        custom.extend(split_csv(value))
    custom.extend(split_csv(os.getenv("KIMI_BASE_URLS") or os.getenv("MOONSHOT_BASE_URLS")))
    defaults = [] if args.no_default_base_urls else DEFAULT_BASE_URLS
    return unique_ordered([*custom, *defaults])


def parse_args():
    parser = argparse.ArgumentParser(description="Kimi / Moonshot AI key checker")
    parser.add_argument("--input", default=INPUT_FILE)
    parser.add_argument("--plain", action="append", default=[])
    parser.add_argument("--proxy-file", default=PROXY_FILE)
    parser.add_argument("--timeout", type=int, default=15)
    parser.add_argument("--max-keys", type=int, default=0)
    parser.add_argument("--base-url", action="append", default=[])
    parser.add_argument("--no-default-base-urls", action="store_true")
    parser.add_argument("--retry-network", action="store_true")
    parser.add_argument("--retry-limited", action="store_true")
    parser.add_argument("--retry-unknown", action="store_true")
    parser.add_argument("--retry-restricted", action="store_true")
    parser.add_argument("--retry-no-balance", action="store_true")
    parser.add_argument("--retry-valid", action="store_true")
    parser.add_argument("--recheck-all", action="store_true")
    parser.add_argument("--debug", action="store_true")
    return parser.parse_args()


def main():
    require_provider_authority(SERVICE)
    args = parse_args()
    ensure_files()
    proxy_cycler = load_proxies(args.proxy_file)
    checked = load_checked_statuses(CHECKED_FILE)
    known = load_known_keys(CHECKED_FILE, STATUS_FILES)
    retry_statuses = retry_statuses_from_args(args)
    retry_files = retry_input_files_from_args(args)
    base_urls = base_urls_from_args(args)
    if not base_urls:
        raise SystemExit("No Kimi/Moonshot base URLs configured")

    print("--- Kimi / Moonshot AI key checker ---")
    print("base_urls: " + ", ".join(endpoint_label(url) for url in base_urls))
    processed = 0
    skipped = 0
    for key, source, finding in iter_candidate_keys(args.input, args.plain, retry_files):
        finding = dict(finding or {})
        candidate_route = str(finding.pop(CANDIDATE_PROVIDER_ROUTE_FIELD, "") or "").lower()
        if should_skip_key(
            key, checked, known, args, retry_statuses, service=SERVICE,
            source=source, finding=finding, detector=DETECTOR,
        ):
            skipped += 1
            continue
        if args.max_keys and processed >= args.max_keys:
            break
        processed += 1
        print(f"\n[{processed}] Kimi/Moonshot candidate {mask_secret(key)} from {source}")
        proxy = next(proxy_cycler) if proxy_cycler else None
        routing_hint = "kimi"
        if keycheck_input_mode() == "postgres":
            if candidate_route == SERVICE:
                routing_hint = SERVICE
            else:
                routing_hint = combined_provider_routing_hint(key, finding_provider_routing_hint(finding))
                if provider_routing_database_failed():
                    raise RuntimeError("provider routing evidence lookup failed closed")
        if routing_hint in (AMBIGUOUS_QWEN_DEEPSEEK_HINT, AMBIGUOUS_GENERIC_SK_HINT):
            result = resolve_provider_key(
                key, finding, proxy, args.timeout, args.debug,
                hint=routing_hint, origin_service=SERVICE,
            )
        else:
            result = check_key(key, base_urls, proxy, args.timeout, args.debug)
        print(f"  STATUS: {result['status']} | {result.get('message', '')[:200]}")
        write_result(key, result, source, finding)
        known.add(key)
        checked[key] = result["status"]

    print(f"\nDone. Processed={processed}, skipped={skipped}, results={RESULTS_FILE}")


if __name__ == "__main__":
    main()