from typing import List, Dict, Union, Optional
import requests
import xml.etree.ElementTree as ET


def search_pubmed_ids(
    query: str,
    retmax: int = 50,
    verbose: bool = False
) -> List[str]:
    """
    Searches PubMed (ESearch) for PMIDs matching a query.
    
    :param query: The PubMed query (full syntax allowed).
    :param retmax: Max number of PMIDs to retrieve.
    :param verbose: If True, print debug info.
    :return: List of PMIDs (strings).
    """
    if verbose:
        print(f"[DEBUG] ESearch for query='{query}', retmax={retmax}")
    base_url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi"
    params = {
        "db": "pubmed",
        "term": query,
        "retmode": "xml",
        "retmax": retmax
    }
    try:
        response = requests.get(base_url, params=params, timeout=30)
        response.raise_for_status()
    except requests.exceptions.RequestException as e:
        if verbose:
            print(f"[ERROR] ESearch failed: {e}")
        return []

    try:
        root = ET.fromstring(response.text)
        id_list = root.find("IdList")
        if id_list is None:
            if verbose:
                print("[DEBUG] No IdList in ESearch response.")
            return []
        pmids = [elem.text for elem in id_list.findall("Id") if elem.text]  # type: ignore
        if verbose:
            print(f"[DEBUG] Found {len(pmids)} PMIDs.")
        return pmids
    except ET.ParseError as e:
        if verbose:
            print(f"[ERROR] Parsing ESearch XML failed: {e}")
        return []


def fetch_pubmed_details_bulk(
    pmids: List[str],
    verbose: bool = False
) -> List[Dict[str, Union[str, List[Dict[str, Union[str, List[str]]]]]]]:
    """
    Fetches article metadata in one EFetch call for multiple PMIDs.
    
    :param pmids: The list of PMIDs to fetch.
    :param verbose: Print debug info if True.
    :return: A list of article dicts with minimal structure:
             [
               {
                 "pmid": str,
                 "title": str,
                 "pub_date": str,
                 "authors": [
                   {
                     "name": str,
                     "affiliations": [str, ...]
                   },
                   ...
                 ],
                 "abstract": str,
               },
               ...
             ]
    """
    if not pmids:
        return []
    if verbose:
        print(f"[DEBUG] EFetch for {len(pmids)} PMIDs.")
    base_url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi"
    joined_pmids = ",".join(pmids)
    params = {
        "db": "pubmed",
        "id": joined_pmids,
        "retmode": "xml"
    }
    try:
        response = requests.get(base_url, params=params, timeout=60)
        response.raise_for_status()
    except requests.exceptions.RequestException as e:
        if verbose:
            print(f"[ERROR] EFetch request failed: {e}")
        return []

    # Parse
    results = []
    try:
        root = ET.fromstring(response.text)
        articles = root.findall(".//PubmedArticle")
        for art in articles:
            item = {
                "pmid": "",
                "title": "",
                "pub_date": "",
                "authors": [],
                "abstract": ""
            }
            pmid_elem = art.find(".//PMID")
            if pmid_elem is not None and pmid_elem.text:
                item["pmid"] = pmid_elem.text.strip()

            title_elem = art.find(".//ArticleTitle")
            if title_elem is not None and title_elem.text:
                item["title"] = title_elem.text.strip()

            # Publication date
            pub_date_elem = art.find(".//PubDate")
            pub_date_parts = []
            if pub_date_elem is not None:
                year = pub_date_elem.find("Year")
                month = pub_date_elem.find("Month")
                day = pub_date_elem.find("Day")
                if year is not None and year.text:
                    pub_date_parts.append(year.text)
                if month is not None and month.text:
                    pub_date_parts.append(month.text)
                if day is not None and day.text:
                    pub_date_parts.append(day.text)
            if pub_date_parts:
                item["pub_date"] = "-".join(pub_date_parts)

            # Abstract
            abstract_texts = art.findall(".//AbstractText")
            if abstract_texts:
                segments = []
                for ab in abstract_texts:
                    if ab.text:
                        segments.append(ab.text.strip())
                if segments:
                    item["abstract"] = " ".join(segments)

            # Authors
            author_list = []
            art_author_nodes = art.findall(".//Author")
            for author_node in art_author_nodes:
                # Gather name
                fore_name = author_node.find("ForeName")
                last_name = author_node.find("LastName")
                collective_name = author_node.find("CollectiveName")
                name_str = ""
                if collective_name is not None and collective_name.text:
                    name_str = collective_name.text.strip()
                else:
                    fn = fore_name.text.strip() if (fore_name is not None and fore_name.text) else ""
                    ln = last_name.text.strip() if (last_name is not None and last_name.text) else ""
                    name_str = f"{fn} {ln}".strip()

                # Gather affiliations
                affils = []
                aff_info_nodes = author_node.findall(".//AffiliationInfo")
                for aff_info in aff_info_nodes:
                    aff_text = aff_info.find("Affiliation")
                    if aff_text is not None and aff_text.text:
                        affils.append(aff_text.text.strip())

                author_list.append({
                    "name": name_str,
                    "affiliations": affils
                })
            item["authors"] = author_list
            results.append(item)

        if verbose:
            print(f"[DEBUG] EFetch returned {len(results)} articles.")
    except ET.ParseError as e:
        if verbose:
            print(f"[ERROR] Parsing EFetch XML failed: {e}")
    return results


