#!/usr/bin/env python3
"""
Count unique request.security() calls in Pine Script files.

TV counts UNIQUE (symbol, timeframe) pairs across the script + all imported
libraries combined.  This tool handles:
  - Direct request.security() calls
  - Calls through thin wrapper functions (e.g. _get_security("SYM", "D", close))
  - Cross-file deduplication when using --strategy mode

TV budget:
  - 40 unique pairs on the basic plan
  - 64 on higher plans

Usage:
    # Full budget view: strategy + all local libraries combined (main usage)
    python3 tools/count_pine_security_calls.py --strategy strategies/strategy_mlp_scores.pine

    # Single file
    python3 tools/count_pine_security_calls.py strategies/strategy_mlp_scores.pine

    # All .pine files in strategies/
    python3 tools/count_pine_security_calls.py

    # Summary only (no per-call listing)
    python3 tools/count_pine_security_calls.py --summary

    # Cross-file duplicate analysis
    python3 tools/count_pine_security_calls.py --cross
"""

import re
import sys
import glob
from collections import Counter
from pathlib import Path

BUDGET             = 40
HIGHER_PLAN_BUDGET = 64
LIBRARIES_DIR      = Path("libraries")

IMPORT_RE   = re.compile(r'^\s*import\s+(\S+)', re.MULTILINE)
LIB_NAME_RE = re.compile(r'/([^/]+)/\d+$')


# ── Text helpers ───────────────────────────────────────────────────────────────

def _strip_line_comments(text: str) -> str:
    """Remove // single-line comments, preserving line count."""
    result = []
    for line in text.splitlines(keepends=True):
        stripped = re.sub(r'//.*', '', line)
        result.append(stripped)
    return ''.join(result)


def _is_commented(text: str, pos: int) -> bool:
    line_start = text.rfind('\n', 0, pos) + 1
    return text[line_start:pos].lstrip().startswith('//')


def _extract_arg(text: str, pos: int) -> tuple[str, int]:
    """Extract one Pine argument starting at pos, return (value, end_pos)."""
    while pos < len(text) and text[pos] in ' \t\n':
        pos += 1
    if pos >= len(text):
        return '<unknown>', pos
    if text[pos] == '"':
        end = text.index('"', pos + 1)
        return text[pos + 1:end], end + 1
    else:
        depth = 0
        start = pos
        while pos < len(text):
            ch = text[pos]
            if ch in '([':
                depth += 1
            elif ch in ')]':
                if depth == 0:
                    break
                depth -= 1
            elif ch == ',' and depth == 0:
                break
            pos += 1
        return text[start:pos].strip(), pos


def _advance_past_comma(text: str, pos: int) -> int:
    """Advance pos past optional whitespace + comma; return new pos."""
    i = pos
    while i < len(text) and text[i] in ' \t\n':
        i += 1
    if i < len(text) and text[i] == ',':
        return i + 1
    return pos  # no comma found, leave pos unchanged


# ── Wrapper detection ──────────────────────────────────────────────────────────

def _find_wrapper_functions(text: str) -> set[str]:
    """
    Return names of functions whose body (excluding comments) contains exactly
    one request.security() call.  These are thin wrappers.
    """
    sec_re = re.compile(r'\brequest\.security\s*\(')
    clean  = _strip_line_comments(text)
    wrappers: set[str] = set()

    # Find all top-level function definitions (column 0)
    fn_starts = list(re.finditer(r'^([a-zA-Z_][a-zA-Z0-9_]*)\s*\(', clean, re.MULTILINE))
    for i, m in enumerate(fn_starts):
        fname = m.group(1)
        # Locate the '=>' that terminates the parameter list
        arrow = clean.find('=>', m.start())
        if arrow == -1:
            continue
        # Sanity: the next fn definition must start AFTER this arrow
        if i + 1 < len(fn_starts) and fn_starts[i + 1].start() < arrow:
            continue
        body_start = arrow + 2
        # Body ends at next top-level identifier (column 0 after newline)
        next_top = re.search(r'\n[a-zA-Z_]', clean[body_start:])
        body_end  = body_start + next_top.start() + 1 if next_top else len(clean)
        body      = clean[body_start:body_end]
        matches   = sec_re.findall(body)
        if len(matches) == 1:
            wrappers.add(fname)
    return wrappers


# ── File parser ────────────────────────────────────────────────────────────────

