#!/usr/bin/env python3
"""Graph-Enabled Index-Map Workspace Intelligence Engine FTS5 [Production Ready]"""

import ast
import fnmatch
import json
import os
import re
import sqlite3
import struct
import sys
import time

LANG_RULES = {
    'rust': {'class': r'\b(?:struct|enum|trait|union)\s+(\w+)', 'func': r'\b(?:async\s+)?(?:pub(?:\([^)]+\))?\s+)?fn\s+(\w+)'},
    'go': {'class': r'\btype\s+(\w+)\s+struct\b', 'func': r'\bfunc\s+(?:\([^)]+\)\s+)?(\w+)'},
    'lua': {'class': r'\b(\w+)\s*=\s*\{\}', 'func': r'\bfunction\s+([\w.:]+)'},
    'javascript': {'class': r'\bclass\s+(\w+)', 'func': r'\b(?:async\s+)?(?:function\s+(\w+)|(?:const|let|var)\s+(\w+)\s*=\s*(?:async\s*)?(?:\([^)]*\)|[^\s=]+)\s*=>)'},
    'typescript': {'class': r'\b(?:class|interface|type)\s+(\w+)', 'func': r'\b(?:async\s+)?(?:function\s+(\w+)|(?:const|let|var)\s+(\w+)\s*=\s*(?:async\s*)?(?:\([^)]*\)|[^\s=]+)\s*=>)'},
    'cpp': {'class': r'\b(?:class|struct)\s+(\w+)', 'func': r'\b\w+(?:<[^>]+>)?\s+(\w+)\s*\([^)]*\)\s*\{'},
    'c': {'class': r'\bstruct\s+(\w+)', 'func': r'\b\w+\s+(\w+)\s*\([^)]*\)\s*\{'}
}
COMPILED_LANG_RULES = {l: {k: re.compile(v) for k, v in r.items()} for l, r in LANG_RULES.items()}

DEFAULT_IGNORED = frozenset({
    '.git', '__pycache__', 'logs', '.auto-f', 'node_modules', 'venv', '.venv',
    'SSD_BACKUPS', '.cargo', '.rustup', '.cache', '.npm', '.local', '.mozilla',
    'target', '.agent', '.idea', '.vscode', 'projects', 'dist', 'build'
})

RE_MD_TITLE = re.compile(r'^#\s+(.*)', re.M)
RE_MD_FRONTMATTER = re.compile(r'^---\s*\n(.*?)\n---\s*\n', re.DOTALL)
RE_MD_TAGS = re.compile(r'tags\s*:\s*[\[\']?([^\]\']+)\]?', re.I)
RE_MD_WIKI_LINKS = re.compile(r'\[\[([^#\]|]+)(?:#[^\]|]+)?(?:\|[^\]]+)?\]\]')
RE_MD_STD_LINKS = re.compile(r'\[[^\]]+\]\(([^)]+\.(?:md|txt|py|sh|rs|go|lua))\)')
RE_SVG_WIDTH = re.compile(r'width\s*=\s*["\']([^"\']+)["\']')
RE_SVG_HEIGHT = re.compile(r'height\s*=\s*["\']([^"\']+)["\']')


def load_gitignore_patterns(target_dir: str) -> list[str]:
    p = os.path.join(target_dir, ".gitignore")
    if os.path.isfile(p):
        try:
            with open(p, 'r', encoding='utf-8', errors='ignore') as f:
                return [l.strip() for l in f if l.strip() and not l.strip().startswith(('#', '!'))]
        except Exception:
            pass
    return []


def should_ignore(rel_path: str, is_dir: bool, gitignore_patterns: list[str]) -> bool:
    parts = rel_path.replace(os.sep, '/').lower().split('/')
    if any(p in DEFAULT_IGNORED for p in parts):
        return True
    fn = parts[-1]
    if (fn in ('history.md', 'tpm.md', 'task_log.md', 'session.jsonl', '.spend_ledger.json', '.state.json', '.ds_store', 'thumbs.db') or
            fn.startswith('index-map-') or
            any(fn.endswith(ext) for ext in ('.db', '.db-wal', '.db-shm', '.sqlite', '.sqlite3', '.log', '.tmp'))):
        return True

    for pattern in gitignore_patterns:
        pat, dir_only = pattern.rstrip('/'), pattern.endswith('/')
        if dir_only and not is_dir:
            continue
        target = '/'.join(parts) if not dir_only else '/'.join(parts).rstrip('/')
        if pat.startswith('/'):
            p_clean = pat[1:]
            if fnmatch.fnmatchcase(target, p_clean) or (len(parts) >= len(p_clean.split('/')) and fnmatch.fnmatchcase('/'.join(parts[:len(p_clean.split('/'))]), p_clean)):
                return True
        else:
            if any(fnmatch.fnmatchcase(part, pat) for part in parts):
                return True
    return False


