#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.10"
# dependencies = ["anthropic"]
# ///
"""
pystr - Python string manipulation for the command line

Usage:
    echo "hello world" | pystr 's.upper()'
    cat file.txt | pystr 's.split()[0]'
    cat file.txt | pystr -a 's.replace("\\n", " ")'

Variables available:
    s  - current line (line mode) or entire input (with -a)
    i  - line number (0-indexed, line mode only)
"""
import sys
import argparse
import os
import math
import re

VERSION = "1.1.0"

MODELS = {
    "haiku": "claude-haiku-4-5-20251001",
    "sonnet": "claude-sonnet-4-5-20250929",
    "opus": "claude-opus-4-5-20251101",
}


def build_prompt(task: str, input_sample: str, is_all_mode: bool, is_grep_mode: bool) -> str:
    """Build the prompt for the LLM."""
    # Truncate input sample to avoid huge prompts
    max_sample = 500
    if len(input_sample) > max_sample:
        input_sample = input_sample[:max_sample] + "..."

    # Build variables section
    if is_all_mode:
        variables = "- s: the entire input text (string)"
    else:
        variables = "- s: the current line (string)\n- i: the line number (0-indexed integer)"

    # Build behavior section
    if is_grep_mode:
        behavior = "The expression is evaluated for each line. Lines where the expression is truthy are printed."
    elif is_all_mode:
        behavior = "The expression will be evaluated and its result printed automatically."
    else:
        behavior = "The expression is evaluated for each line and its result printed automatically."

    return f"""You are generating Python code for a command-line text processing tool.

IMPORTANT: Return ONLY executable Python code. Do not return the expected output - return code that produces the output. Do not use import statements - math and re modules are already available.

You can write readable multi-line code with helper functions if it helps clarity. The last line must be an expression that produces the result.

Example of multi-line code:
```
def is_valid(x):
    pattern = r'^[0-9]+$'
    return bool(re.match(pattern, x))
is_valid(s)
```

Available variables and modules:
{variables}
- math: the math module (e.g., math.sqrt, math.floor, math.pi)
- re: the re module for regex (e.g., re.match, re.search, re.sub, re.findall)
- All Python builtins are available (len, int, str, list, sum, sorted, filter, map, etc.)

{behavior}

Sample input:
{input_sample}

Return ONLY Python code, no explanation, no markdown code fences.

Task: {task}"""


def generate_code(prompt: str, model: str) -> str:
    """Call Claude API to generate code."""
    import anthropic

    api_key = os.environ.get("ANTHROPIC_API_KEY")
    if not api_key:
        print("Error: ANTHROPIC_API_KEY environment variable not set", file=sys.stderr)
        print("Get your API key at https://console.anthropic.com/", file=sys.stderr)
        sys.exit(1)

    client = anthropic.Anthropic()

    try:
        message = client.messages.create(
            model=model,
            max_tokens=5000,
            messages=[{"role": "user", "content": prompt}],
        )
        code = message.content[0].text.strip()
        # Remove markdown code blocks if present
        if code.startswith("```"):
            lines = code.split("\n")
            # Remove first line (```python or ```) and last line (```)
            lines = [l for l in lines if not l.startswith("```")]
            code = "\n".join(lines).strip()
        return code
    except anthropic.APIError as e:
        print(f"Error calling Claude API: {e}", file=sys.stderr)
        sys.exit(1)