def parse_file(path: Path) -> dict:
    text     = path.read_text(encoding='utf-8', errors='replace')
    sec_re   = re.compile(r'\brequest\.security\s*\(')
    calls: list[dict] = []

    # Step 1 — direct request.security() calls (skip commented lines)
    for m in sec_re.finditer(text):
        if _is_commented(text, m.start()):
            continue
        lineno = text[:m.start()].count('\n') + 1
        cursor = m.end()
        symbol, cursor = _extract_arg(text, cursor)
        cursor = _advance_past_comma(text, cursor)
        timeframe, _ = _extract_arg(text, cursor)
        # Skip wrapper definition bodies where symbol is a bare variable
        # (no ':' means it's not a TV ticker; not a known dynamic symbol either)
        is_tv_symbol = ':' in symbol
        is_dynamic   = symbol in ('syminfo.tickerid',)
        if not (is_tv_symbol or is_dynamic):
            continue
        calls.append({
            'line': lineno, 'symbol': symbol, 'timeframe': timeframe,
            'key': _canonical_key(symbol, timeframe),
            'commented': False, 'via': 'direct',
        })

    # Step 2 — wrapper call sites
    wrappers = _find_wrapper_functions(text)
    for fname in wrappers:
        pat = re.compile(r'(?<![a-zA-Z0-9_])' + re.escape(fname) + r'\s*\(')
        for m in pat.finditer(text):
            if _is_commented(text, m.start()):
                continue
            lineno = text[:m.start()].count('\n') + 1
            cursor = m.end()
            symbol, cursor = _extract_arg(text, cursor)
            cursor = _advance_past_comma(text, cursor)
            timeframe, _ = _extract_arg(text, cursor)

            # Skip the function definition itself (args look like type declarations)
            if any(kw in symbol for kw in
                   ('series ', 'simple ', 'int ', 'float ', 'string ', 'bool ')):
                continue
            # Only count calls where symbol is a TV ticker (contains ':')
            # or is a well-known dynamic symbol variable
            is_tv_symbol   = ':' in symbol
            is_dynamic     = symbol in ('syminfo.tickerid',)
            if not (is_tv_symbol or is_dynamic):
                continue

            calls.append({
                'line': lineno, 'symbol': symbol, 'timeframe': timeframe,
                'key': _canonical_key(symbol, timeframe),
                'commented': False, 'via': fname,
            })

    imports = IMPORT_RE.findall(text)
    return {'path': path, 'calls': calls, 'imports': imports, 'wrappers': wrappers}


def _canonical_key(symbol: str, timeframe: str) -> tuple[str, str]:
    """
    Normalize (symbol, timeframe) for deduplication.
    TV resolves runtime variable values — common local variable names for
    daily data ('period', 'tf', 'timeframe') resolve to 'D' in this codebase.
    """
    tf = timeframe.strip()
    if tf in ('period', 'tf', 'i_tf'):
        tf = 'D'
    return (symbol, tf)


# ── Dedup ──────────────────────────────────────────────────────────────────────

def _dedup(calls: list[dict]) -> tuple[list, list]:
    seen: dict = {}
    unique, dupes = [], []
    for c in calls:
        k = c['key']
        if k not in seen:
            seen[k] = c
            unique.append(c)
        else:
            dupes.append((c, seen[k]))
    return unique, dupes


# ── Library resolution ─────────────────────────────────────────────────────────

def _import_to_local_path(import_str: str):
    m = LIB_NAME_RE.search(import_str)
    if not m:
        return None
    name = m.group(1)
    for f in LIBRARIES_DIR.glob('*.pine'):
        if f.stem.lower() == name.lower():
            return f
    return None


# ── Strategy mode ──────────────────────────────────────────────────────────────