def is_binary(filepath: str) -> bool:
    try:
        with open(filepath, 'rb') as f:
            return b'\x00' in f.read(1024)
    except Exception:
        return True


def get_file_summary(filepath: str, limit: int = 10) -> str:
    """Streams first N lines to extract high-signal file headers with O(1) memory."""
    if not os.path.isfile(filepath) or is_binary(filepath):
        return ""
    try:
        with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
            for idx, line in enumerate(f):
                if idx >= limit:
                    break
                s = line.strip()
                if s and not line.startswith("#!") and not s.lower().startswith(("file:", "path:", "description:", "author:", "license:")):
                    return s.strip('#"\'/ \t')[:100].strip()
    except Exception:
        pass
    return ""


class PythonGraphVisitor(ast.NodeVisitor):
    def __init__(self, filepath: str):
        self.filepath, self.classes, self.functions, self.imports = filepath, [], [], []

    def visit_Import(self, node):
        for alias in node.names:
            self.imports.append((alias.name, alias.asname or alias.name))
        self.generic_visit(node)

    def visit_ImportFrom(self, node):
        if node.module:
            for alias in node.names:
                self.imports.append((f"{node.module}.{alias.name}", alias.asname or alias.name))
        self.generic_visit(node)

    def visit_ClassDef(self, node):
        bases = []
        for b in node.bases:
            if isinstance(b, ast.Name):
                bases.append(b.id)
            elif isinstance(b, ast.Attribute):
                bases.append(f"{getattr(b.value, 'id', '')}.{b.attr}")
        doc = ast.get_docstring(node) or ""
        self.classes.append((node.name, node.lineno, getattr(node, 'end_lineno', node.lineno), bases, doc))
        self.generic_visit(node)

    def visit_FunctionDef(self, node):
        self.visit_any_func(node)

    def visit_AsyncFunctionDef(self, node):
        self.visit_any_func(node)

    def visit_any_func(self, node):
        doc = ast.get_docstring(node) or ""
        args = [a.arg for a in node.args.args]
        self.functions.append((node.name, node.lineno, getattr(node, 'end_lineno', node.lineno), args, doc))
        self.generic_visit(node)


def parse_non_python_file(filepath: str, ext: str) -> dict:
    lang = {'.rs': 'rust', '.go': 'go', '.lua': 'lua', '.js': 'javascript', '.ts': 'typescript', '.cpp': 'cpp', '.hpp': 'cpp', '.c': 'c', '.h': 'c'}.get(ext)
    empty = {'classes': [], 'functions': []}
    if not lang or lang not in COMPILED_LANG_RULES:
        return empty
    rules, classes, functions = COMPILED_LANG_RULES[lang], [], []

    try:
        with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
            lines = f.readlines()
    except Exception:
        return empty

    for idx, line in enumerate(lines, 1):
        clean = line.strip()
        if not clean or clean[:2] in ('//', '/*', '--') or clean[0] == '#':
            continue
        if cm := rules['class'].search(clean):
            classes.append((cm.group(1), idx, idx, [], ""))
        if fm := rules['func'].search(clean):
            if fn_name := next((g for g in fm.groups() if g), None):
                functions.append((fn_name, idx, idx, [], ""))

    return {'classes': classes, 'functions': functions}


