import requests
from typing import Optional, Literal, List
import os
from dotenv import load_dotenv

load_dotenv()

CompanyType = Literal[
    "Public Company",
    "Educational",
    "Self Employed",
    "Government Agency",
    "Non Profit",
    "Self Owned",
    "Privately Held",
    "Partnership",
]


class PersanaClient:
    def __init__(self, api_key: Optional[str] = None):
        self.api_key = api_key or os.getenv("PERSANA_KEY")
        if not self.api_key:
            raise ValueError("PERSANA_KEY is not set")

    def people_search(
        self,
        include_job_titles: Optional[list[str]] = None,
        exclude_job_titles: Optional[list[str]] = None,
        include_companies: Optional[list[str]] = None,
        exclude_companies: Optional[list[str]] = None,
        company_types: Optional[list[CompanyType]] = None,
        company_include_keywords: Optional[list[str]] = None,
        company_exclude_keywords: Optional[list[str]] = None,
        include_industries: Optional[list[str]] = None,
        exclude_industries: Optional[list[str]] = None,
    ):
        """
        Runs a persana people search with the given filter parameters and returns the results
        Args:
            include_job_titles: The job titles to include in the search. (Optional) should stay set to None unless requested
            exclude_job_titles: The job titles to exclude from the search. (Optional) should stay set to None unless requested
            include_companies: The companies to include in the search. (Optional) should stay set to None unless requested
            exclude_companies: The companies to exclude from the search. (Optional) should stay set to None unless requested
            company_type: The type of the company. (Optional) should stay set to None unless requested
            company_include_keywords: The keywords to use to search for the company. (Optional) should stay set to None unless requested
            company_exclude_keywords: The keywords to exclude from the search for the company. (Optional) should stay set to None unless requested
            include_industries: The industries to include in the search. (Optional) should stay set to None unless requested
            exclude_industries: The industries to exclude from the search. (Optional) should stay set to None unless requested
        Returns:
            The results of the people search
        """
        params = {
            "title_includes": include_job_titles,
            "title_excludes": exclude_job_titles,
            "companies_includes": include_companies,
            "companies_excludes": exclude_companies,
            "company_types": company_types,
            "company_keywords_includes": company_include_keywords,
            "company_keywords_excludes": company_exclude_keywords,
            "industries_includes": include_industries,
            "industries_excludes": exclude_industries,
        }
        params = {k: v for k, v in params.items() if v is not None}
        response = requests.post(
            "https://prod.api.persana.ai/api/v1/people/search",
            headers={
                "Content-Type": "application/json",
                "x-api-key": self.api_key,
            },
            json=params,
        )
        json_data = response.json()
        return json_data["data"]["profiles"]


if __name__ == "__main__":
    client = PersanaClient()
    result = client.people_search(
        include_job_titles=["Software Engineer"],
        exclude_job_titles=["CTO"],
        include_companies=["Google"],
        exclude_companies=["Apple"],
        company_types=["Public Company"],
        company_include_keywords=["Software"],
        company_exclude_keywords=["Data"],
        include_industries=["Technology"],
        exclude_industries=["Finance"],
    )
    print(result)


def truncate_profiles(profiles: List[dict]) -> List[dict]:
    new_profiles = []
    experience_data_keep_keys = [
        "title",
        "company_name",
        "company_company_headline",
        "company_hero_image",
        "company_description",
        "company_type",
    ]
    profile_keep_keys = ["name", "headline", "company", "title"]
    for profile in profiles:
        new_experience_data = {
            k: v
            for k, v in profile["experience_data"].items()
            if k in experience_data_keep_keys
        }
        new_profile = {k: v for k, v in profile.items() if k in profile_keep_keys}
        new_profile["experience_data"] = new_experience_data
        new_profiles.append(new_profile)
    return new_profiles