def strategy_budget_report(strategy_path: Path, summary_only: bool = False) -> None:
    """Parse strategy + all imported local libraries; show combined TV budget."""
    result      = parse_file(strategy_path)
    strat_calls = result['calls']
    strat_unique, _ = _dedup(strat_calls)

    lib_results: list[tuple[str, dict]] = []
    unmapped:    list[str] = []
    for imp in result['imports']:
        local = _import_to_local_path(imp)
        if local:
            lib_results.append((imp, parse_file(local)))
        else:
            unmapped.append(imp)

    all_calls = list(strat_calls)
    for _, lr in lib_results:
        all_calls.extend(lr['calls'])

    unique_all, _ = _dedup(all_calls)
    total     = len(unique_all)
    remaining = BUDGET - total
    status    = 'OK ✅' if remaining >= 0 else '⚠  OVER BUDGET'

    print(f"\n{'═'*65}")
    print(f"TV Budget Analysis: {strategy_path.name}")
    print(f"{'═'*65}")
    print(f"  Strategy direct calls           : {len(strat_unique)}")
    for imp, lr in lib_results:
        lib_unique, _ = _dedup(lr['calls'])
        if not lib_unique:
            continue
        m    = LIB_NAME_RE.search(imp)
        name = m.group(1) if m else imp
        print(f"  {name:<38}: {len(lib_unique)} unique calls")
    cross_saved = len(all_calls) - total
    if cross_saved:
        print(f"  Cross-file duplicates (saved)   : {cross_saved}")
    print(f"  {'─'*55}")
    print(f"  TOTAL unique calls              : {total} / {BUDGET}  [{status}]")
    if remaining < 0:
        print(f"  ⚡ Need to remove               : {-remaining} call(s) to reach limit")
    else:
        print(f"  Slots remaining                 : {remaining}")
    print(f"  (Higher plan limit              : {HIGHER_PLAN_BUDGET})")

    if not summary_only:
        print(f"\n  All unique (symbol, timeframe) pairs counted by TV:")
        for i, c in enumerate(unique_all, 1):
            src     = c.get('via', 'direct')
            src_tag = f"  [via {src}]" if src != 'direct' else ''
            print(f"    {i:3d}. {c['symbol']!r:45s} tf={c['timeframe']}{src_tag}")

        if unmapped:
            print(f"\n  ⚠  Imports not found locally (calls not counted):")
            for imp in unmapped:
                print(f"       {imp}")

        if lib_results:
            print(f"\n{'─'*65}")
            print("Per-library breakdown (non-zero only):")
            for imp, lr in lib_results:
                lib_unique, _ = _dedup(lr['calls'])
                if not lib_unique:
                    continue
                m    = LIB_NAME_RE.search(imp)
                name = m.group(1) if m else imp
                print(f"\n  {name} ({len(lib_unique)} unique):")
                for c in lib_unique:
                    src     = c.get('via', 'direct')
                    src_tag = f"[via {src}]" if src != 'direct' else ''
                    print(f"    line {c['line']:4d}  {c['symbol']!r:43s} tf={c['timeframe']} {src_tag}")


# ── Single-file report ─────────────────────────────────────────────────────────

def print_report(results: list[dict], summary_only: bool = False,
                 cross: bool = False) -> None:
    for r in results:
        calls  = r['calls']
        unique, dupes = _dedup(calls)

        print(f"\n{'─'*60}\n{r['path']}")
        if r.get('wrappers'):
            print(f"  Wrappers: {', '.join(sorted(r['wrappers']))}")
        print(f"  {len(calls)} effective calls → {len(unique)} unique / {BUDGET} budget")

        if r['imports']:
            print("  Imports:")
            for imp in r['imports']:
                local = _import_to_local_path(imp)
                tag   = f" → {local}" if local else ' → (TV-hosted, not counted locally)'
                print(f"    {imp}{tag}")

        if not summary_only and unique:
            print("  Unique (symbol, tf) pairs:")
            for c in unique:
                src     = c.get('via', 'direct')
                src_tag = f" [via {src}]" if src != 'direct' else ''
                print(f"    line {c['line']:4d}  {c['symbol']!r:45s} tf={c['timeframe']}{src_tag}")

        remaining = BUDGET - len(unique)
        status    = 'OK' if remaining >= 0 else 'OVER BUDGET'
        print(f"  → {remaining} slot(s) remaining  [{status}]")

    if cross and len(results) > 1:
        print(f"\n{'='*60}")
        print("Cross-file duplicate analysis:")
        key_to_files: dict = {}
        cnt: Counter = Counter()
        for r in results:
            unique, _ = _dedup(r['calls'])
            for c in unique:
                k = c['key']
                cnt[k] += 1
                key_to_files.setdefault(k, []).append((r['path'].name, c['line']))
        shared = {k: v for k, v in key_to_files.items() if cnt[k] > 1}
        if shared:
            for (sym, tf), locs in sorted(shared.items()):
                print(f"  {sym!r} tf={tf}")
                for fname, ln in locs:
                    print(f"    {fname}  line {ln}")
        else:
            print("  No duplicate (symbol, tf) pairs across these files.")

    print(f"\n{'='*60}\nAnalysed {len(results)} file(s).")


def main():
    args          = sys.argv[1:]
    summary_only  = '--summary'  in args
    cross         = '--cross'    in args
    strategy_mode = '--strategy' in args
    args = [a for a in args if not a.startswith('--')]

    if strategy_mode:
        sp = Path(args[0]) if args else Path('strategies/strategy_mlp_scores.pine')
        if not sp.exists():
            print(f"Strategy file not found: {sp}", file=sys.stderr)
            sys.exit(1)
        strategy_budget_report(sp, summary_only=summary_only)
        return

    if args:
        paths = [Path(p) for pattern in args for p in glob.glob(pattern)]
    else:
        paths = sorted(Path('strategies').glob('*.pine'))

    if not paths:
        print('No .pine files found.')
        sys.exit(1)

    results = [parse_file(p) for p in paths]
    print_report(results, summary_only=summary_only, cross=cross)


if __name__ == '__main__':
    main()