def analyze_markdown_file(filepath: str) -> dict:
    try:
        with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
            content = f.read()
        tm, tags, links = RE_MD_TITLE.search(content), set(), set()
        if fm_match := RE_MD_FRONTMATTER.search(content):
            if tags_match := RE_MD_TAGS.search(fm_match.group(1)):
                tags.update(t.strip().lower() for t in tags_match.group(1).split(',') if t.strip())
        for w in RE_MD_WIKI_LINKS.findall(content):
            links.add(w.strip())
        for ml in RE_MD_STD_LINKS.findall(content):
            if not ml.startswith(("http", "https")):
                links.add(os.path.basename(ml))
        return {"type": "markdown", "title": tm.group(1).strip() if tm else os.path.splitext(os.path.basename(filepath))[0], "tags": sorted(list(tags)), "links": sorted(list(links))}
    except Exception:
        return {"type": "markdown", "title": os.path.splitext(os.path.basename(filepath))[0]}


def analyze_image_file(filepath: str) -> dict:
    ext, size_bytes = os.path.splitext(filepath)[1].lstrip('.').lower(), os.path.getsize(filepath)
    size_str = f"{size_bytes / 1024:.1f}KB" if size_bytes < 1024 * 1024 else f"{size_bytes / (1024 * 1024):.1f}MB"
    try:
        with open(filepath, 'rb') as f:
            head = f.read(512)
        w, h = 0, 0
        if ext == "png" and head.startswith(b'\x89PNG\r\n\x1a\n'):
            w, h = struct.unpack('>II', head[16:24])
        elif ext == "gif" and head.startswith((b'GIF87a', b'GIF89a')):
            w, h = struct.unpack('<HH', head[6:10])
        elif ext in ("jpg", "jpeg") and head.startswith(b'\xff\xd8'):
            idx = 2
            while idx < len(head):
                marker, length = struct.unpack('>HH', head[idx:idx + 4])
                if 0xffc0 <= marker <= 0xffc3:
                    w, h = struct.unpack('>HH', head[idx + 5:idx + 9])
                    break
                idx += length + 2
        elif ext == "svg":
            text = head.decode("utf-8", errors="ignore")
            wm, hm = RE_SVG_WIDTH.search(text), RE_SVG_HEIGHT.search(text)
            return {"type": "image/svg", "dim": f"{wm.group(1) if wm else ''}x{hm.group(1) if hm else ''}", "size": size_str}
        return {"type": f"image/{ext}", "dim": f"{w}x{h}", "size": size_str} if w and h else {"type": f"image/{ext}", "size": size_str}
    except Exception:
        return {"type": f"image/{ext}", "size": size_str}


def analyze_json_file(filepath: str) -> dict:
    try:
        with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
            data = json.load(f)
        return {"type": "json", "keys": list(data.keys()) if isinstance(data, dict) else []}
    except Exception:
        return {"type": "json", "error": "Malformed JSON"}


def analyze_shell_file(filepath: str) -> dict:
    return {"type": "shell", "summary": get_file_summary(filepath, 10) or "Shell Module"}


def analyze_generic_file(filepath: str) -> dict:
    return {"type": f"binary/{os.path.splitext(filepath)[1].lstrip('.') or 'bin'}"} if is_binary(filepath) else {"type": os.path.splitext(filepath)[1].lstrip('.') or "file", "summary": get_file_summary(filepath, 15) or "Empty file"}


def init_db(db_path: str) -> sqlite3.Connection:
    os.makedirs(os.path.dirname(db_path), exist_ok=True)
    conn = sqlite3.connect(db_path)
    cur = conn.cursor()
    cur.execute("PRAGMA journal_mode=WAL;")
    cur.execute("DROP TABLE IF EXISTS nodes;")
    cur.execute("CREATE VIRTUAL TABLE nodes USING fts5(id, name, type, filepath, language, signature, doc, lines);")
    conn.commit()
    return conn


def find_db_path(start_dir: str | None = None) -> str | None:
    curr = os.path.abspath(start_dir or os.getcwd())
    while True:
        for check in (os.path.join(curr, ".agent"), curr):
            if os.path.isdir(check):
                if found := next((os.path.join(check, f) for f in os.listdir(check) if f.startswith("index-map-memory-") and f.endswith(".db")), None):
                    return found
        if (parent := os.path.dirname(curr)) == curr:
            break
        curr = parent
    return None


