#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Timestamp: "2025-11-12 12:57:38 (ywatanabe)"


"""
Generate AI2 Asta prompt from manuscript files for finding related papers

Functionalities:
  - Extracts title, keywords, authors, and abstract from manuscript .tex files
  - Generates formatted prompt for AI2 Asta
  - Can save to file or print to stdout
  - Supports both co-author paper search and general related paper search
  - Flexible section inclusion/exclusion
  - Optional bibliography statistics
  - Multiple output formats

Dependencies:
  - packages:
    - pathlib
    - re
    - argparse
    - yaml

IO:
  - input-files:
    - 00_shared/title.tex
    - 00_shared/keywords.tex
    - 00_shared/authors.tex
    - 01_manuscript/contents/abstract.tex
    - 00_shared/**/*.bib (optional, for stats)

  - output-files:
    - Prompt text to stdout or file
"""

import argparse
import json
import re
from pathlib import Path
from typing import Any, Dict, List, Set

import yaml
from _logging import getLogger

HEADER = """# Literature Search Request
We are preparing a manuscript with the information provided below.

1. Please identify related papers that may be relevant to our work.
2. Comprehensive results are welcome, as we will evaluate all suggestions for relevance.
3. Your contribution to advancing scientific research is greatly appreciated.
4. If possible, please output as a BibTeX file (.bib).
"""


def load_config(config_path: Path = None) -> Dict[str, Any]:
    """Load configuration from YAML file.

    Args:
        config_path: Path to config file. If None, uses default location.

    Returns:
        Configuration dictionary
    """
    if config_path is None:
        # Default to ../config/config_manuscript.yaml relative to script location
        script_dir = Path(__file__).resolve().parent
        config_path = script_dir.parent.parent / "config" / "config_manuscript.yaml"

    if not config_path.exists():
        raise FileNotFoundError(f"Config file not found: {config_path}")

    with open(config_path, "r") as f:
        config = yaml.safe_load(f)

    return config


def read_tex_content(tex_path: Path) -> str:
    """Read raw content from .tex file, removing comments.

    Args:
        tex_path: Path to .tex file

    Returns:
        Raw tex content without comments (empty string if file doesn't exist)
    """
    if not tex_path.exists():
        return ""

    content = tex_path.read_text(encoding="utf-8")

    # Remove comment lines (lines starting with %)
    lines = content.split("\n")
    lines = [line for line in lines if not line.strip().startswith("%")]

    return "\n".join(lines).strip()


def clean_latex_content(content: str) -> str:
    """Clean LaTeX commands from content, keeping only the text.

    Args:
        content: Raw LaTeX content

    Returns:
        Cleaned text content
    """
    import re

    # Remove PDF bookmarks first
    content = re.sub(r"\\pdfbookmark\[[^\]]*\]\{[^}]*\}\{[^}]*\}", "", content)

    # Remove correlation references
    content = re.sub(r"\\corref\{[^}]*\}", "", content)

    # Remove author numbering like \author[1]{Name}
    content = re.sub(r"\\author\[[^\]]*\]\{([^}]*)\}", r"\1", content)

    # Remove addresses
    content = re.sub(r"\\address\[[^\]]*\]\{[^}]*\}", "", content)

    # Remove cortext
    content = re.sub(r"\\cortext\[[^\]]*\]\{[^}]*\}", "", content)

    # Remove environment markers
    content = re.sub(r"\\begin\{[^}]+\}", "", content)
    content = re.sub(r"\\end\{[^}]+\}", "", content)

    # Remove standalone commands with optional arguments
    content = re.sub(r"\\[a-zA-Z]+\[[^\]]*\]", "", content)

    # Remove common LaTeX commands but keep their content (iteratively)
    # \command{content} -> content
    for _ in range(3):  # Multiple passes for nested commands
        content = re.sub(r"\\[a-zA-Z]+\{([^{}]*)\}", r"\1", content)

    # Remove any remaining backslash commands
    content = re.sub(r"\\[a-zA-Z]+", "", content)

    # Remove special characters
    content = re.sub(r"\\&", "&", content)
    content = re.sub(r"\\_", "_", content)
    content = re.sub(r"\{", "", content)
    content = re.sub(r"\}", "", content)

    # Clean up keyword separators
    content = re.sub(r"\\sep", ", ", content)

    # Clean up multiple spaces and newlines
    content = re.sub(r"\n\s*\n\s*\n+", "\n\n", content)
    content = re.sub(r" +", " ", content)
    content = re.sub(r"^\s+", "", content, flags=re.MULTILINE)

    return content.strip()


