import requests
from typing import Any, Dict, List, Literal, Optional, Union
from urllib.parse import urljoin


class FridayClient:
    """
    Python SDK for the Friday API.

    Args:
        api_key (str): Your Friday API key.
        base_url (str, optional): Base URL for the API.
    """

    def __init__(self, api_key: str, base_url: str = "https://api.fridaydata.com"):
        self.api_key = api_key
        self.base_url = base_url.rstrip("/")
        self.session = requests.Session()
        self.session.headers.update(
            {
                "X-API-Key": api_key,
                "Content-Type": "application/json",
            }
        )

    def _make_request(
        self,
        method: str,
        endpoint: str,
        *,
        use_api_key: bool = True,
        bearer_token: Optional[str] = None,
        **kwargs: Any,
    ) -> dict:
        """Make a request to the API."""
        url = urljoin(self.base_url + "/", endpoint.lstrip("/"))

        headers = dict(kwargs.pop("headers", {}) or {})
        if use_api_key:
            headers.setdefault("X-API-Key", self.api_key)
        else:
            headers.pop("X-API-Key", None)

        if bearer_token:
            headers["Authorization"] = f"Bearer {bearer_token}"

        response = self.session.request(method, url, headers=headers, **kwargs)
        response.raise_for_status()
        return response.json()

    def get_status(self) -> dict:
        return self._make_request("GET", "/status")

    def get_profile(
        self,
        profile_url: str,
        realtime: bool = False,
        prompt: Optional[str] = None,
        custom_schema: Optional[str] = None,
    ) -> dict:
        params: Dict[str, Any] = {
            "profile_url": profile_url,
            "realtime": realtime,
        }
        if prompt is not None:
            params["prompt"] = prompt
        if custom_schema is not None:
            params["custom_schema"] = custom_schema

        return self._make_request("GET", "/profile", params=params)

    def analyze_company(self, linkedin_url: str, realtime: bool = False) -> dict:
        return self._make_request(
            "POST",
            "/analyze-company",
            json={"linkedin_url": linkedin_url, "realtime": realtime},
        )

    def scrape(
        self,
        url: str,
        formats: Optional[List[Literal["html", "markdown", "links", "text"]]] = None,
    ) -> dict:
        return self._make_request(
            "POST",
            "/scrape",
            json={"url": url, "formats": formats or ["html"]},
        )

    def crawl(
        self,
        url: str,
        formats: Optional[List[Literal["html", "markdown", "links", "text"]]] = None,
        max_pages: int = 10,
    ) -> dict:
        return self._make_request(
            "POST",
            "/crawl",
            json={
                "url": url,
                "formats": formats or ["html", "markdown", "links"],
                "max_pages": max_pages,
            },
        )

    def extract(
        self,
        url: str,
        query: str,
        custom_schema: Optional[Union[Dict[str, Any], List[Dict[str, Any]]]] = None,
    ) -> dict:
        payload: Dict[str, Any] = {
            "url": url,
            "query": query,
        }
        if custom_schema is not None:
            payload["custom_schema"] = custom_schema

        return self._make_request("POST", "/extract", json=payload)

    def search(self, query: str, location: str = "US", num_results: int = 15) -> dict:
        return self._make_request(
            "POST",
            "/search",
            json={
                "query": query,
                "location": location,
                "num_results": num_results,
            },
        )

    def reddit_search(
        self,
        query: str,
        sort: Literal["top", "hot", "new", "relevance"] = "top",
        time_filter: Literal["hour", "day", "week", "month", "year", "all"] = "all",
        limit: int = 25,
    ) -> dict:
        return self._make_request(
            "POST",
            "/reddit/search",
            json={
                "query": query,
                "sort": sort,
                "time_filter": time_filter,
                "limit": limit,
            },
        )

    def reddit_subreddit_search(
        self,
        subreddit: str,
        query: str,
        sort: Literal["top", "hot", "new", "relevance"] = "top",
        time_filter: Literal["hour", "day", "week", "month", "year", "all"] = "all",
        limit: int = 25,
    ) -> dict:
        return self._make_request(
            "POST",
            "/reddit/subreddit/search",
            json={
                "subreddit": subreddit,
                "query": query,
                "sort": sort,
                "time_filter": time_filter,
                "limit": limit,
            },
        )

    def reddit_user_posts(self, username: str, limit: int = 25) -> dict:
        return self._make_request(
            "POST",
            "/reddit/user/posts",
            json={"username": username, "limit": limit},
        )

    def scrape_twitter_profile(self, profile_url: str, tweet_count: int = 10) -> dict:
        return self._make_request(
            "POST",
            "/twitter/profile",
            json={"profile_url": profile_url, "tweet_count": tweet_count},
        )

    def x_latest(self, handle: str, limit: int = 10, refresh: Optional[str] = None) -> dict:
        params: Dict[str, Any] = {"handle": handle, "limit": limit}
        if refresh is not None:
            params["refresh"] = refresh
        return self._make_request("GET", "/x/latest", params=params)

    def x_search(
        self,
        q: Optional[str] = None,
        query: Optional[str] = None,
        limit: int = 20,
        refresh: Optional[str] = None,
    ) -> dict:
        params: Dict[str, Any] = {"limit": limit}
        if q is not None:
            params["q"] = q
        if query is not None:
            params["query"] = query
        if refresh is not None:
            params["refresh"] = refresh
        return self._make_request("GET", "/x/search", params=params)

    def x_news(
        self,
        limit: int = 10,
        ai_only: Optional[str] = None,
        with_tweets: Optional[str] = None,
        tabs: Optional[str] = None,
        refresh: Optional[str] = None,
    ) -> dict:
        params: Dict[str, Any] = {"limit": limit}
        if ai_only is not None:
            params["aiOnly"] = ai_only
        if with_tweets is not None:
            params["withTweets"] = with_tweets
        if tabs is not None:
            params["tabs"] = tabs
        if refresh is not None:
            params["refresh"] = refresh
        return self._make_request("GET", "/x/news", params=params)

    def x_about(self, handle: str, refresh: Optional[str] = None) -> dict:
        params = {"refresh": refresh} if refresh is not None else None
        return self._make_request("GET", f"/x/about/{handle}", params=params)

    def x_whoami(self, refresh: Optional[str] = None) -> dict:
        params = {"refresh": refresh} if refresh is not None else None
        return self._make_request("GET", "/x/whoami", params=params)

    def linkedin_profile_posts(self, profile_url: str) -> dict:
        return self._make_request(
            "POST",
            "/linkedin/profile/posts",
            json={"profile_url": profile_url},
        )

    def linkedin_company_posts(self, company_url: str) -> dict:
        return self._make_request(
            "POST",
            "/linkedin/company/posts",
            json={"company_url": company_url},
        )

    def create_api_key(
        self,
        client: str,
        rate_limit: int,
        key_type: Literal["free_trial", "basic", "pro", "growth"],
        description: Optional[str] = None,
    ) -> dict:
        return self._make_request(
            "POST",
            "/admin/create-api-key",
            json={
                "client": client,
                "rate_limit": rate_limit,
                "key_type": key_type,
                "description": description,
            },
        )

    def revoke_api_key(self, key: str) -> dict:
        return self._make_request("DELETE", f"/admin/revoke-api-key/{key}")

    def reset_rate_limits(self) -> dict:
        return self._make_request("POST", "/admin/reset")

    def create_user_api_key(self, key_type: str, jwt_token: str) -> dict:
        return self._make_request(
            "POST",
            "/user/create-api-key",
            params={"key_type": key_type},
            use_api_key=False,
            bearer_token=jwt_token,
        )

    def get_user_api_keys(self, jwt_token: str) -> dict:
        return self._make_request(
            "GET",
            "/user/api-keys",
            use_api_key=False,
            bearer_token=jwt_token,
        )

    def get_subscription_plans(self, jwt_token: Optional[str] = None) -> dict:
        return self._make_request(
            "GET",
            "/subscription-plans",
            use_api_key=False,
            bearer_token=jwt_token,
        )

    def regenerate_user_api_key(self, old_key: str, jwt_token: str) -> dict:
        return self._make_request(
            "POST",
            "/user/regenerate-api-key",
            params={"old_key": old_key},
            use_api_key=False,
            bearer_token=jwt_token,
        )
