"""Pyparse: A tree-sitter based parser for querying Python programs."""

from __future__ import annotations

import argparse
import json
import pdb
import sys
from pathlib import Path

import tree_sitter_python as tspython
from tree_sitter import Language, Parser, Query, QueryCursor

PY_LANGUAGE = Language(tspython.language())

_DEFAULT_QUERY = """
(module
  (function_definition name: (identifier) @fname
  body: (block
    (comment)*
    (string)?
    (comment)*
    .
    (return_statement
      (call function: (identifier)
        arguments: (argument_list) @_args
      )
    )
    .
    (comment)*
  )
  ) @func
)
"""

_QUERIES = {
    "functions": """(function_definition
  name: (identifier) @func_name)""",
    "classes": """(class_definition
  name: (identifier) @class_name)""",
    "calls": """(call
  function: (identifier) @call_name)""",
    "imports": """[
  (import_statement
    name: (dotted_name) @import_name)
  (import_from_statement
    module_name: (dotted_name) @import_from)
  (import_from_statement
    name: (dotted_name) @import_name)
  ]""",
}


class QueryResult:
    def __init__(self):
        self.matches = []


def parse_file(filepath: Path | str, query_str: str | None = None) -> QueryResult:
    """Parse a Python file and run a tree-sitter query on it."""
    filepath = Path(filepath)
    if not filepath.exists():
        raise FileNotFoundError(f"File not found: {filepath}")

    query_str = _DEFAULT_QUERY

    parser = Parser(PY_LANGUAGE)
    code = filepath.read_text()
    tree = parser.parse(code.encode())
    root = tree.root_node

    query = Query(PY_LANGUAGE, query_str)
    result = QueryResult()

    query_cursor = QueryCursor(query)
    captures_dict = query_cursor.captures(root)
    pdb.set_trace()
    for capture_name, nodes in captures_dict.items():
        for node in nodes:
            text_bytes = node.text
            if text_bytes:
                text = text_bytes.decode()
            else:
                text = ""
            result.matches.append(
                {
                    "capture": capture_name,
                    "text": text,
                    "line": node.start_point[0] + 1,
                    "column": node.start_point[1],
                }
            )

    return result


def cli():
    parser = argparse.ArgumentParser(
        prog="pyparse", description="Query and parse Python programs using tree-sitter"
    )
    for arg in [
        ("file", {"help": "Python file to parse"}),
        ("-q", "--query", {"help": "Custom tree-sitter query string"}),
        (
            "--preset",
            {
                "choices": ["functions", "classes", "calls", "imports"],
                "help": "Use a preset query",
            },
        ),
        ("--json", {"action": "store_true", "help": "Output results as JSON"}),
    ]:
        parser.add_argument(*arg[:-1], **arg[-1])

    args = parser.parse_args()

    query_str = args.query
    if args.preset:
        query_str = _QUERIES[args.preset]

    try:
        result = parse_file(args.file, query_str)
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)

    if args.json:
        print(json.dumps(result.matches, indent=2))
    else:
        for match in result.matches:
            line = match.get("line", "?")
            column = match.get("column", "?")
            capture = match.get("capture", "?")
            text = match.get("text", "?")
            print(f"{line}:{int(column) + 1}\t{capture}\t{text}")


if __name__ == "__main__":
    cli()
