import argparse
import re
import sys
from . import __version__
from .core import (
    process_latex_source,
    get_default_cache_dir,
    list_sections,
    extract_section,
    extract_arxiv_id,
    count_tokens,
    parse_section_tree,
    format_section_tree,
    find_all_by_name,
)

def main():
    default_cache = str(get_default_cache_dir())

    parser = argparse.ArgumentParser(
        description="Download and display LaTeX source from arXiv papers or process local TeX files."
    )
    parser.add_argument("--version", action="version", version=f"%(prog)s {__version__}")
    parser.add_argument(
        "arxiv_id",
        nargs="?",
        default=None,
        help="The arXiv ID (e.g. 2303.08774) or URL (e.g. https://arxiv.org/abs/2303.08774). Not needed if --local-folder is provided."
    )
    parser.add_argument(
        "--no-comments",
        action="store_true",
        help="Remove LaTeX comments from the output"
    )
    parser.add_argument(
        "--cache-dir",
        type=str,
        help=f"Custom directory to store downloaded files (default: {default_cache})",
        default=None
    )
    parser.add_argument(
        "--no-appendix",
        action="store_true",
        help="Remove the appendix section and everything after it"
    )
    parser.add_argument(
        "--local-folder",
        type=str,
        help="Path to a local folder containing TeX files (alternative to arxiv_id)",
        default=None
    )
    parser.add_argument(
        "--list-sections",
        action="store_true",
        help="List all section names in the document"
    )
    parser.add_argument(
        "--section",
        type=str,
        action="append",
        help="Extract only the specified section(s). Can be used multiple times."
    )
    parser.add_argument(
        "-c", "--copy",
        action="store_true",
        default=False,
        help="Copy the output to the clipboard instead of printing it.",
    )
    parser.add_argument(
        "--force-download",
        action="store_true",
        default=False,
        help="Re-download the paper from arXiv even if a cached version exists.",
    )
    parser.add_argument(
        "--lock-timeout",
        type=float,
        default=120.0,
        help="Seconds to wait for the per-paper cache lock when another process is downloading",
    )
    parser.add_argument(
        "--figure-paths",
        action="store_true",
        default=False,
        help="Output resolved figure file paths instead of LaTeX text",
    )

    parser.add_argument(
        "--abstract",
        action="store_true",
        default=False,
        help="Output only the abstract text",
    )

    parser.add_argument(
        "--expand-macros",
        action="store_true",
        default=False,
        help="Expand \\newcommand and related macro definitions inline",
    )

    parser.add_argument(
        "--token-count",
        action="store_true",
        default=False,
        help="Print the tiktoken token count to stderr (requires pip install 'arxiv-to-prompt[tokens]')",
    )

    args = parser.parse_args()

    # Validate that either arxiv_id or local_folder is provided
    if not args.arxiv_id and not args.local_folder:
        parser.error("Either provide an arXiv ID or use --local-folder to specify a local folder")

    if args.arxiv_id and args.local_folder:
        parser.error("Cannot specify both arXiv ID and --local-folder")

    # Validate incompatible flag combinations
    if args.abstract and args.figure_paths:
        parser.error("Cannot use both --abstract and --figure-paths")
    if args.abstract and args.no_comments:
        parser.error("--abstract already excludes comments to avoid extracting commented-out abstracts")
    if args.abstract and args.no_appendix:
        parser.error("Cannot use both --abstract and --no-appendix")
    if args.abstract and args.section:
        parser.error("Cannot use both --abstract and --section")
    if args.abstract and args.list_sections:
        parser.error("Cannot use both --abstract and --list-sections")
    if args.figure_paths and args.section:
        parser.error("Cannot use both --figure-paths and --section")
    if args.figure_paths and args.list_sections:
        parser.error("Cannot use both --figure-paths and --list-sections")

    arxiv_id = extract_arxiv_id(args.arxiv_id) if args.arxiv_id else None

    content = process_latex_source(
        arxiv_id=arxiv_id,
        keep_comments=not args.no_comments,
        cache_dir=args.cache_dir,
        use_cache=not args.force_download,
        remove_appendix_section=args.no_appendix,
        local_folder=args.local_folder,
        lock_timeout_seconds=args.lock_timeout,
        figure_paths_only=args.figure_paths,
        abstract_only=args.abstract,
        expand_macros_flag=args.expand_macros,
    )
    if not content:
        return

    if args.list_sections:
        tree = parse_section_tree(content)
        output = format_section_tree(tree)
    elif args.section:
        tree = parse_section_tree(content)
        extracted = []
        for section_path in args.section:
            # Check for ambiguity only if not using path notation
            if " > " not in section_path:
                matching_paths = find_all_by_name(tree, section_path)
                if len(matching_paths) > 1:
                    print(f"Warning: '{section_path}' is ambiguous. Found at:", file=sys.stderr)
                    for path in matching_paths:
                        print(f"  - {path}", file=sys.stderr)
                    print("Use path notation to disambiguate.", file=sys.stderr)
                    continue

            section_content = extract_section(content, section_path)
            if section_content:
                extracted.append(section_content)
            else:
                print(f"Warning: Section '{section_path}' not found", file=sys.stderr)
        output = "\n\n".join(extracted) if extracted else None
    else:
        output = content

    if output is None:
        return

    if args.token_count:
        try:
            tokens = count_tokens(output)
            print(f"Token count: {tokens}", file=sys.stderr)
        except ImportError as e:
            print(f"Error: {e}", file=sys.stderr)
            return

    if args.copy:
        try:
            import pyperclip
            pyperclip.copy(output)
            print("Copied to clipboard.", file=sys.stderr)
        except Exception as e:
            print(f"Error: Could not copy to clipboard: {e}", file=sys.stderr)
    else:
        print(output)

if __name__ == "__main__":
    main()
