#!/usr/bin/env python3
"""graph-ask: find where something lives in a codebase for almost no tokens.

Graphify builds a code graph of the repository locally (tree-sitter, no LLM). Jev
(TypeSafe System One, a fast and very cheap classifier) reads a plain-language
question and ranks every file in scope, then ranks the symbols inside the winning
files. The answer is a short list of "path (score)" plus "L<line> symbol" pointers
and how they connect, so a coding assistant opens the right lines instead of
grepping and reading whole files. Without a key, or with --no-jev, it ranks by
keyword overlap alone. Python standard library only, Python 3.9 or later.

Setup (once per machine):
    uv tool install graphifyy      # or: pip install graphifyy
    graphify update .              # builds graphify-out/graph.json (gitignored)

Key: graph-ask reads TYPESAFE_API_KEY from the environment, or a line
TYPESAFE_API_KEY=... in ~/.typesafe.env. The key is never printed. Set
GRAPH_ASK_ENDPOINT to override the TypeSafe endpoint URL.

Optional config: a .graph-ask.json file at the repository root names the areas of
the codebase, so Jev can first pick the area a question is about:
    {
      "areas": {
        "core":   {"description": "The main application source", "prefixes": ["src/", "lib/"]},
        "docs":   {"description": "Written documentation and specs", "prefixes": ["docs/"]}
      },
      "default_area": "core",
      "test_markers": ["tests/", "test_", "_test.", ".test.", ".spec."]
    }
"default_area" and "test_markers" are optional; the list above is the built-in
default for "test_markers". Markers match case-insensitively, so a path like
Tests/FooTests.swift counts as a test file. With no config, or no "areas", graph-ask treats the
whole repository as one area, skips the area-pick request, and still works; --area
then reports a clear error. Files that match no prefix fall into an area "other".

Usage:
    graph-ask.py "where does the app validate a user login"
    graph-ask.py --files 8 --json "how does the client talk to the server"
    graph-ask.py --area core "where is the request router"

Released under the Unlicense (public domain). No warranty.
"""
import argparse
import concurrent.futures
import json
import math
import os
import re
import subprocess
import sys
import time
import urllib.error
import urllib.request
from collections import Counter, defaultdict

ENDPOINT = os.environ.get("GRAPH_ASK_ENDPOINT", "https://api.typesafe.ai/v1/systemone")
SHARD = 250          # Jev allows 255 options per choice question
ADVANCE = 8          # files that advance from each shard to the final round
PRICE_PER_MTOK = 0.042

SINGLE_AREA = "all"                 # the implicit area when no config names any
OTHER_AREA = "other"                # files that match no configured prefix
DEFAULT_TEST_MARKERS = ("tests/", "test_", "_test.", ".test.", ".spec.")


def load_config(path=".graph-ask.json"):
    """Return (areas, default_area, test_markers). areas maps name -> (description, prefixes)."""
    try:
        with open(path, encoding="utf-8") as handle:
            raw = json.load(handle)
    except OSError:
        return {}, None, DEFAULT_TEST_MARKERS
    except json.JSONDecodeError as error:
        sys.exit(f"could not parse {path}: {error}")
    areas = {}
    for name, spec in (raw.get("areas") or {}).items():
        spec = spec or {}
        areas[name] = (spec.get("description", ""), tuple(spec.get("prefixes") or ()))
    markers = tuple(raw.get("test_markers") or DEFAULT_TEST_MARKERS)
    return areas, raw.get("default_area"), markers


def load_key():
    key = os.environ.get("TYPESAFE_API_KEY")
    if key:
        return key
    try:
        with open(os.path.expanduser("~/.typesafe.env"), encoding="utf-8") as handle:
            match = re.search(r'TYPESAFE_API_KEY="?([^"\n]+)"?', handle.read())
            return match.group(1) if match else None
    except OSError:
        return None


