import argparse
import time
import re
from typing import Set, Optional
from pathlib import Path
import sys
import traceback
import logging
import contextlib
from io import StringIO
from urllib.parse import quote
import requests
from bs4 import BeautifulSoup
import bibtexparser


logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)-10s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)

GREEN = "\033[92m"
RED = "\033[91m"
YELLOW = "\033[93m"
RESET = "\033[0m"


def get_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Automatically fetch BibTeX entries from ADS, arXiv, and INSPIRE-HEP",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
        python auto_bib_fetcher.py document.log bibliography.bib
        python auto_bib_fetcher.py --scan-tex document.tex bibliography.bib
        python auto_bib_fetcher.py --delay 2.0 document.log bibliography.bib
        """,
    )

    parser.add_argument("input", help="Input file (LaTeX log or TeX source)")
    parser.add_argument("output", help="Output BibTeX file")
    parser.add_argument(
        "--scan-tex",
        action="store_true",
        help="Scan TeX source file instead of log file",
    )
    parser.add_argument(
        "--delay", type=float, default=1.0, help="Delay between requests (seconds)"
    )
    parser.add_argument(
        "--append",
        action="store_true",
        help="Append to existing bibliography file",
        default=True,
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        default=False,
        help="Show verbose output",
    )

    return parser


class BibFetcher:
    """Main class for fetching bibliography entries from various sources."""

    def __init__(self, delay: float = 1.0, verbose: bool = False):
        """Initialize with rate limiting delay."""
        self.delay = delay
        self.verbose = verbose
        self.session = requests.Session()
        self.session.headers.update(
            {"User-Agent": "Auto-Bibliography-Fetcher/1.0 (Academic Research Tool)"}
        )

    def extract_citations_from_log(self, log_file: Path) -> Set[str]:
        """Extract missing citation keys from LaTeX log file."""
        citations = set()

        if not log_file.exists():
            logging.warning(f"{RED}Log file {log_file} not found{RESET}")
            return citations

        with open(log_file, "r", encoding="utf-8", errors="ignore") as f:
            content = f.read()

        # only searching for undefined citations warning generated by LaTeX
        for match in re.finditer(
            r"LaTeX Warning: Citation '([^']+)' on page \d+ undefined on input line",
            content,
        ):
            citations.add(match.group(1).strip())

        return citations

    def extract_citations_from_tex(self, tex_file: Path) -> Set[str]:
        """Extract all citation keys from LaTeX source file."""
        citations = set()

        if not tex_file.exists():
            logging.info(f"Warning: TeX file {tex_file} not found")
            return citations

        with open(tex_file, "r", encoding="utf-8", errors="ignore") as f:
            content = f.read()

        content = re.sub(r"(?<!\\)%.*$", "", content, flags=re.MULTILINE)

        citation_patterns = [
            r"\\cite(?:\[[^\]]*\])?(?:\[[^\]]*\])?\{([^}]+)\}",
            r"\\citep(?:\[[^\]]*\])?(?:\[[^\]]*\])?\{([^}]+)\}",
            r"\\citet(?:\[[^\]]*\])?(?:\[[^\]]*\])?\{([^}]+)\}",
            r"\\autocite(?:\[[^\]]*\])?(?:\[[^\]]*\])?\{([^}]+)\}",
            r"\\parencite(?:\[[^\]]*\])?(?:\[[^\]]*\])?\{([^}]+)\}",
            r"\\textcite(?:\[[^\]]*\])?(?:\[[^\]]*\])?\{([^}]+)\}",
        ]

        for pattern in citation_patterns:
            for match in re.finditer(pattern, content):
                keys = [key.strip() for key in match.group(1).split(",")]
                citations.update(keys)

        return citations

    def get_existing_citations(self, bib_file: Path) -> Set[str]:
        """Extract citation keys from an existing BibTeX file."""
        existing_citations = set()

        if not bib_file.exists():
            return existing_citations

        try:
            with open(bib_file, "r", encoding="utf-8", errors="ignore") as f:
                bib_database = bibtexparser.load(f)
                for entry in bib_database.entries:
                    if "ID" in entry:
                        existing_citations.add(entry["ID"])
        except Exception as e:
            logging.error(f"Warning: Could not parse BibTeX file {bib_file}: {e}")

        return existing_citations

    def identify_citation_source(self, citation_key: str) -> str:
        """Identify which database a citation key belongs to."""
        citation_key = citation_key.strip()

        # arXiv patterns
        if re.match(r"^arXiv:", citation_key) or re.match(
            r"^\d{4}\.\d{4,5}$", citation_key
        ):
            return "arxiv"

        # INSPIRE-HEP patterns (often have specific formats)
        elif re.match(r"^[A-Za-z]+:\d{4}[a-z]{2,3}$", citation_key):
            return "inspire"
        # ADS. I hate ADS. Anyone from ADS reading this, please fix your slow website, you bibtex citation keys.
        else:
            return "ads"

    def fetch_from_ads(self, citation_key: str) -> Optional[str]:
        """Fetch BibTeX entry from NASA ADS."""
        try:
            # Try direct bibcode lookup
            url = f"https://ui.adsabs.harvard.edu/abs/{quote(citation_key)}/exportcitation"
            if self.verbose:
                logging.info(f"{YELLOW}Fetching from adsabs: {url}{RESET}")

            response = self.session.get(url, timeout=10)

            if response.status_code == 200:
                soup = BeautifulSoup(response.text, "html.parser")
                textarea = soup.find("textarea", class_="export-textarea")
                if textarea:
                    return textarea.text.strip()

            # Fallback: search by citation key
            search_url = f"https://ui.adsabs.harvard.edu/search/q={quote(citation_key)}"
            logging.warning(
                f"{RED}Could not fetch directly, try manual lookup: {search_url}{RESET}"
            )

        except Exception as e:
            logging.error(f"Error fetching from ADS for {citation_key}: {e}")
            logging.error(traceback.print_exc())

        return None

    def fetch_from_arxiv(self, citation_key: str) -> Optional[str]:
        """Fetch BibTeX entry from arXiv."""
        try:
            arxiv_id = citation_key.replace("arXiv:", "").strip()
            api_url = f"https://arxiv.org/bibtex/{arxiv_id}"
            if self.verbose:
                logging.info(f"{YELLOW}Fetching from arXiv: {api_url}{RESET}")
            response = self.session.get(api_url, timeout=10)

            if response.status_code == 200:
                content = response.text.strip()
                _replace_key = bibtexparser.loads(content)
                _replace_key.entries[0]["ID"] = citation_key

                output_buffer = StringIO()
                with contextlib.redirect_stdout(output_buffer):
                    bibtexparser.dump(_replace_key, sys.stdout)
                content = output_buffer.getvalue()
                return content

        except Exception as e:
            logging.error(f"Error fetching from arXiv for {citation_key}: {e}")

        return None

    def fetch_from_inspire(self, citation_key: str) -> Optional[str]:
        """Fetch BibTeX entry from INSPIRE-HEP."""
        try:
            api_url = f"https://inspirehep.net/api/literature?q={quote(citation_key)}&format=bibtex"
            if self.verbose:
                logging.info(f"{YELLOW}Fetching from inspirehep: {api_url}{RESET}")
            response = self.session.get(api_url, timeout=10)

            if response.status_code == 200:
                content = response.text.strip()
                if content.startswith("@") and content.endswith("}"):
                    return content

        except Exception as e:
            logging.error(f"Error fetching from INSPIRE for {citation_key}: {e}")
            traceback.print_exc()

        return None

    def fetch_bibtex_entry(self, citation_key: str) -> Optional[str]:
        """Fetch BibTeX entry for a given citation key."""
        source = self.identify_citation_source(citation_key)

        if source == "arxiv":
            result = self.fetch_from_arxiv(citation_key)
        elif source == "inspire":
            result = self.fetch_from_inspire(citation_key)
        elif source == "ads":
            result = self.fetch_from_ads(citation_key)
        else:
            result = None

        # Rate limiting
        time.sleep(self.delay)

        return result

    def process_citations(self, citations: Set[str], output_file: Path, append: bool):
        """Process a set of citations and write to output file."""
        # Check existing citations in the output file
        existing_citations = self.get_existing_citations(output_file)
        citations_to_fetch = citations - existing_citations

        if len(citations_to_fetch) == 0:
            color = GREEN
        else:
            color = YELLOW

        logging.info(
            f"{color}Found {len(citations)} citations, {len(citations_to_fetch)} need fetching: {', '.join(sorted(citations_to_fetch))}{RESET}"
        )

        if not citations_to_fetch:
            logging.info("All citations already exist in the BibTeX file!")
            return

        mode = "a" if append else "w"

        with open(output_file, mode, encoding="utf-8") as f:
            if not append:
                f.write("% Auto-generated bibliography entries\n\n")

            successful = 0
            failed = []

            for citation_key in sorted(citations_to_fetch):
                bibtex_entry = self.fetch_bibtex_entry(citation_key)

                if bibtex_entry:
                    f.write(bibtex_entry)
                    f.write("\n\n")
                    successful += 1
                    logging.info(
                        f"{GREEN}✓ Successfully fetched: {citation_key}{RESET}"
                    )
                else:
                    failed.append(citation_key)
                    logging.warning(f"{RED}✗ Failed to fetch: {citation_key}{RESET}")

            failed_color = RED if len(failed) != 0 else GREEN

            logging.info(
                f"{GREEN}Summary: {successful} successful{RESET}, {failed_color}{len(failed)} failed{RESET}"
            )

            if failed:
                f.write(
                    "% Failed to fetch:\nTry checking the citation keys manually.\n"
                )
                for key in failed:
                    f.write(f"% {key}\n")
                logging.warning(f"{RED}Failed citations: {', '.join(failed)}{RESET}")


def main():
    parser = get_parser()
    args = parser.parse_args()

    fetcher = BibFetcher(delay=args.delay, verbose=args.verbose)
    output_path = Path(args.output)

    if args.scan_tex:
        citations = fetcher.extract_citations_from_tex(Path(args.input))
    else:
        citations = fetcher.extract_citations_from_log(Path(args.input))

    if not citations:
        logging.error(f"{RED}No citations found!{RESET}")
        sys.exit(1)

    fetcher.process_citations(citations, output_path, args.append)

    logging.info(f"{GREEN}BibTeX entries written to: {output_path}{RESET}")


if __name__ == "__main__":
    main()