def compile_ascii_tree(target_dir: str) -> str:
    files = []
    for root in (target_dir, os.path.join(target_dir, "modules")):
        if os.path.isdir(root):
            sub = "modules" if "modules" in root else ""
            files.extend((sub, f, os.path.join(root, f)) for f in sorted(os.listdir(root)) if os.path.isfile(os.path.join(root, f)) and not f.startswith("index-map-") and not f.endswith(('.db', '.tmp', '.log')))

    if not files:
        return ""
    lines = ["# Codebase Modules & Active Architecture:"]
    for idx, (subdir, name, path) in enumerate(files):
        disp, desc = os.path.join(subdir, name) if subdir else name, get_file_summary(path)
        connector = "└── " if idx == len(files) - 1 else "├── "
        lines.append(f"#   {connector}{disp:<24} - {desc}" if desc else f"#   {connector}{disp}")
    return "\n".join(lines) + "\n# -----------------------------------------------------------------------------------------\n"


def compile_flat_shorthand(skeleton: dict) -> str:
    lines = []
    for rel_path, data in sorted(skeleton.items()):
        t = data.get("type", "file")
        if t == "python":
            parts = [f"imp:{','.join(data['imports'])}" if data.get('imports') else "", f"cls:{','.join(data['classes'])}" if data.get('classes') else "", f"fn:{','.join(data['functions'])}" if data.get('functions') else ""]
        elif t in ("rust", "go", "lua", "javascript", "typescript", "cpp", "c"):
            parts = [f"cls:{','.join(data['classes'])}" if data.get('classes') else "", f"fn:{','.join(data['functions'])}" if data.get('functions') else ""]
        elif t == "markdown":
            parts = [f"title:{data['title']}" if data.get('title') else "", f"tags:{','.join(data['tags'])}" if data.get('tags') else "", f"links:{','.join(data['links'])}" if data.get('links') else ""]
        elif t.startswith("image/"):
            parts = [f"dim:{data['dim']}" if data.get('dim') else "", f"size:{data['size']}" if data.get('size') else ""]
        elif t == "json":
            parts = [f"keys:{','.join(data['keys'])}" if data.get('keys') else ""]
        else:
            parts = [f"sum:{data.get('summary', '').replace('\n', ' ').strip()}" if data.get('summary') else ""]

        meta = " | ".join(filter(None, parts))
        lines.append(f"[{t}:{rel_path}] {meta}" if meta else f"[{t}:{rel_path}]")
    return "\n".join(lines)


