import argparse
import os

import requests

CONTEXT7_API_KEY = os.environ.get("CONTEXT7_API_KEY")

headers = {"Authorization": f"Bearer {CONTEXT7_API_KEY}"}


def search_library(library: str, query: str) -> str:
    # Step 1: Search for the library
    search_response = requests.get(
        "https://context7.com/api/v2/libs/search",
        headers=headers,
        params={"libraryName": library, "query": query},
    )
    search_response.raise_for_status()
    data = search_response.json()
    best_match = data["results"][0]
    print(f"Found: {best_match['title']} ({best_match['id']})")
    return best_match["id"]


def search_snippet(library: str, query: str) -> None:
    # Step 2: Get documentation context
    context_response = requests.get(
        "https://context7.com/api/v2/context",
        headers=headers,
        params={
            "libraryId": library,
            "query": query,
            "type": "json",
        },
    )
    context_response.raise_for_status()
    docs = context_response.json()

    for snippet in docs["codeSnippets"]:
        print(f"Title: {snippet['codeTitle']}")
        for code in snippet["codeList"]:
            print(f"Code: {code['code']}")

    for info in docs["infoSnippets"]:
        print(f"Content: {info['content']}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("library", help="Library name to search for")
    parser.add_argument("query", help="Documentation query to run")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    library_id = search_library(args.library, args.query)
    search_snippet(library_id, args.query)


if __name__ == "__main__":
    main()