def main():
    parser = argparse.ArgumentParser(
        description="Python string manipulation for the command line",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
BEHAVIOR FLAGS (can be combined):
  (default)   Evaluate expression per line. One output line per input line.
  -a, --all   s = entire stdin instead of current line. Single evaluation, single output.
  -g, --grep  Print original line when expression is truthy (instead of printing result).
  -p, --prompt  Expression is natural language. Claude generates the Python code.
  -n, --no-print  Suppress auto-print (for side-effect expressions).

AVAILABLE VARIABLES:
  s      Current line (line mode) or entire input (-a mode). Always a string.
  i      0-indexed line number (line mode only, not available in -a mode).
  math   The math module (math.sqrt, math.floor, math.pi, etc.)
  re     The re module (re.match, re.search, re.sub, re.findall, etc.)
  All Python builtins are available (len, int, str, list, sum, sorted, etc.)

OUTPUT FORMAT:
  Line mode:  One line per input line. Each line is str(result) of the expression.
  All mode:   Single print of str(result).
  Grep mode:  Original input lines (unchanged) where expression was truthy.
  Errors go to stderr. Exit code 1 on error, 0 on success.

EXAMPLES WITH OUTPUT:
  $ echo "hello world" | pystr 's.upper()'
  HELLO WORLD

  $ printf "alice\\nbob\\ncharlie" | pystr 'f"{i}: {s}"'
  0: alice
  1: bob
  2: charlie

  $ seq 5 | pystr -g 'int(s) % 2 == 0'
  2
  4

  $ printf "line1\\nline2\\nline3" | pystr -a 'len(s.splitlines())'
  3

  $ echo "john,doe,30" | pystr -p "get the second comma-separated field"
  doe
        """,
    )
    parser.add_argument("-v", "--version", action="version", version=f"pystr {VERSION}")
    parser.add_argument("expression", help="Python expression to evaluate (result is printed)")
    parser.add_argument(
        "-a", "--all",
        action="store_true",
        help="Read all input at once instead of line-by-line",
    )
    parser.add_argument(
        "-n", "--no-print",
        action="store_true",
        help="Don't auto-print (use for expressions with side effects)",
    )
    parser.add_argument(
        "-g", "--grep",
        action="store_true",
        help="Grep mode: print lines where expression is truthy",
    )
    parser.add_argument(
        "-p", "--prompt",
        action="store_true",
        help="Interpret expression as natural language prompt for Claude",
    )
    parser.add_argument(
        "--show",
        action="store_true",
        help="Show generated Python code before executing (use with -p)",
    )
    parser.add_argument(
        "--confirm",
        action="store_true",
        help="Ask for confirmation before executing generated code (use with -p)",
    )
    parser.add_argument(
        "--model",
        choices=["haiku", "sonnet", "opus"],
        default="haiku",
        help="Claude model to use for code generation (use with -p, default: haiku)",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Show the LLM prompt without executing (use with -p)",
    )

    args = parser.parse_args()

    # Read input
    stdin_text = sys.stdin.read()

    # Handle prompt mode: translate natural language to Python expression
    expression = args.expression
    if args.prompt:
        llm_prompt = build_prompt(args.expression, stdin_text, args.all, args.grep)

        if args.dry_run:
            print("=== LLM Prompt ===", file=sys.stderr)
            print(llm_prompt, file=sys.stderr)
            print("==================", file=sys.stderr)
            sys.exit(0)

        expression = generate_code(llm_prompt, MODELS[args.model])

        if args.show:
            print(f"Generated code: {expression}", file=sys.stderr)

        if args.confirm:
            print(f"Generated code: {expression}", file=sys.stderr)
            response = input("Execute? [y/N] ")
            if response.lower() not in ("y", "yes"):
                print("Aborted.", file=sys.stderr)
                sys.exit(0)

    # Compile code - try as expression first, then as multi-line code
    is_multiline = False
    try:
        code = compile(expression, "<expression>", "eval")
    except SyntaxError:
        # Try as multi-line code: exec all but last line, eval last line
        lines = expression.strip().split("\n")
        if len(lines) > 1:
            setup_code = "\n".join(lines[:-1])
            final_expr = lines[-1].strip()
            try:
                setup_compiled = compile(setup_code, "<setup>", "exec")
                final_compiled = compile(final_expr, "<expression>", "eval")
                code = (setup_compiled, final_compiled)
                is_multiline = True
            except SyntaxError as e:
                if args.prompt:
                    print(f"Generated invalid code:\n{expression}", file=sys.stderr)
                    print(f"Syntax error: {e}", file=sys.stderr)
                else:
                    print(f"Syntax error: {e}", file=sys.stderr)
                sys.exit(1)
        else:
            if args.prompt:
                print(f"Generated invalid code: {expression}", file=sys.stderr)
            print(f"Syntax error in expression", file=sys.stderr)
            sys.exit(1)

    def run_code(code, context, line_info=""):
        """Run code with error handling."""
        try:
            if is_multiline:
                setup_compiled, final_compiled = code
                exec(setup_compiled, context)
                return eval(final_compiled, context)
            else:
                return eval(code, context)
        except Exception as e:
            if args.prompt:
                print(f"Generated code failed:\n{expression}", file=sys.stderr)
                print(f"{type(e).__name__}: {e}{line_info}", file=sys.stderr)
            else:
                print(f"{type(e).__name__}: {e}{line_info}", file=sys.stderr)
            sys.exit(1)

    if args.all:
        # Whole input mode
        s = stdin_text
        result = run_code(code, {"s": s, "math": math, "re": re, "__builtins__": __builtins__})
        if not args.no_print:
            print(result)
    else:
        # Line-by-line mode
        lines = stdin_text.splitlines()
        for i, line in enumerate(lines):
            s = line
            result = run_code(code, {"s": s, "i": i, "math": math, "re": re, "__builtins__": __builtins__}, f" (line {i})")
            if args.grep:
                # Grep mode: print line if expression is truthy
                if result:
                    print(s)
            elif not args.no_print:
                print(result)


if __name__ == "__main__":
    main()