import argparse
import json
import os
from urllib.error import HTTPError
from urllib.parse import urlsplit
from urllib.request import HTTPRedirectHandler, ProxyHandler, Request, build_opener


class NoRedirects(HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


def local_base(value):
    url = urlsplit(value)
    if (
        url.scheme != "http"
        or url.hostname not in {"localhost", "127.0.0.1", "::1"}
        or url.username is not None
        or url.password is not None
        or url.path.rstrip("/") != "/v1"
        or url.query
        or url.fragment
    ):
        raise argparse.ArgumentTypeError("Use a loopback HTTP address ending in /v1.")
    return value.rstrip("/")


def request(base, endpoint, payload=None, *, timeout=120):
    headers = {"Content-Type": "application/json"}
    key = os.environ.get("LOCAL_MODEL_API_KEY")
    if key:
        headers["Authorization"] = f"Bearer {key}"
    data = json.dumps(payload).encode("utf-8") if payload is not None else None
    message = Request(base + endpoint, data=data, headers=headers)
    opener = build_opener(ProxyHandler({}), NoRedirects())
    with opener.open(message, timeout=timeout) as response:
        return json.load(response)


def main():
    parser = argparse.ArgumentParser(description="Talk to a local Chat Completions endpoint.")
    parser.add_argument("--base", type=local_base, default="http://127.0.0.1:11434/v1")
    parser.add_argument("--list", action="store_true", help="Print available model IDs")
    parser.add_argument("--model", help="Exact model ID from --list")
    parser.add_argument("prompt", nargs="?")
    args = parser.parse_args()
    if not args.list and (not args.model or not args.prompt):
        parser.error("Use --list, or provide --model and a prompt.")

    try:
        if args.list:
            for model in request(args.base, "/models")["data"]:
                print(model["id"])
            return

        result = request(args.base, "/chat/completions", {
            "model": args.model,
            "messages": [{"role": "user", "content": args.prompt}],
            "temperature": 0,
            "max_tokens": 96,
            "stream": False,
        })
        text = result["choices"][0]["message"].get("content")
        if not isinstance(text, str) or not text.strip():
            raise ValueError("The server returned no visible reply. Check the model and output limit.")
        print(text)
    except HTTPError as error:
        raise SystemExit(f"Server returned HTTP {error.code}. Check the model ID, port, and API key.") from None
    except (OSError, KeyError, IndexError, TypeError, ValueError) as error:
        raise SystemExit(f"Could not read a local reply: {error}") from None


if __name__ == "__main__":
    main()