"""Python 3, standard library only. Import never sends a request.

python router.py --base-url YOUR_SERVICE_URL             # routing only
python router.py --base-url YOUR_SERVICE_URL --execute   # billable inference

Own provider credentials stay in this process, never go to the routing service.
Fallbacks are YOUR policy and can differ from the routing service's limits.
"""
import argparse
from datetime import datetime, timezone
import json
import os
import time
from urllib.parse import urlencode, urlsplit
from urllib.request import Request, urlopen


MODELS = {
    "openai": "gpt-4o",
    "fireworks": "accounts/fireworks/models/gpt-oss-120b",
    "cerebras": "gpt-oss-120b",
}
PROVIDERS = {
    "openai": ("https://api.openai.com/v1/chat/completions", "OPENAI_API_KEY"),
    "fireworks": ("https://api.fireworks.ai/inference/v1/chat/completions", "FIREWORKS_API_KEY"),
    "cerebras": ("https://api.cerebras.ai/v1/chat/completions", "CEREBRAS_API_KEY"),
}
PARAMETERS = {
    "task_type", "latency_sensitivity", "cost_sensitivity", "provider",
    "max_latency_ms", "max_cost_per_1k_tokens_usd",
}


class Router:
    def __init__(self, base_url, fallback, models=None, routing_api_key=None):
        url = urlsplit(base_url)
        if (url.scheme not in ("http", "https") or not url.netloc or url.username or url.password
                or url.query or url.fragment):
            raise ValueError("Use an http(s) service URL without credentials")
        self.base_url = base_url.rstrip("/")
        self.models = dict(models or MODELS)
        if self.models.get(fallback["provider"]) != fallback["model"]:
            raise ValueError("Fallback must match your supported provider/model map")
        self.fallback = dict(fallback)
        self.routing_api_key = routing_api_key
        self.cache = {}

    def _fallback(self, reason):
        return {**self.fallback, "source": "caller_fallback", "reason": reason}

    def choose(self, **params):
        if not set(params).issubset(PARAMETERS):
            return self._fallback("invalid_parameters")
        key = urlencode(sorted(params.items()))
        now = time.time()
        cached = self.cache.get(key)
        if cached and cached[0] > now:
            choice = cached[1]
            if self.models.get(choice["provider"]) == choice["model"]:
                return {**choice, "source": "cache"}
        self.cache.pop(key, None)
        headers = {"Accept": "application/json"}
        if self.routing_api_key:
            headers["X-API-Key"] = self.routing_api_key
        try:
            request = Request(f"{self.base_url}/v1/should-call?{key}", headers=headers)
            with urlopen(request, timeout=2) as response:
                data = json.load(response)
            if data.get("should_call") is not True:
                return self._fallback("no_eligible_choice")
            provider, model = data["recommended_provider"], data["recommended_model"]
            if self.models.get(provider) != model:
                return self._fallback("unsupported_model")
            measurement = data["measurement"]
            expiry = datetime.fromisoformat(measurement["expires_at"].replace("Z", "+00:00"))
            if expiry.tzinfo is None or measurement["sample_count"] < 1:
                return self._fallback("invalid_evidence")
            until = min(time.time() + 30, expiry.astimezone(timezone.utc).timestamp() - 1)
            if until <= time.time():
                return self._fallback("expired_evidence")
            choice = {"provider": provider, "model": model, "source": "routing_service"}
            if len(self.cache) >= 32:
                self.cache.clear()
            self.cache[key] = (until, choice)
            return choice
        except Exception:
            # Do not log URLs, headers or provider error bodies.
            return self._fallback("routing_unavailable")

    def run(self, prompt, call_provider, **params):
        choice = self.choose(**params)
        try:
            return call_provider(choice["provider"], choice["model"], prompt)
        except Exception:
            if (choice["provider"], choice["model"]) == (self.fallback["provider"], self.fallback["model"]):
                raise
            self.cache.clear()
            # At most one fallback attempt, never an unbounded retry loop.
            return call_provider(self.fallback["provider"], self.fallback["model"], prompt)


def call_provider(provider, model, prompt):
    endpoint, key_name = PROVIDERS[provider]
    api_key = os.environ.get(key_name)
    if not api_key:
        raise RuntimeError(f"Configure your own {key_name}")
    request = Request(endpoint, data=json.dumps({
        "model": model, "messages": [{"role": "user", "content": prompt}], "max_tokens": 256,
    }).encode(), headers={"Content-Type": "application/json", "Authorization": f"Bearer {api_key}"})
    with urlopen(request, timeout=30) as response:
        result = json.load(response)
    content = result["choices"][0]["message"]["content"]
    if not isinstance(content, str) or not content.strip():
        raise RuntimeError("Provider returned no usable text")
    return content


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--base-url", required=True)
    parser.add_argument("--execute", action="store_true", help="Accept billable provider inference")
    parser.add_argument("--prompt", default="Explain what an AI agent is in one sentence.")
    args = parser.parse_args()
    available = {key: model for key, model in MODELS.items() if os.environ.get(PROVIDERS[key][1])}
    if args.execute and not available:
        parser.error("Configure at least one of your own provider keys before --execute")
    models = available if args.execute else MODELS
    fallback_provider = next(iter(models))
    router = Router(args.base_url, {"provider": fallback_provider, "model": models[fallback_provider]},
                    models=models, routing_api_key=os.environ.get("INFERENCELATENCY_API_KEY"))
    if args.execute:
        try:
            print(router.run(args.prompt, call_provider))
        except Exception:
            parser.exit(1, "Provider inference failed; check your own credentials, quota and model support.\n")
    else:
        print(json.dumps(router.choose(), indent=2))


if __name__ == "__main__":
    main()