"""PubChem search engine for chemical compound information."""

from typing import Any, Dict, List, Optional
from urllib.parse import quote

from langchain_core.language_models import BaseLLM
from loguru import logger

from ...constants import USER_AGENT
from ...security.safe_requests import safe_get
from ..rate_limiting import RateLimitError
from ..search_engine_base import BaseSearchEngine


class PubChemSearchEngine(BaseSearchEngine):
    """
    PubChem search engine for chemical compound information.

    Provides access to chemical structures, properties, and bioactivity data.
    No authentication required.
    """

    is_public = True
    is_generic = False
    is_scientific = True
    is_code = False
    is_lexical = True
    needs_llm_relevance_filter = True

    def __init__(
        self,
        max_results: int = 10,
        include_synonyms: bool = True,
        llm: Optional[BaseLLM] = None,
        max_filtered_results: Optional[int] = None,
        settings_snapshot: Optional[Dict[str, Any]] = None,
        **kwargs,
    ):
        """
        Initialize the PubChem search engine.

        Args:
            max_results: Maximum number of search results
            include_synonyms: Whether to include compound synonyms
            llm: Language model for relevance filtering
            max_filtered_results: Maximum results after filtering
            settings_snapshot: Settings snapshot for thread context
        """
        super().__init__(
            llm=llm,
            max_filtered_results=max_filtered_results,
            max_results=max_results,
            settings_snapshot=settings_snapshot,
            **kwargs,
        )

        self.include_synonyms = include_synonyms
        self.base_url = "https://pubchem.ncbi.nlm.nih.gov/rest/pug"
        self.autocomplete_url = (
            "https://pubchem.ncbi.nlm.nih.gov/rest/autocomplete"
        )

        # User-Agent header for API requests
        self.headers = {"User-Agent": USER_AGENT}

    def _search_compounds(self, query: str) -> List[str]:
        """Search for compound names matching the query."""
        try:
            url = (
                f"{self.autocomplete_url}/compound/{quote(query, safe='')}/json"
            )
            params = {"limit": self.max_results * 2}  # Get extra for filtering

            response = safe_get(
                url, params=params, headers=self.headers, timeout=30
            )
            self._raise_if_rate_limit(response.status_code)
            response.raise_for_status()
            data = response.json()

            terms: list[str] = data.get("dictionary_terms", {}).get(
                "compound", []
            )
            return terms

        except RateLimitError:
            raise
        except Exception:
            logger.exception("PubChem autocomplete search failed")
            return []

    def _get_compound_by_name(self, name: str) -> Optional[Dict[str, Any]]:
        """Get compound information by name."""
        try:
            self.rate_tracker.apply_rate_limit(self.engine_type)
            # Get CID first
            url = f"{self.base_url}/compound/name/{quote(name, safe='')}/cids/JSON"
            response = safe_get(url, headers=self.headers, timeout=30)

            if response.status_code == 404:
                return None
            self._raise_if_rate_limit(response.status_code)

            response.raise_for_status()
            data = response.json()
            cids = data.get("IdentifierList", {}).get("CID", [])

            if not cids:
                return None

            cid = cids[0]

            # Get compound properties
            properties = self._get_compound_properties(cid)

            # Get compound description
            description = self._get_compound_description(cid)

            return {
                "cid": cid,
                "name": name,
                "properties": properties,
                "description": description,
            }

        except RateLimitError:
            raise
        except Exception:
            logger.exception(f"Error fetching PubChem compound: {name}")
            return None

    def _get_compound_properties(self, cid: int) -> Dict[str, Any]:
        """Get properties for a compound by CID."""
        try:
            self.rate_tracker.apply_rate_limit(self.engine_type)
            properties_list = [
                "MolecularFormula",
                "MolecularWeight",
                "IUPACName",
                "CanonicalSMILES",
                "IsomericSMILES",
                "InChI",
                "InChIKey",
                "XLogP",
                "TPSA",
                "Complexity",
                "Charge",
                "HBondDonorCount",
                "HBondAcceptorCount",
                "RotatableBondCount",
                "HeavyAtomCount",
            ]

            url = f"{self.base_url}/compound/cid/{cid}/property/{','.join(properties_list)}/JSON"
            response = safe_get(url, headers=self.headers, timeout=30)
            self._raise_if_rate_limit(response.status_code)
            response.raise_for_status()
            data = response.json()

            props = data.get("PropertyTable", {}).get("Properties", [])
            return props[0] if props else {}

        except RateLimitError:
            raise
        except Exception:
            logger.exception(f"Error fetching PubChem properties for CID {cid}")
            return {}

    def _get_compound_description(self, cid: int) -> str:
        """Get description for a compound by CID."""
        try:
            self.rate_tracker.apply_rate_limit(self.engine_type)
            url = f"{self.base_url}/compound/cid/{cid}/description/JSON"
            response = safe_get(url, headers=self.headers, timeout=30)

            if response.status_code == 404:
                return ""
            self._raise_if_rate_limit(response.status_code)

            response.raise_for_status()
            data = response.json()

            descriptions = data.get("InformationList", {}).get(
                "Information", []
            )
            for desc in descriptions:
                if desc.get("Description"):
                    return desc.get("Description", "")  # type: ignore[no-any-return]

            return ""

        except RateLimitError:
            raise
        except Exception:
            logger.exception(
                f"Error fetching PubChem description for CID {cid}"
            )
            return ""

    def _get_compound_synonyms(self, cid: int, limit: int = 10) -> List[str]:
        """Get synonyms for a compound by CID."""
        try:
            self.rate_tracker.apply_rate_limit(self.engine_type)
            url = f"{self.base_url}/compound/cid/{cid}/synonyms/JSON"
            response = safe_get(url, headers=self.headers, timeout=30)

            if response.status_code == 404:
                return []
            self._raise_if_rate_limit(response.status_code)

            response.raise_for_status()
            data = response.json()

            info = data.get("InformationList", {}).get("Information", [])
            if info:
                synonyms = info[0].get("Synonym", [])
                return synonyms[:limit]  # type: ignore[no-any-return]
            return []

        except RateLimitError:
            raise
        except Exception:
            logger.exception(f"Error fetching PubChem synonyms for CID {cid}")
            return []

    def _get_previews(self, query: str) -> List[Dict[str, Any]]:
        """
        Get preview information for PubChem compounds.

        Args:
            query: The search query (compound name)

        Returns:
            List of preview dictionaries
        """
        logger.info(f"Getting PubChem previews for query: {query}")

        # Apply rate limiting
        self._last_wait_time = self.rate_tracker.apply_rate_limit(
            self.engine_type
        )

        # Search for matching compound names
        compound_names = self._search_compounds(query)

        if not compound_names:
            # Try direct lookup
            compound = self._get_compound_by_name(query)
            if compound:
                compound_names = [query]
            else:
                logger.info("No PubChem compounds found")
                return []

        logger.info(f"Found {len(compound_names)} potential compounds")

        previews: list[dict[str, Any]] = []
        seen_cids = set()
        for name in compound_names:
            if len(previews) >= self.max_results:
                break

            try:
                compound = self._get_compound_by_name(name)
                if not compound:
                    continue

                cid = compound["cid"]

                # Deduplicate by CID (autocomplete may return
                # case variants like "Caffeine" and "caffeine")
                if cid in seen_cids:
                    continue
                seen_cids.add(cid)
                properties = compound.get("properties", {})
                description = compound.get("description", "")

                # Build compound URL
                compound_url = (
                    f"https://pubchem.ncbi.nlm.nih.gov/compound/{cid}"
                )

                # Get key properties
                molecular_formula = properties.get("MolecularFormula", "")
                molecular_weight = properties.get("MolecularWeight", "")
                iupac_name = properties.get("IUPACName", "")
                smiles = (
                    properties.get("CanonicalSMILES", "")
                    or properties.get("SMILES", "")
                    or properties.get("IsomericSMILES", "")
                    or properties.get("ConnectivitySMILES", "")
                )

                # Get drug-relevant properties
                xlogp = properties.get("XLogP")
                hbond_donors = properties.get("HBondDonorCount")
                hbond_acceptors = properties.get("HBondAcceptorCount")

                # Build snippet
                snippet_parts = []
                if molecular_formula:
                    snippet_parts.append(f"Formula: {molecular_formula}")
                if molecular_weight:
                    snippet_parts.append(f"MW: {molecular_weight}")
                if xlogp is not None:
                    snippet_parts.append(f"XLogP: {xlogp}")
                if hbond_donors is not None or hbond_acceptors is not None:
                    hbond_info = []
                    if hbond_donors is not None:
                        hbond_info.append(f"H-Donors: {hbond_donors}")
                    if hbond_acceptors is not None:
                        hbond_info.append(f"H-Acceptors: {hbond_acceptors}")
                    snippet_parts.append(", ".join(hbond_info))
                if iupac_name:
                    snippet_parts.append(f"IUPAC: {iupac_name}")
                if description:
                    snippet_parts.append(description[:200])
                snippet = ". ".join(snippet_parts)

                preview = {
                    "id": str(cid),
                    "cid": cid,
                    "title": name,
                    "link": compound_url,
                    "snippet": snippet,
                    "molecular_formula": molecular_formula,
                    "molecular_weight": molecular_weight,
                    "iupac_name": iupac_name,
                    "smiles": smiles,
                    "inchi_key": properties.get("InChIKey", ""),
                    "description": description,
                    "source": "PubChem",
                    "_raw": {
                        "properties": properties,
                        "description": description,
                    },
                }

                previews.append(preview)

            except RateLimitError:
                raise
            except Exception:
                logger.exception(f"Error processing PubChem compound: {name}")
                continue

        return previews

    def _get_full_content(
        self, relevant_items: List[Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        """
        Get full content for the relevant PubChem compounds.

        Args:
            relevant_items: List of relevant preview dictionaries

        Returns:
            List of result dictionaries with full content
        """
        logger.info(
            f"Getting full content for {len(relevant_items)} PubChem compounds"
        )

        results = []
        for item in relevant_items:
            result = item.copy()

            cid = item.get("cid")
            if cid and self.include_synonyms:
                # Get synonyms
                synonyms = self._get_compound_synonyms(cid)
                result["synonyms"] = synonyms

            raw = item.get("_raw", {})
            if raw:
                properties = raw.get("properties", {})
                description = raw.get("description", "")

                # Build content summary
                content_parts = []
                content_parts.append(
                    f"Compound: {result.get('title', 'Unknown')}"
                )
                if cid is not None:
                    content_parts.append(f"CID: {cid}")

                if result.get("molecular_formula"):
                    content_parts.append(
                        f"Molecular Formula: {result['molecular_formula']}"
                    )
                if result.get("molecular_weight"):
                    content_parts.append(
                        f"Molecular Weight: {result['molecular_weight']} g/mol"
                    )
                if result.get("iupac_name"):
                    content_parts.append(f"IUPAC Name: {result['iupac_name']}")
                if result.get("smiles"):
                    content_parts.append(f"SMILES: {result['smiles']}")
                if result.get("inchi_key"):
                    content_parts.append(f"InChIKey: {result['inchi_key']}")

                # Additional properties
                if properties.get("XLogP") is not None:
                    content_parts.append(f"XLogP: {properties['XLogP']}")
                if properties.get("TPSA") is not None:
                    content_parts.append(f"TPSA: {properties['TPSA']} Å²")
                if properties.get("HBondDonorCount") is not None:
                    content_parts.append(
                        f"H-Bond Donors: {properties['HBondDonorCount']}"
                    )
                if properties.get("HBondAcceptorCount") is not None:
                    content_parts.append(
                        f"H-Bond Acceptors: {properties['HBondAcceptorCount']}"
                    )

                if result.get("synonyms"):
                    content_parts.append(
                        f"\nSynonyms: {', '.join(result['synonyms'][:5])}"
                    )

                if description:
                    content_parts.append(f"\nDescription: {description}")

                result["content"] = "\n".join(content_parts)

            # Clean up internal fields
            if "_raw" in result:
                del result["_raw"]

            results.append(result)

        return results

    def get_compound(self, cid: int) -> Optional[Dict[str, Any]]:
        """
        Get a specific compound by CID.

        Args:
            cid: The PubChem compound ID

        Returns:
            Compound dictionary or None
        """
        try:
            properties = self._get_compound_properties(cid)
            description = self._get_compound_description(cid)
            synonyms = self._get_compound_synonyms(cid)

            return {
                "cid": cid,
                "properties": properties,
                "description": description,
                "synonyms": synonyms,
            }
        except RateLimitError:
            raise
        except Exception:
            logger.exception(f"Error fetching PubChem compound {cid}")
            return None

    def search_by_formula(self, formula: str) -> List[Dict[str, Any]]:
        """
        Search compounds by molecular formula.

        Args:
            formula: Molecular formula (e.g., "C6H12O6")

        Returns:
            List of matching compounds
        """
        try:
            url = f"{self.base_url}/compound/fastformula/{quote(formula, safe='')}/cids/JSON"
            response = safe_get(url, headers=self.headers, timeout=30)

            if response.status_code == 404:
                return []
            self._raise_if_rate_limit(response.status_code)

            response.raise_for_status()
            data = response.json()
            cids = data.get("IdentifierList", {}).get("CID", [])

            results = []
            for cid in cids[: self.max_results]:
                compound = self.get_compound(cid)
                if compound:
                    results.append(compound)

            return results

        except RateLimitError:
            raise
        except Exception:
            logger.exception(f"Error searching by formula: {formula}")
            return []