def run_indexing(target_dir: str, force: bool = False, is_agent: bool = False, no_db: bool = False, print_output: bool = False) -> None:
    if os.path.basename(target_dir) == ".agent":
        target_dir = os.path.dirname(target_dir)
    proj_name = os.path.basename(target_dir)

    if is_agent:
        out_dir = os.path.join(target_dir, ".agent")
    else:
        out_dir = target_dir
        if not any(flag in sys.argv for flag in ("--db", "--with-db")):
            no_db = True

    os.makedirs(out_dir, exist_ok=True)
    output_file = os.path.join(out_dir, f"index-map-{proj_name}.txt")

    db_conn = None
    if not no_db:
        db_path = os.path.join(out_dir, f"index-map-memory-{proj_name}.db")
        db_conn = init_db(db_path)

    patterns, skeleton = load_gitignore_patterns(target_dir), {}
    db_rows = []

    for root, dirs, files in os.walk(target_dir):
        dirs[:] = [d for d in dirs if not should_ignore(os.path.relpath(os.path.join(root, d), target_dir), True, patterns)]
        for file in files:
            fp = os.path.join(root, file)
            rel_path = os.path.relpath(fp, target_dir)
            if should_ignore(rel_path, False, patterns):
                continue

            ext = os.path.splitext(file)[1].lower()
            is_py_shebang = False
            if not ext and not is_binary(fp):
                try:
                    with open(fp, 'r', encoding='utf-8', errors='ignore') as f:
                        first_line = f.readline()
                        second_line = f.readline()
                        if first_line.startswith("#!") and ("python" in first_line or "python" in second_line):
                            is_py_shebang = True
                except Exception:
                    pass

            if ext == '.py' or is_py_shebang:
                try:
                    with open(fp, 'r', encoding='utf-8', errors='replace') as f:
                        content = f.read()
                    tree = ast.parse(content, filename=fp)
                    v = PythonGraphVisitor(rel_path)
                    v.visit(tree)
                    skeleton[rel_path] = {
                        "type": "python",
                        "imports": sorted(list({i[0].split('.')[0] for i in v.imports})),
                        "classes": [c[0] for c in v.classes],
                        "functions": [f[0] for f in v.functions]
                    }
                    if db_conn:
                        for cname, start, end, bases, doc in v.classes:
                            db_rows.append((f"class:{rel_path}:{cname}", cname, "class", rel_path, "python", f"class {cname}({', '.join(bases)})", doc, f"{start}-{end}"))
                        for fname, start, end, args, doc in v.functions:
                            db_rows.append((f"func:{rel_path}:{fname}", fname, "function", rel_path, "python", f"{fname}({', '.join(args)})", doc, f"{start}-{end}"))
                except Exception as e:
                    skeleton[rel_path] = {"type": "python", "error": f"AST error: {e}"}
            elif ext in ('.rs', '.go', '.lua', '.js', '.ts', '.cpp', '.hpp', '.c', '.h'):
                parsed = parse_non_python_file(fp, ext)
                skeleton[rel_path] = {"type": ext.lstrip('.'), "classes": [c[0] for c in parsed['classes']], "functions": [f[0] for f in parsed['functions']]}
                if db_conn:
                    for cname, start, end, _, doc in parsed['classes']:
                        db_rows.append((f"class:{rel_path}:{cname}", cname, "class", rel_path, ext.lstrip('.'), cname, doc, f"{start}-{end}"))
                    for fname, start, end, _, doc in parsed['functions']:
                        db_rows.append((f"func:{rel_path}:{fname}", fname, "function", rel_path, ext.lstrip('.'), fname, doc, f"{start}-{end}"))
            else:
                dispatch = {'.md': analyze_markdown_file, '.sh': analyze_shell_file, '.json': analyze_json_file, '.png': analyze_image_file, '.jpg': analyze_image_file, '.svg': analyze_image_file}
                res = dispatch.get(ext, analyze_generic_file)(fp)
                skeleton[rel_path] = res
                if db_conn and res.get("title"):
                    db_rows.append((f"doc:{rel_path}", res.get("title", ""), "markdown", rel_path, "markdown", res.get("title", ""), ", ".join(res.get("tags", [])), "1-1"))

    if db_conn and db_rows:
        db_conn.cursor().executemany("INSERT INTO nodes VALUES (?, ?, ?, ?, ?, ?, ?, ?)", db_rows)
        db_conn.commit()
        db_conn.close()

    disp_dir = target_dir.replace(os.path.expanduser("~"), "~", 1) if target_dir.startswith(os.path.expanduser("~")) else target_dir
    header = f"# CODEBASE INDEX MAP | Generated: {time.strftime('%Y-%m-%d %H:%M')} | Workspace: {proj_name} | Path: {disp_dir}\n# Active symbols and structural connections mapped in the local index.\n# -----------------------------------------------------------------------------------------\n"
    header += compile_ascii_tree(target_dir) or "# -----------------------------------------------------------------------------------------\n\n"
    flat_map = header + compile_flat_shorthand(skeleton)

    try:
        tmp = f"{output_file}.tmp"
        with open(tmp, 'w', encoding='utf-8') as f:
            f.write(flat_map)
        os.replace(tmp, output_file)
        if print_output:
            sys.stdout.write(flat_map)
            sys.stdout.flush()
        elif not is_agent:
            sys.stderr.write(f"\r\x1b[2K\033[1;32m[ok] Mapping complete! [{proj_name} index-map {'& SQLite database updated' if db_conn else 'updated'}]\033[0m\n")
    except Exception as e:
        sys.stderr.write(f"\r\x1b[2K\033[1;31m[error] Failed to write map: {e}\033[0m\n")