def parse_and_filter_papers(
    articles: List[Dict[str, Union[str, List[Dict[str, Union[str, List[str]]]]]]],
    verbose: bool = False
) -> List[Dict[str, str]]:
    """
    Given a list of articles from `fetch_pubmed_details_bulk`, identify those
    with at least one author affiliated with a pharma/biotech company.
    Collect final data for CSV:
      PubmedID, Title, Publication Date, Non-academic Author(s),
      Company Affiliation(s), Corresponding Author Email

    :param articles: The list of article dicts.
    :param verbose: If True, prints debug info.
    :return: A list of dicts with keys:
       "PubmedID", "Title", "Publication Date",
       "Non-academic Author(s)", "Company Affiliation(s)",
       "Corresponding Author Email"
    """
    # Define some naive keywords
    company_keywords = ["pharma", "pharmaceutical", "biotech", "inc", "ltd", "gmbh", "corp", "co.", "company"]
    academic_keywords = ["university", "univ.", "college", "school", "institute", "academia", "hospital"]

    final_rows = []

    for art in articles:
        pmid = str(art["pmid"])
        title = str(art["title"])
        pub_date = str(art["pub_date"])

        authors = art["authors"]  # type: ignore
        # We'll collect data
        has_company = False
        company_affiliations = set()
        non_academic_authors = set()
        corresponding_email = ""

        for author_info in authors:
            name = author_info["name"]
            affs = author_info["affiliations"]
            # Attempt to find email in these affiliations
            if not corresponding_email:
                for aff in affs:
                    if "@" in aff:
                        # naive approach: pick first substring with "@"
                        for part in aff.split():
                            if "@" in part:
                                # remove punctuation
                                c_email = part.strip("(),;")
                                corresponding_email = c_email
                                break
                    if corresponding_email:
                        break

            # Check if any aff is "company"
            is_company_author = False
            is_academic = False
            for aff in affs:
                lower_aff = aff.lower()
                if any(ck in lower_aff for ck in company_keywords):
                    is_company_author = True
                    company_affiliations.add(aff)
                if any(ak in lower_aff for ak in academic_keywords):
                    is_academic = True

            if is_company_author:
                has_company = True
            # If no academic affiliation found, treat them as non-academic
            if not is_academic:
                non_academic_authors.add(name)

        if has_company:
            # We keep this paper
            row = {
                "PubmedID": pmid,
                "Title": title,
                "Publication Date": pub_date,
                "Non-academic Author(s)": "; ".join(sorted(non_academic_authors)),
                "Company Affiliation(s)": "; ".join(sorted(company_affiliations)),
                "Corresponding Author Email": corresponding_email
            }
            final_rows.append(row)

    if verbose:
        print(f"[DEBUG] parse_and_filter_papers: returning {len(final_rows)} rows with pharma affiliation.")
    return final_rows
