from . import Plugin, tool
from typing import Annotated, Literal
import os
from tavily import AsyncTavilyClient


class Search(Plugin):
    def __init__(
        self,
        tavily_api_key: str | None = None,
    ):
        super().__init__()
        api_key = tavily_api_key or os.environ.get("TAVILY_API_KEY")
        if not api_key:
            raise ValueError("TAVILY_API_KEY is required for SearchPlugin")
        self.__tavily = AsyncTavilyClient(api_key=api_key)

    @tool
    async def web_search(
        self,
        query: Annotated[
            str, "The search query. Please be as specific and verbose as possible."
        ],
        topic: Annotated[
            Literal["general", "news", "finance"] | None,
            "The topic of the search. Default is 'general'.",
        ] = None,
        time_range: Annotated[
            Literal["day", "week", "month", "year"] | None,
            "The time range back from the current date based on publish date or last updated date.",
        ] = None,
        max_results: Annotated[
            int | None,
            "The maximum number of search results to return. Default is 5. It must be between 0 and 20.",
        ] = None,
    ):
        """
        Perform web search on the given query.
        Returning the top related search results in json format.
        When necessary, you need to combine this tool with the get_webpage_content tools (if available), to browse the web in depth by jumping through links.
        """

        tavily_results = await self.__tavily.search(
            query=query,
            search_depth="advanced",
            # max_results=10,
            include_answer=True,
            include_images=True,
            include_image_descriptions=True,
            topic=topic,  # type: ignore
            time_range=time_range,  # type: ignore
            max_results=max_results,  # type: ignore
        )
        return tavily_results


_ALL_COUNTRIES = [
    # Default country
    "Australia",
    # Other countries
    "Afghanistan",
    "Albania",
    "Antarctica",
    "Algeria",
    "American Samoa",
    "Andorra",
    "Angola",
    "Antigua and Barbuda",
    "Azerbaijan",
    "Argentina",
    "Austria",
    "The Bahamas",
    "Bahrain",
    "Bangladesh",
    "Armenia",
    "Barbados",
    "Belgium",
    "Bhutan",
    "Bolivia",
    "Bosnia and Herzegovina",
    "Botswana",
    "Brazil",
    "Belize",
    "Solomon Islands",
    "Brunei",
    "Bulgaria",
    "Myanmar (Burma)",
    "Burundi",
    "Cambodia",
    "Cameroon",
    "Canada",
    "Cabo Verde",
    "Central African Republic",
    "Sri Lanka",
    "Chad",
    "Chile",
    "China",
    "Christmas Island",
    "Cocos (Keeling) Islands",
    "Colombia",
    "Comoros",
    "Republic of the Congo",
    "Democratic Republic of the Congo",
    "Cook Islands",
    "Costa Rica",
    "Croatia",
    "Cyprus",
    "Czechia",
    "Benin",
    "Denmark",
    "Dominica",
    "Dominican Republic",
    "Ecuador",
    "El Salvador",
    "Equatorial Guinea",
    "Ethiopia",
    "Eritrea",
    "Estonia",
    "South Georgia and the South Sandwich Islands",
    "Fiji",
    "Finland",
    "France",
    "French Polynesia",
    "French Southern and Antarctic Lands",
    "Djibouti",
    "Gabon",
    "Georgia",
    "The Gambia",
    "Germany",
    "Ghana",
    "Kiribati",
    "Greece",
    "Grenada",
    "Guam",
    "Guatemala",
    "Guinea",
    "Guyana",
    "Haiti",
    "Heard Island and McDonald Islands",
    "Vatican City",
    "Honduras",
    "Hungary",
    "Iceland",
    "India",
    "Indonesia",
    "Iraq",
    "Ireland",
    "Israel",
    "Italy",
    "Cote d'Ivoire",
    "Jamaica",
    "Japan",
    "Kazakhstan",
    "Jordan",
    "Kenya",
    "South Korea",
    "Kuwait",
    "Kyrgyzstan",
    "Laos",
    "Lebanon",
    "Lesotho",
    "Latvia",
    "Liberia",
    "Libya",
    "Liechtenstein",
    "Lithuania",
    "Luxembourg",
    "Madagascar",
    "Malawi",
    "Malaysia",
    "Maldives",
    "Mali",
    "Malta",
    "Mauritania",
    "Mauritius",
    "Mexico",
    "Monaco",
    "Mongolia",
    "Moldova",
    "Montenegro",
    "Morocco",
    "Mozambique",
    "Oman",
    "Namibia",
    "Nauru",
    "Nepal",
    "Netherlands",
    "Curacao",
    "Sint Maarten",
    "Caribbean Netherlands",
    "New Caledonia",
    "Vanuatu",
    "New Zealand",
    "Nicaragua",
    "Niger",
    "Nigeria",
    "Niue",
    "Norfolk Island",
    "Norway",
    "Northern Mariana Islands",
    "United States Minor Outlying Islands",
    "Micronesia",
    "Marshall Islands",
    "Palau",
    "Pakistan",
    "Panama",
    "Papua New Guinea",
    "Paraguay",
    "Peru",
    "Philippines",
    "Pitcairn Islands",
    "Poland",
    "Portugal",
    "Guinea-Bissau",
    "Timor-Leste",
    "Qatar",
    "Romania",
    "Rwanda",
    "Saint Helena, Ascension and Tristan da Cunha",
    "Saint Kitts and Nevis",
    "Saint Lucia",
    "Saint Pierre and Miquelon",
    "Saint Vincent and the Grenadines",
    "San Marino",
    "Sao Tome and Principe",
    "Saudi Arabia",
    "Senegal",
    "Serbia",
    "Seychelles",
    "Sierra Leone",
    "Singapore",
    "Slovakia",
    "Vietnam",
    "Slovenia",
    "Somalia",
    "South Africa",
    "Zimbabwe",
    "Spain",
    "Suriname",
    "Eswatini",
    "Sweden",
    "Switzerland",
    "Tajikistan",
    "Thailand",
    "Togo",
    "Tokelau",
    "Tonga",
    "Trinidad and Tobago",
    "United Arab Emirates",
    "Tunisia",
    "Turkiye",
    "Turkmenistan",
    "Tuvalu",
    "Uganda",
    "Ukraine",
    "North Macedonia",
    "Egypt",
    "United Kingdom",
    "Guernsey",
    "Jersey",
    "Isle of Man",
    "Tanzania",
    "United States",
    "Burkina Faso",
    "Uruguay",
    "Uzbekistan",
    "Venezuela",
    "Wallis and Futuna",
    "Samoa",
    "Yemen",
    "Zambia",
]