class Jev:
    def __init__(self, key):
        self.key = key
        self.requests = 0
        self.tokens = 0
        self.ms = 0.0

    def ask(self, state, questions):
        body = json.dumps({"state": state, "model": "jev-latest", "questions": questions}).encode()
        request = urllib.request.Request(ENDPOINT, data=body, headers={
            "Authorization": f"Bearer {self.key}", "Content-Type": "application/json"})
        start = time.time()
        try:
            with urllib.request.urlopen(request, timeout=20) as response:
                reply = json.load(response)
        except urllib.error.HTTPError as error:
            raise RuntimeError(f"Jev returned {error.code}: {error.read()[:300].decode(errors='replace')} (request was {len(body)} bytes)") from None
        self.ms += (time.time() - start) * 1000
        self.requests += 1
        self.tokens += len(body) // 4   # estimate; good enough for a cost line
        return reply.get("answers", {})


def words(text):
    text = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", text)
    return [w for w in re.split(r"[^a-z0-9]+", text.lower()) if len(w) > 2]


def is_test(path, markers):
    lowered = path.lower()
    return any(marker.lower() in lowered for marker in markers)


def area_of(path, areas_config):
    if not areas_config:
        return SINGLE_AREA
    for name, (_description, prefixes) in areas_config.items():
        if prefixes and path.startswith(prefixes):
            return name
    return OTHER_AREA


class Graph:
    def __init__(self, path):
        with open(path, encoding="utf-8") as handle:
            raw = json.load(handle)
        self.built_at = raw.get("built_at_commit")
        self.nodes = {n["id"]: n for n in raw["nodes"] if n.get("source_file")}
        self.links = [l for l in raw["links"] if l["source"] in self.nodes and l["target"] in self.nodes]
        self.degree = Counter()
        for link in self.links:
            self.degree[link["source"]] += 1
            self.degree[link["target"]] += 1
        self.by_file = defaultdict(list)
        for node in self.nodes.values():
            self.by_file[node["source_file"]].append(node)
        for members in self.by_file.values():
            members.sort(key=lambda n: -self.degree[n["id"]])

    def symbols(self, path, limit):
        stem = os.path.splitext(os.path.basename(path))[0]
        seen, out = set(), []
        for node in self.by_file[path]:
            label = node["label"]
            if label in seen or label == os.path.basename(path) or label == stem:
                continue
            seen.add(label)
            out.append(node)
            if len(out) == limit:
                break
        return out

    def describe(self, path, limit):
        return {"path": path, "defines": ", ".join(n["label"] for n in self.symbols(path, limit))[:28 * limit]}


def line_of(node):
    match = re.search(r"\d+", node.get("source_location") or "")
    return int(match.group()) if match else 1


def commits_behind(commit):
    if not commit:
        return 0
    try:
        out = subprocess.run(["git", "rev-list", "--count", f"{commit}..HEAD"], capture_output=True, text=True, timeout=5)
        return int(out.stdout.strip() or 0) if out.returncode == 0 else 0
    except (OSError, ValueError, subprocess.SubprocessError):
        return 0


def keyword_rank(graph, question, paths):
    """Fallback and tie breaker: IDF weighted overlap between the question and a file's names."""
    terms = set(words(question))
    bags = {p: Counter(words(p) + [w for n in graph.by_file[p][:40] for w in words(n["label"])]) for p in paths}
    document_frequency = Counter(w for bag in bags.values() for w in bag)
    total = len(paths) or 1
    score = {p: sum(math.log(total / document_frequency[w]) for w in terms if w in bag) for p, bag in bags.items()}
    return sorted(paths, key=lambda p: -score[p])


def rank(jev, question, instructions, options):
    """One Jev choice question over up to 255 options. Returns [(key, probability)] best first."""
    answer = jev.ask(question, {"pick": {"type": "choice", "instructions": instructions, "criteria": options}})["pick"]
    probabilities = answer.get("probabilities") or {answer.get("choice"): answer.get("confidence", 0)}
    return sorted(probabilities.items(), key=lambda item: -item[1])