def extract_bib_keys(bib_path: Path) -> Set[str]:
    """Extract all citation keys from a .bib file.

    Args:
        bib_path: Path to .bib file

    Returns:
        Set of citation keys found in the file
    """
    if not bib_path.exists():
        return set()

    content = bib_path.read_text(encoding="utf-8")

    # Match @article{key, @book{key, etc.
    pattern = r"@\w+\s*\{\s*([^,\s]+)"
    keys = re.findall(pattern, content)

    return set(keys)


def get_bibliography_stats(bib_dir: Path) -> Dict[str, int]:
    """Get statistics about bibliography files.

    Args:
        bib_dir: Path to bibliography directory

    Returns:
        Dictionary with bibliography statistics
    """
    bib_files = list(bib_dir.rglob("*.bib"))
    all_keys = set()

    for bib_file in bib_files:
        keys = extract_bib_keys(bib_file)
        all_keys.update(keys)

    return {
        "total_references": len(all_keys),
        "total_files": len(bib_files),
    }


def add_title(parts: List[str], title: str) -> None:
    """Add title section to prompt parts.

    Args:
        parts: List to append prompt parts to
        title: Title content (LaTeX will be cleaned)
    """
    title_clean = clean_latex_content(title) if title else ""
    if title_clean:
        parts.append(f"## Title\n{title_clean}")
        parts.append("")


def add_keywords(parts: List[str], keywords: str) -> None:
    """Add keywords section to prompt parts.

    Args:
        parts: List to append prompt parts to
        keywords: Keywords content (LaTeX will be cleaned)
    """
    keywords_clean = clean_latex_content(keywords) if keywords else ""
    if keywords_clean:
        parts.append(f"## Keywords\n{keywords_clean}")
        parts.append("")


def add_authors(parts: List[str], authors: str) -> None:
    """Add authors section to prompt parts.

    Args:
        parts: List to append prompt parts to
        authors: Authors content (LaTeX will be cleaned)
    """
    authors_clean = clean_latex_content(authors) if authors else ""
    if authors_clean:
        parts.append(f"## Authors\n{authors_clean}")
        parts.append("")


def add_abstract(parts: List[str], abstract: str) -> None:
    """Add abstract section to prompt parts.

    Args:
        parts: List to append prompt parts to
        abstract: Abstract content (LaTeX will be cleaned)
    """
    abstract_clean = clean_latex_content(abstract) if abstract else ""
    if abstract_clean:
        parts.append(f"## Abstract\n{abstract_clean}")


def add_citation_info(
    parts: List[str],
    citation_data: Dict[str, Any],
    bib_filters: Dict[str, bool],
) -> None:
    """Add citation information sections to prompt parts.

    Args:
        parts: List to append prompt parts to
        citation_data: Citation data from check_cited_states.py
        bib_filters: Dictionary indicating which citation lists to include
    """
    if not citation_data or not any(bib_filters.values()):
        return

    parts.append("")
    parts.append("---")
    parts.append("")

    details = citation_data.get("details", {})

    # Add cited references
    if bib_filters.get("cited", False):
        cited_refs = details.get("successfully_cited", [])
        if cited_refs:
            parts.append(f"### Already Cited References ({len(cited_refs)})")
            for ref in cited_refs:
                parts.append(f"- `{ref}`")
            parts.append("")

    # Add uncited references
    if bib_filters.get("uncited", False):
        uncited_refs = details.get("uncited_references", [])
        if uncited_refs:
            parts.append(
                f"### Uncited References in Our Bibliography ({len(uncited_refs)})"
            )
            parts.append("*These might be relevant to cite*")
            for ref in uncited_refs:
                parts.append(f"- `{ref}`")
            parts.append("")

    # Add missing references
    if bib_filters.get("missing", False):
        missing_refs = details.get("missing_references", [])
        if missing_refs:
            parts.append(f"### Missing References ({len(missing_refs)})")
            parts.append("*Cited but not in our bibliography - need to find*")
            for ref in missing_refs:
                parts.append(f"- `{ref}`")
            parts.append("")