def search_symbols(pattern: str, workspace: str | None = None) -> str:
    db = find_db_path(workspace)
    if not db:
        return "[error] Index database not found. Initialize with 'ai init' or 'index-map init --db'."
    conn = sqlite3.connect(db)
    cur = conn.cursor()

    clean_query = re.sub(r'[^a-zA-Z0-9_* ]', ' ', pattern).strip()
    if not clean_query:
        conn.close()
        return "[error] Empty search pattern."

    fts_terms = " OR ".join(f"{w}*" for w in clean_query.split())
    sql = "SELECT id, type, filepath, signature, doc, lines, bm25(nodes) as rank FROM nodes WHERE nodes MATCH ? ORDER BY rank LIMIT 15"
    try:
        rows = cur.execute(sql, (fts_terms,)).fetchall()
    except Exception:
        rows = cur.execute("SELECT id, type, filepath, signature, doc, lines, 0 FROM nodes WHERE name LIKE ? LIMIT 15", (f"%{clean_query}%",)).fetchall()
    conn.close()

    if not rows:
        return f"[no matches found for '{pattern}']"

    res = [f"\n\033[1;32m[Matches for: '{pattern}']\033[0m"]
    for nid, ntype, fpath, sig, doc, lines, _ in rows:
        doc_snippet = f" • [italic]{doc.splitlines()[0][:60]}[/italic]" if doc else ""
        res.append(f"  • \033[1;36m{sig}\033[0m [{ntype}] in \033[33m{fpath}:{lines}\033[0m{doc_snippet}")
    return "\n".join(res)


def extract_snippet(symbol: str, workspace: str | None = None) -> str:
    db = find_db_path(workspace)
    if not db:
        return "[error] Database not found."
    conn = sqlite3.connect(db)
    cur = conn.cursor()
    rows = cur.execute("SELECT filepath, lines, type, name FROM nodes WHERE name = ? OR id LIKE ? LIMIT 5", (symbol, f"%:{symbol}")).fetchall()
    conn.close()

    if not rows:
        return f"[error] Symbol '{symbol}' not found."

    res = []
    root_dir = os.path.dirname(os.path.dirname(db)) if os.path.basename(os.path.dirname(db)) == ".agent" else os.path.dirname(db)
    for filepath, lines_str, node_type, name in rows:
        if not lines_str or "-" not in lines_str:
            continue
        l_start, l_end = map(int, lines_str.split("-"))
        full_path = os.path.join(root_dir, filepath)
        if not os.path.isfile(full_path):
            continue

        res.append(f"\n\033[1;36m[Snippet] {node_type.upper()} {name} in {filepath} (Lines {l_start} to {l_end}):\033[0m")
        try:
            with open(full_path, 'r', encoding='utf-8', errors='ignore') as f:
                lines = f.readlines()
            res.append("".join(lines[max(0, l_start - 1):l_end]))
        except Exception as e:
            res.append(f"[error] Extract failed: {e}")
    return "\n".join(res)


def main():
    raw_args = sys.argv[1:]
    is_agent = "--agent" in raw_args
    force = "--force" in raw_args or "-f" in raw_args
    no_db = "--no-db" in raw_args or "--txt-only" in raw_args
    print_cat = any(f in raw_args for f in ("--cat", "-p", "--print", "--show", "cat", "show"))
    args = [a for a in raw_args if a not in ("--agent", "--db", "--with-db", "--no-db", "--txt-only", "--force", "-f", "--cat", "-p", "--print", "--show", "cat", "show")]

    if not args:
        run_indexing(os.getcwd(), force=force, is_agent=is_agent, no_db=no_db, print_output=print_cat)
        return

    cmd = args[0].lower()
    if cmd in ("--help", "-h"):
        print("Usage: index-map [init | search <pattern> | snippet <symbol>] [--db] [--force]")
    elif cmd == "init":
        target = os.path.abspath(os.path.expanduser(args[1])) if len(args) > 1 else os.getcwd()
        run_indexing(target, force=force, is_agent=is_agent, no_db=no_db)
    elif cmd == "search":
        print(search_symbols(args[1], workspace=args[2] if len(args) > 2 else None)) if len(args) > 1 else print("[error] Search pattern required.")
    elif cmd == "snippet":
        print(extract_snippet(args[1], workspace=args[2] if len(args) > 2 else None)) if len(args) > 1 else print("[error] Symbol required.")
    else:
        path_test = os.path.abspath(os.path.expanduser(args[0]))
        if os.path.isdir(path_test):
            run_indexing(path_test, force=force, is_agent=is_agent, no_db=no_db)
        else:
            print(f"[error] Unknown command: {args[0]}")


if __name__ == "__main__":
    main()