def main():
    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    parser.add_argument("question")
    parser.add_argument("--graph", default="graphify-out/graph.json")
    parser.add_argument("--files", type=int, default=5, help="files to report (default 5)")
    parser.add_argument("--symbols", type=int, default=6, help="symbols per file (default 6)")
    parser.add_argument("--area", action="append", help="restrict to configured areas; needs .graph-ask.json (repeatable)")
    parser.add_argument("--tests", action="store_true", help="include test files")
    parser.add_argument("--json", action="store_true")
    parser.add_argument("--no-jev", action="store_true", help="keyword ranking only")
    parser.add_argument("--config", default=".graph-ask.json", help="area config file (default .graph-ask.json)")
    args = parser.parse_args()
    if hasattr(sys.stdout, "reconfigure"):   # Windows consoles are not UTF-8 by default
        sys.stdout.reconfigure(encoding="utf-8", errors="replace")

    areas_config, default_area, test_markers = load_config(args.config)
    all_area_names = (list(areas_config) + [OTHER_AREA]) if areas_config else [SINGLE_AREA]

    if args.area:
        if not areas_config:
            sys.exit("--area needs an \"areas\" map in .graph-ask.json, but none was found. "
                     "Remove --area to scan the whole repository.")
        unknown = [a for a in args.area if a not in areas_config]
        if unknown:
            sys.exit(f"unknown area {', '.join(repr(a) for a in unknown)}. "
                     f"Valid areas: {', '.join(sorted(areas_config))}.")

    if not os.path.exists(args.graph):
        sys.exit(f"no graph at {args.graph}. Build it first (about 45 s, local, free): graphify update .")
    started = time.time()
    graph = Graph(args.graph)
    key = None if args.no_jev else load_key()
    jev = Jev(key) if key else None
    question = args.question

    # 1. Which part of the repo, and are tests wanted? One request, only when areas are configured.
    areas, want_tests = args.area, args.tests
    if jev and not areas and areas_config:
        criteria = {}
        for name, (description, _prefixes) in areas_config.items():
            text = description
            if name == default_area:
                text = (text + " This is the default area when the question names no specific area.").strip()
            criteria[name] = text
        answers = jev.ask(question, {
            "area": {"type": "choice", "instructions": "A developer asks this about a codebase. Which part of the repository holds the answer?",
                     "criteria": criteria},
            "tests": {"type": "noul", "instructions": "Is the developer asking about tests, test coverage, or a failing test?"},
        })
        probabilities = answers["area"].get("probabilities") or {answers["area"]["choice"]: 1.0}
        # Unsure which area? Scan more of the repo. A wider scan costs a fraction of a cent.
        top = max(probabilities.values())
        areas = all_area_names if top < 0.5 else [name for name, p in probabilities.items() if p >= 0.1]
        want_tests = want_tests or answers["tests"].get("noul", 0) > 0.6
    areas = areas or all_area_names

    # A question about tests scans every test file in the repo, whatever area was picked.
    paths = [p for p in graph.by_file if os.path.exists(p) and
             ((area_of(p, areas_config) in areas and (want_tests or not is_test(p, test_markers))) or (want_tests and is_test(p, test_markers)))]
    ordered = keyword_rank(graph, question, paths)

    # 2. Every file in the area goes through Jev, in parallel shards. Winners meet in a final round.
    if jev and paths:
        instructions = ("A developer asks a question about a codebase. Each option is one source file with the main things it defines. "
                        "Pick the file most likely to contain the answer.")
        shards = [ordered[i::max(1, math.ceil(len(ordered) / SHARD))] for i in range(max(1, math.ceil(len(ordered) / SHARD)))]

        def run_shard(shard):
            options = {f"f{i}": graph.describe(p, 8) for i, p in enumerate(shard)}
            try:
                ranked = rank(jev, question, instructions, options)
            except RuntimeError as error:   # over Jev's request size cap: halve and retry
                if "max_tokens_exceeded" not in str(error) or len(shard) < 2:
                    raise
                return run_shard(shard[::2]) + run_shard(shard[1::2])
            return [(shard[int(k[1:])], prob) for k, prob in ranked[:ADVANCE] if k.startswith("f")]

        finalists = []
        if len(shards) == 1:
            finalists = [p for p, _ in run_shard(shards[0])]
        else:
            with concurrent.futures.ThreadPoolExecutor(max_workers=6) as pool:
                for result in pool.map(run_shard, shards):
                    finalists += [p for p, _ in result]
        # Guard against a wrong area pick: the best keyword matches from the rest of the repo join the final.
        outside = [p for p in graph.by_file if p not in set(paths) and (want_tests or not is_test(p, test_markers) or "test" in question.lower())]
        finalists += keyword_rank(graph, question, outside)[:16]
        finalists = list(dict.fromkeys(finalists))[:120]
        options = {f"f{i}": graph.describe(p, 20) for i, p in enumerate(finalists)}
        final = rank(jev, question, instructions, options)
        chosen = [(finalists[int(k[1:])], prob) for k, prob in final if k.startswith("f")][:args.files]
    else:
        chosen = [(p, 0.0) for p in ordered[:args.files]]

    # 3. Inside the winning files, which symbols matter? One request for all of them.
    picked = {}
    pool_nodes = [n for path, _ in chosen for n in graph.symbols(path, 250 // max(1, len(chosen)))]
    if jev and pool_nodes:
        options = {f"s{i}": {"name": n["label"], "file": os.path.basename(n["source_file"])} for i, n in enumerate(pool_nodes)}
        ranked = rank(jev, question, "Each option is a function, type or heading in a codebase. Pick the one most relevant to the developer's question.", options)
        for k, prob in ranked:
            node = pool_nodes[int(k[1:])]
            bucket = picked.setdefault(node["source_file"], [])
            if len(bucket) < args.symbols and (prob >= 0.01 or len(bucket) < 2):
                bucket.append(node)
    for path, _ in chosen:
        picked.setdefault(path, graph.symbols(path, args.symbols))

    # 4. How the chosen symbols connect to each other across files.
    chosen_ids = {n["id"]: n for nodes in picked.values() for n in nodes}
    edges = [(chosen_ids[l["source"]], l["relation"], chosen_ids[l["target"]]) for l in graph.links
             if l["source"] in chosen_ids and l["target"] in chosen_ids and l["relation"] not in ("contains", "method")
             and chosen_ids[l["source"]]["source_file"] != chosen_ids[l["target"]]["source_file"]][:12]

    stats = {"areas": areas, "tests": want_tests, "files_scanned": len(paths), "seconds": round(time.time() - started, 2),
             "jev_requests": jev.requests if jev else 0, "jev_tokens_est": jev.tokens if jev else 0,
             "jev_cost_usd_est": round((jev.tokens if jev else 0) * PRICE_PER_MTOK / 1e6, 5)}
    if args.json:
        print(json.dumps({"files": [{"path": p, "score": round(s, 3), "symbols": [{"name": n["label"], "line": line_of(n)} for n in picked[p]]}
                                    for p, s in chosen],
                          "edges": [[a["label"], r, b["label"]] for a, r, b in edges], "stats": stats}, indent=1))
        return
    for path, score in chosen:
        print(f"{path}  ({score:.2f})" if jev else path)
        for node in sorted(picked[path], key=line_of):
            print(f"  L{line_of(node)}  {node['label']}")
    if edges:
        print("links:")
        for a, relation, b in edges:
            print(f"  {a['label']} --{relation}--> {b['label']}")
    behind = commits_behind(graph.built_at)
    if behind:
        print(f"note: the graph is {behind} commits behind HEAD. Refresh it with: graphify update .")
    print(f"[{'+'.join(areas)}{' +tests' if want_tests else ''}: {stats['files_scanned']} files ranked in {stats['seconds']} s, "
          f"{stats['jev_requests']} Jev requests, about ${stats['jev_cost_usd_est']}]")


if __name__ == "__main__":
    main()