def generate_ai2_prompt(
    title: str,
    keywords: str,
    authors: str,
    abstract: str,
    sections: List[str] = None,
    citation_data: Dict[str, Any] = None,
    bib_filters: Dict[str, bool] = None,
) -> str:
    """Generate AI2 Asta prompt in markdown format.

    Args:
        title: Paper title
        keywords: Keywords
        authors: Author names
        abstract: Abstract text
        sections: List of sections to include (default: all)
        citation_data: Citation data from check_cited_states.py (optional)
        bib_filters: Dictionary indicating which citation lists to include

    Returns:
        Formatted prompt for AI2 Asta in markdown
    """
    if sections is None:
        sections = ["title", "keywords", "authors", "abstract"]

    if bib_filters is None:
        bib_filters = {}

    # Build the prompt header
    header = HEADER
    parts = [header, ""]

    # Add requested sections
    if "title" in sections:
        add_title(parts, title)
    if "keywords" in sections:
        add_keywords(parts, keywords)
    if "authors" in sections:
        add_authors(parts, authors)
    if "abstract" in sections:
        add_abstract(parts, abstract)

    # Add citation information
    add_citation_info(parts, citation_data, bib_filters)

    return "\n".join(parts).strip()


def main():
    parser = argparse.ArgumentParser(
        description="Generate AI2 Asta prompt from manuscript files",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Generate prompt for related papers (default: all sections, markdown format)
  %(prog)s

  # Only title and abstract
  %(prog)s --title --abstract

  # Everything except authors (for blind review)
  %(prog)s --title --keywords --abstract

  # Include uncited refs to get suggestions for citing them (recommended!)
  %(prog)s --bib-uncited

  # Include all citation info
  %(prog)s --bib-cited --bib-uncited --bib-missing

  # Save to custom file
  %(prog)s --output ai2_prompt.txt

  # Verbose output
  %(prog)s --verbose
        """,
    )

    # Section selection (if none specified, include all)
    parser.add_argument(
        "--title",
        action="store_true",
        help="Include title section",
    )
    parser.add_argument(
        "--keywords",
        action="store_true",
        help="Include keywords section",
    )
    parser.add_argument(
        "--authors",
        action="store_true",
        help="Include authors section",
    )
    parser.add_argument(
        "--abstract",
        action="store_true",
        help="Include abstract section",
    )

    # Bibliography options
    parser.add_argument(
        "--bib-cited",
        action="store_true",
        help="Include list of successfully cited references in the prompt",
    )
    parser.add_argument(
        "--bib-uncited",
        action="store_true",
        help="Include list of uncited references in the prompt",
    )
    parser.add_argument(
        "--bib-missing",
        action="store_true",
        help="Include list of missing references in the prompt",
    )

    # Other options
    parser.add_argument(
        "--config",
        type=Path,
        help="Path to config YAML file (default: auto-detect from script location)",
    )
    parser.add_argument(
        "--output",
        "-o",
        type=Path,
        help="Output file path (default: print to stdout)",
    )
    parser.add_argument(
        "--verbose",
        "-v",
        action="store_true",
        help="Enable verbose logging",
    )

    args = parser.parse_args()

    # Setup logger
    logger = getLogger(__name__, args.verbose)

    # Determine which sections to include
    # If any section flag is specified, include only those
    # Otherwise, include all sections
    section_flags = {
        "title": args.title,
        "keywords": args.keywords,
        "authors": args.authors,
        "abstract": args.abstract,
    }

    if any(section_flags.values()):
        sections = [name for name, enabled in section_flags.items() if enabled]
    else:
        sections = ["title", "keywords", "authors", "abstract"]

    # Load configuration
    load_config(args.config)

    # Get project root (where the config file is located)
    if args.config:
        project_root = args.config.parent.parent
    else:
        script_dir = Path(__file__).resolve().parent
        project_root = script_dir.parent.parent

    # Extract information from manuscript files
    # Use paths relative to project root
    title_path = project_root / "00_shared" / "title.tex"
    keywords_path = project_root / "00_shared" / "keywords.tex"
    authors_path = project_root / "00_shared" / "authors.tex"
    abstract_path = project_root / "01_manuscript" / "contents" / "abstract.tex"

    logger.info("Reading manuscript files...")

    title = read_tex_content(title_path)
    keywords = read_tex_content(keywords_path)
    authors = read_tex_content(authors_path)
    abstract = read_tex_content(abstract_path)

    # Validate required content
    if not title and "title" in sections:
        logger.warning("No title found in 00_shared/title.tex")
    if not abstract and "abstract" in sections:
        logger.warning("No abstract found in 01_manuscript/contents/abstract.tex")

    # Load citation data if any bib options are specified
    citation_data = None
    if args.bib_cited or args.bib_uncited or args.bib_missing:
        citation_json_path = (
            project_root / "00_shared" / "bib_files" / "cited_states.json"
        )
        if citation_json_path.exists():
            try:
                citation_data = json.loads(
                    citation_json_path.read_text(encoding="utf-8")
                )
                logger.info(f"Loaded citation data from {citation_json_path.name}")
            except Exception as e:
                logger.warning(f"Could not load citation data: {e}")
        else:
            logger.warning("Citation data not found. Run check_cited_states.py first.")
            logger.warning(f"Expected: {citation_json_path}")

    # Determine which citation filters are active
    bib_filters = {
        "cited": args.bib_cited,
        "uncited": args.bib_uncited,
        "missing": args.bib_missing,
    }

    # Generate prompt in markdown format
    prompt = generate_ai2_prompt(
        title,
        keywords,
        authors,
        abstract,
        sections,
        citation_data,
        bib_filters,
    )

    # Create structured data for JSON output
    data = {
        "metadata": {
            "title": title,
            "keywords": keywords,
            "authors": authors,
            "abstract": abstract,
        },
        "prompt": prompt,
        "sections_included": sections,
        "citation_filters": bib_filters if any(bib_filters.values()) else None,
        "citation_data": (citation_data.get("summary") if citation_data else None),
    }

    # Always save to default locations unless custom output specified
    default_json_path = project_root / "00_shared" / "ai2_prompt_data.json"
    default_md_path = project_root / "00_shared" / "ai2_prompt.md"
    default_json_path.parent.mkdir(parents=True, exist_ok=True)

    # Save JSON and markdown to default locations (always)
    if not args.output:
        # Save JSON data
        json_str = json.dumps(data, indent=2, ensure_ascii=False)
        default_json_path.write_text(json_str, encoding="utf-8")

        # Save markdown prompt
        default_md_path.write_text(prompt, encoding="utf-8")

    # Output
    if args.output:
        args.output.write_text(prompt, encoding="utf-8")
        logger.success(f"Saved to: {args.output}")
        logger.info("\nNext steps:")
        logger.info("1. Visit https://asta.allen.ai/chat/")
        logger.info(f"2. Paste the prompt from {args.output}")
        logger.info("3. Click 'Export All Citations' to download BibTeX file")
    else:
        print("\n" + "=" * 80)
        print("AI2 ASTA PROMPT")
        print("=" * 80)
        print()
        print(prompt)
        print()
        print("=" * 80)
        logger.info("\nNext steps:")
        logger.info("1. Visit https://asta.allen.ai/chat/")
        logger.info("2. Copy and paste the prompt above")
        logger.info("3. Click 'Export All Citations' to download BibTeX file")
        logger.success(f"Saved to: {default_md_path}")
        logger.success(f"Saved to: {default_json_path}")


if __name__ == "__main__":
    main()

# EOF