# ================================================================
# yui_code_graph.py — CODE KNOWLEDGE GRAPH v1.0
# Inspired by CodeGraph (Colby McInny) — adapted for Yui/free-tier
#
# Quét toàn bộ Python project → build Knowledge Graph → SQLite
# Khi AI cần hiểu code, query graph thay vì đọc lại file từ đầu
#
# HIỆU QUẢ (theo CodeGraph research):
#   • Tiết kiệm ~59% token (không đọc lại file)
#   • Giảm ~70% tool calls (1 query thay vì scan nhiều file)
#   • Nhanh hơn ~49% overall
#
# HOÀN TOÀN FREE:
#   • SQLite local (không cần DB server)
#   • Pure Python stdlib (ast, sqlite3, pathlib)
#   • Không cần API, không cần key
#
# SCHEMA:
#   nodes(id, type, name, file, line, code_snippet, docstring)
#   edges(from_id, to_id, relation)  -- calls, imports, inherits, uses
#   file_summaries(file, summary, last_modified, token_count)
#   search_index(term, node_id, weight)
# ================================================================

import os, ast, re, sqlite3, time, hashlib, json
from pathlib import Path
from typing import Optional, List, Dict, Tuple, Any
from dataclasses import dataclass, field


# ─── SCHEMA ──────────────────────────────────────────────────────
_SCHEMA = """
CREATE TABLE IF NOT EXISTS nodes (
    id           TEXT PRIMARY KEY,
    type         TEXT NOT NULL,    -- function|class|method|import|variable|constant
    name         TEXT NOT NULL,
    file         TEXT NOT NULL,
    line         INTEGER,
    code_snippet TEXT,             -- Chỉ signature, không phải toàn bộ body
    docstring    TEXT,
    decorators   TEXT,             -- JSON list
    parent       TEXT,             -- class chứa method
    created_at   REAL
);

CREATE TABLE IF NOT EXISTS edges (
    from_id  TEXT NOT NULL,
    to_id    TEXT NOT NULL,
    relation TEXT NOT NULL,        -- calls|imports|inherits|uses|defines
    PRIMARY KEY (from_id, to_id, relation)
);

CREATE TABLE IF NOT EXISTS file_summaries (
    file          TEXT PRIMARY KEY,
    summary       TEXT,            -- AI-generated hoặc auto-extract
    lines         INTEGER,
    functions     INTEGER,
    classes       INTEGER,
    imports       TEXT,            -- JSON list
    last_modified REAL,
    hash          TEXT
);

CREATE TABLE IF NOT EXISTS search_index (
    term    TEXT NOT NULL,
    node_id TEXT NOT NULL,
    weight  REAL DEFAULT 1.0,
    PRIMARY KEY (term, node_id)
);

CREATE INDEX IF NOT EXISTS idx_nodes_name   ON nodes(name);
CREATE INDEX IF NOT EXISTS idx_nodes_file   ON nodes(file);
CREATE INDEX IF NOT EXISTS idx_nodes_type   ON nodes(type);
CREATE INDEX IF NOT EXISTS idx_edges_from   ON edges(from_id);
CREATE INDEX IF NOT EXISTS idx_edges_to     ON edges(to_id);
CREATE INDEX IF NOT EXISTS idx_search_term  ON search_index(term);
"""


# ─── NODE ID GENERATOR ───────────────────────────────────────────
def _node_id(file: str, name: str, line: int = 0) -> str:
    key = f"{file}::{name}::{line}"
    return hashlib.md5(key.encode()).hexdigest()[:12]


# ─── AST VISITOR ─────────────────────────────────────────────────
@dataclass
class ParsedNode:
    nid:       str
    ntype:     str
    name:      str
    file:      str
    line:      int
    snippet:   str
    docstring: str
    decs:      List[str] = field(default_factory=list)
    parent:    str = ""
    calls:     List[str] = field(default_factory=list)
    imports:   List[str] = field(default_factory=list)


class CodeVisitor(ast.NodeVisitor):
    """AST visitor tách nodes + relationships."""

    def __init__(self, filepath: str, source: str):
        self.filepath = filepath
        self.source   = source
        self.lines    = source.splitlines()
        self.nodes: List[ParsedNode] = []
        self._class_stack: List[str] = []  # stack khi vào class
        self._current_fn: Optional[str] = None

    # ── IMPORTS ──────────────────────────────────────────────────
    def visit_Import(self, node: ast.Import):
        for alias in node.names:
            name = alias.asname or alias.name
            nid  = _node_id(self.filepath, f"import:{alias.name}", node.lineno)
            self.nodes.append(ParsedNode(
                nid=nid, ntype="import", name=name,
                file=self.filepath, line=node.lineno,
                snippet=f"import {alias.name}",
                docstring="",
                imports=[alias.name],
            ))
        self.generic_visit(node)

    def visit_ImportFrom(self, node: ast.ImportFrom):
        module = node.module or ""
        for alias in node.names:
            full = f"{module}.{alias.name}" if module else alias.name
            nid  = _node_id(self.filepath, f"from:{full}", node.lineno)
            self.nodes.append(ParsedNode(
                nid=nid, ntype="import", name=alias.asname or alias.name,
                file=self.filepath, line=node.lineno,
                snippet=f"from {module} import {alias.name}",
                docstring="",
                imports=[full],
            ))
        self.generic_visit(node)

    # ── CLASSES ──────────────────────────────────────────────────
    def visit_ClassDef(self, node: ast.ClassDef):
        nid    = _node_id(self.filepath, node.name, node.lineno)
        bases  = [self._name(b) for b in node.bases]
        decs   = [self._name(d) for d in node.decorator_list]
        doc    = ast.get_docstring(node) or ""
        snippet = f"class {node.name}({', '.join(bases)}):"

        pnode = ParsedNode(
            nid=nid, ntype="class", name=node.name,
            file=self.filepath, line=node.lineno,
            snippet=snippet, docstring=doc[:200],
            decs=decs,
            imports=bases,  # bases treated as "uses" relationship
            parent=self._class_stack[-1] if self._class_stack else "",
        )
        self.nodes.append(pnode)

        self._class_stack.append(node.name)
        self.generic_visit(node)
        self._class_stack.pop()

    # ── FUNCTIONS ─────────────────────────────────────────────────
    def visit_FunctionDef(self, node: ast.FunctionDef):
        self._handle_fn(node)

    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef):
        self._handle_fn(node, is_async=True)

    def _handle_fn(self, node, is_async=False):
        parent_class = self._class_stack[-1] if self._class_stack else ""
        ntype = "method" if parent_class else "function"

        # Build signature
        args     = self._get_args(node.args)
        ret      = self._name(node.returns) if node.returns else ""
        ret_str  = f" -> {ret}" if ret else ""
        prefix   = "async " if is_async else ""
        snippet  = f"{prefix}def {node.name}({', '.join(args)}){ret_str}:"
        doc      = ast.get_docstring(node) or ""
        decs     = [self._name(d) for d in node.decorator_list]

        nid = _node_id(self.filepath, f"{parent_class}.{node.name}" if parent_class else node.name, node.lineno)

        # Tìm các hàm được gọi trong body
        calls = self._find_calls(node)

        pnode = ParsedNode(
            nid=nid, ntype=ntype, name=node.name,
            file=self.filepath, line=node.lineno,
            snippet=snippet, docstring=doc[:200],
            decs=decs, parent=parent_class,
            calls=calls,
        )
        self.nodes.append(pnode)

        prev = self._current_fn
        self._current_fn = f"{parent_class}.{node.name}" if parent_class else node.name
        self.generic_visit(node)
        self._current_fn = prev

    # ── CONSTANTS ─────────────────────────────────────────────────
    def visit_Assign(self, node: ast.Assign):
        # Chỉ bắt module-level constants (ALL_CAPS)
        if not self._class_stack and not self._current_fn:
            for target in node.targets:
                if isinstance(target, ast.Name) and target.id.isupper():
                    try:
                        val = ast.literal_eval(node.value)
                        snippet = f"{target.id} = {repr(val)[:60]}"
                    except:
                        snippet = f"{target.id} = ..."
                    nid = _node_id(self.filepath, target.id, node.lineno)
                    self.nodes.append(ParsedNode(
                        nid=nid, ntype="constant", name=target.id,
                        file=self.filepath, line=node.lineno,
                        snippet=snippet, docstring="",
                    ))
        self.generic_visit(node)

    # ── HELPERS ───────────────────────────────────────────────────
    def _name(self, node) -> str:
        if node is None: return ""
        if isinstance(node, ast.Name):      return node.id
        if isinstance(node, ast.Attribute): return f"{self._name(node.value)}.{node.attr}"
        if isinstance(node, ast.Constant):  return repr(node.value)
        if isinstance(node, ast.Subscript): return f"{self._name(node.value)}[...]"
        return ast.dump(node)[:30]

    def _get_args(self, args: ast.arguments) -> List[str]:
        result = []
        for arg in args.args:
            ann = f": {self._name(arg.annotation)}" if arg.annotation else ""
            result.append(f"{arg.arg}{ann}")
        if args.vararg:
            result.append(f"*{args.vararg.arg}")
        if args.kwarg:
            result.append(f"**{args.kwarg.arg}")
        return result[:6]  # Max 6 args in snippet

    def _find_calls(self, fn_node) -> List[str]:
        calls = []
        for node in ast.walk(fn_node):
            if isinstance(node, ast.Call):
                name = self._name(node.func)
                if name and '.' not in name[:1] and name not in calls:
                    calls.append(name)
        return calls[:15]


# ─── FILE PARSER ─────────────────────────────────────────────────
def parse_file(filepath: str) -> Tuple[List[ParsedNode], Dict]:
    """Parse một Python file → nodes + file summary."""
    try:
        with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
            source = f.read()
    except Exception as e:
        return [], {}

    try:
        tree = ast.parse(source, filename=filepath)
    except SyntaxError:
        return [], {}

    visitor = CodeVisitor(filepath, source)
    visitor.visit(tree)

    # File summary
    all_imports = []
    fn_count    = 0
    cls_count   = 0
    for n in visitor.nodes:
        if n.ntype == "import":
            all_imports.extend(n.imports)
        elif n.ntype in ("function", "method"):
            fn_count += 1
        elif n.ntype == "class":
            cls_count += 1

    summary = {
        "file":     filepath,
        "lines":    len(source.splitlines()),
        "functions":fn_count,
        "classes":  cls_count,
        "imports":  list(set(all_imports))[:20],
        "last_modified": os.path.getmtime(filepath),
        "hash":     hashlib.md5(source[:2000].encode()).hexdigest()[:8],
    }
    return visitor.nodes, summary


# ─── GRAPH BUILDER ───────────────────────────────────────────────
class CodeGraph:
    """
    Knowledge Graph cho Python codebase.
    Lưu vào SQLite, query nhanh không cần đọc file.
    """

    def __init__(self, db_path: str = "/tmp/yui_code_graph.db"):
        self.db_path = db_path
        self._conn   = None
        self._init_db()

    def _conn_(self) -> sqlite3.Connection:
        if not self._conn:
            self._conn = sqlite3.connect(self.db_path, check_same_thread=False)
            self._conn.row_factory = sqlite3.Row
            self._conn.execute("PRAGMA journal_mode=WAL")
            self._conn.execute("PRAGMA synchronous=NORMAL")
        return self._conn

    def _init_db(self):
        conn = self._conn_()
        conn.executescript(_SCHEMA)
        conn.commit()
        print(f"[CODEGRAPH] DB: {self.db_path}")

    # ── BUILD / UPDATE ────────────────────────────────────────────
    def scan_directory(
        self,
        root: str,
        extensions: List[str] = [".py"],
        exclude_dirs: List[str] = ["__pycache__", ".git", "venv", "node_modules", ".pyc"],
        max_file_kb: int = 500,
    ) -> Dict[str, int]:
        """
        Quét toàn bộ directory → build graph.
        Returns: stats dict
        """
        stats = {"files": 0, "nodes": 0, "edges": 0, "skipped": 0}
        root_path = Path(root)

        py_files = []
        for ext in extensions:
            for f in root_path.rglob(f"*{ext}"):
                skip = False
                for ex in exclude_dirs:
                    if ex in str(f):
                        skip = True; break
                if not skip and f.stat().st_size < max_file_kb * 1024:
                    py_files.append(str(f))

        print(f"[CODEGRAPH] Tìm thấy {len(py_files)} files, đang quét...")

        for fp in py_files:
            try:
                n_new, n_updated = self._index_file(fp)
                stats["files"]  += 1
                stats["nodes"]  += n_new
            except Exception as e:
                stats["skipped"] += 1
                print(f"[CODEGRAPH] Skip {fp}: {e}")

        stats["edges"] = self._conn_().execute("SELECT COUNT(*) FROM edges").fetchone()[0]
        print(f"[CODEGRAPH] ✅ Scan xong: {stats}")
        return stats

    def _index_file(self, filepath: str) -> Tuple[int, int]:
        """Index một file. Trả về (nodes_added, nodes_updated)."""
        conn = self._conn_()

        # Check if file changed
        try:
            mtime = os.path.getmtime(filepath)
            row = conn.execute(
                "SELECT last_modified, hash FROM file_summaries WHERE file=?",
                (filepath,)
            ).fetchone()
            if row:
                with open(filepath,'r',encoding='utf-8',errors='ignore') as f:
                    src = f.read()
                h = hashlib.md5(src[:2000].encode()).hexdigest()[:8]
                if abs(row["last_modified"] - mtime) < 1 and row["hash"] == h:
                    return 0, 0  # File chưa thay đổi
        except:
            pass

        nodes, summary = parse_file(filepath)

        # Clear old data for this file
        old_ids = [r[0] for r in conn.execute(
            "SELECT id FROM nodes WHERE file=?", (filepath,)
        ).fetchall()]
        if old_ids:
            conn.execute(
                f"DELETE FROM edges WHERE from_id IN ({','.join('?'*len(old_ids))})",
                old_ids
            )
            conn.execute("DELETE FROM nodes WHERE file=?", (filepath,))
            conn.execute("DELETE FROM search_index WHERE node_id IN "
                         f"({','.join('?'*len(old_ids))})", old_ids)

        now = time.time()
        # Insert nodes
        for n in nodes:
            conn.execute("""
                INSERT OR REPLACE INTO nodes
                (id, type, name, file, line, code_snippet, docstring, decorators, parent, created_at)
                VALUES (?,?,?,?,?,?,?,?,?,?)
            """, (
                n.nid, n.ntype, n.name, n.file, n.line,
                n.snippet, n.docstring,
                json.dumps(n.decs), n.parent, now
            ))

            # Build call edges
            for called in n.calls:
                edge_to = _node_id(filepath, called, 0)
                conn.execute(
                    "INSERT OR IGNORE INTO edges VALUES (?,?,?)",
                    (n.nid, edge_to, "calls")
                )

            # Build import edges
            for imp in n.imports:
                edge_to = _node_id(filepath, f"import:{imp}", 0)
                conn.execute(
                    "INSERT OR IGNORE INTO edges VALUES (?,?,?)",
                    (n.nid, edge_to, "imports")
                )

            # Search index — tokenize name + docstring
            terms = set()
            for word in re.split(r'[_\s\./]+', n.name.lower()):
                if len(word) >= 2: terms.add(word)
            for word in re.split(r'\W+', (n.docstring or "").lower()):
                if len(word) >= 3: terms.add(word)
            for term in terms:
                conn.execute(
                    "INSERT OR IGNORE INTO search_index VALUES (?,?,?)",
                    (term, n.nid, 1.0)
                )

        # File summary
        conn.execute("""
            INSERT OR REPLACE INTO file_summaries
            (file, summary, lines, functions, classes, imports, last_modified, hash)
            VALUES (?,?,?,?,?,?,?,?)
        """, (
            summary.get("file",""), "",
            summary.get("lines",0), summary.get("functions",0),
            summary.get("classes",0), json.dumps(summary.get("imports",[])),
            summary.get("last_modified", now), summary.get("hash",""),
        ))

        conn.commit()
        return len(nodes), 0

    # ── QUERY API ────────────────────────────────────────────────
    def find(self, name: str, ntype: str = "") -> List[Dict]:
        """Tìm node theo tên. Hỗ trợ partial match."""
        conn  = self._conn_()
        where = "WHERE LOWER(name) LIKE ?"
        args  = [f"%{name.lower()}%"]
        if ntype:
            where += " AND type=?"
            args.append(ntype)
        rows = conn.execute(
            f"SELECT * FROM nodes {where} ORDER BY LENGTH(name) LIMIT 15", args
        ).fetchall()
        return [dict(r) for r in rows]

    def who_calls(self, func_name: str) -> List[Dict]:
        """Ai gọi hàm này?"""
        conn = self._conn_()
        rows = conn.execute("""
            SELECT n.* FROM nodes n
            JOIN edges e ON e.from_id = n.id
            WHERE e.relation='calls'
              AND e.to_id IN (
                  SELECT id FROM nodes WHERE LOWER(name)=LOWER(?)
              )
            LIMIT 20
        """, (func_name,)).fetchall()
        return [dict(r) for r in rows]

    def what_calls(self, func_name: str) -> List[Dict]:
        """Hàm này gọi những gì?"""
        conn = self._conn_()
        row  = conn.execute(
            "SELECT id FROM nodes WHERE LOWER(name)=LOWER(?) LIMIT 1",
            (func_name,)
        ).fetchone()
        if not row: return []
        rows = conn.execute("""
            SELECT n.* FROM nodes n
            JOIN edges e ON e.to_id = n.id
            WHERE e.from_id=? AND e.relation='calls'
            LIMIT 20
        """, (row["id"],)).fetchall()
        return [dict(r) for r in rows]

    def file_overview(self, filepath: str) -> Dict:
        """Tổng quan 1 file: functions, classes, imports."""
        conn = self._conn_()
        summary = conn.execute(
            "SELECT * FROM file_summaries WHERE file LIKE ?",
            (f"%{filepath}%",)
        ).fetchone()
        nodes   = conn.execute(
            "SELECT type, name, line, code_snippet FROM nodes WHERE file LIKE ? ORDER BY line",
            (f"%{filepath}%",)
        ).fetchall()
        return {
            "summary": dict(summary) if summary else {},
            "nodes":   [dict(n) for n in nodes],
        }

    def search(self, query: str, limit: int = 10) -> List[Dict]:
        """Full-text search trên name + docstring."""
        conn  = self._conn_()
        terms = [t for t in re.split(r'\W+', query.lower()) if len(t) >= 2]
        if not terms: return []

        # Score nodes by how many terms match
        placeholders = ",".join("?" * len(terms))
        rows = conn.execute(f"""
            SELECT n.*, SUM(si.weight) as score
            FROM search_index si
            JOIN nodes n ON n.id = si.node_id
            WHERE si.term IN ({placeholders})
            GROUP BY n.id
            ORDER BY score DESC, LENGTH(n.name)
            LIMIT ?
        """, terms + [limit]).fetchall()
        return [dict(r) for r in rows]

    def get_context_for_query(self, user_query: str, max_tokens: int = 2000) -> str:
        """
        Main API cho AI: Nhận query → trả về context từ graph.
        Thay thế việc đọc lại file từ đầu.
        """
        results = self.search(user_query, limit=12)
        if not results:
            # Fallback: exact name match
            for word in re.split(r'\W+', user_query):
                if len(word) >= 4:
                    results = self.find(word)
                    if results: break

        if not results:
            return ""

        lines = ["[CODEGRAPH CONTEXT]"]
        token_estimate = 0

        # Group by file
        by_file: Dict[str, List] = {}
        for r in results:
            f = r.get("file", "?")
            by_file.setdefault(f, []).append(r)

        for filepath, nodes in by_file.items():
            fname = os.path.basename(filepath)
            lines.append(f"\n📄 {fname}")
            for n in nodes[:5]:
                snippet = n.get("code_snippet", "")
                doc     = n.get("docstring", "")
                line    = n.get("line", 0)
                lines.append(f"  L{line} [{n['type']}] {snippet}")
                if doc:
                    lines.append(f"    → {doc[:80]}")
                token_estimate += len(snippet) // 4 + 5
                if token_estimate > max_tokens:
                    break
            if token_estimate > max_tokens:
                break

        lines.append(f"\n[{len(results)} nodes found | ~{token_estimate} tokens]")
        return "\n".join(lines)

    def stats(self) -> Dict:
        """Stats của graph."""
        conn = self._conn_()
        return {
            "nodes":        conn.execute("SELECT COUNT(*) FROM nodes").fetchone()[0],
            "edges":        conn.execute("SELECT COUNT(*) FROM edges").fetchone()[0],
            "files":        conn.execute("SELECT COUNT(*) FROM file_summaries").fetchone()[0],
            "functions":    conn.execute("SELECT COUNT(*) FROM nodes WHERE type='function'").fetchone()[0],
            "methods":      conn.execute("SELECT COUNT(*) FROM nodes WHERE type='method'").fetchone()[0],
            "classes":      conn.execute("SELECT COUNT(*) FROM nodes WHERE type='class'").fetchone()[0],
            "db_size_kb":   os.path.getsize(self.db_path) // 1024 if os.path.exists(self.db_path) else 0,
        }

    def export_summary(self) -> str:
        """Tạo text summary cho AI về toàn bộ codebase."""
        conn  = self._conn_()
        files = conn.execute(
            "SELECT * FROM file_summaries ORDER BY functions DESC LIMIT 20"
        ).fetchall()

        lines = ["=== CODEBASE OVERVIEW ==="]
        for f in files:
            fname = os.path.basename(f["file"])
            lines.append(
                f"  {fname}: {f['lines']}L, {f['functions']} funcs, {f['classes']} classes"
            )

        top_fns = conn.execute("""
            SELECT name, file, line FROM nodes
            WHERE type IN ('function','method')
            ORDER BY rowid DESC LIMIT 20
        """).fetchall()
        lines.append("\n=== KEY FUNCTIONS ===")
        for fn in top_fns:
            lines.append(f"  {fn['name']}() @ {os.path.basename(fn['file'])}:{fn['line']}")

        return "\n".join(lines)


# ─── SINGLETON + AUTO-SCAN ────────────────────────────────────────
_graph: Optional[CodeGraph] = None
_last_scan: float = 0
_SCAN_INTERVAL = 300   # Re-scan mỗi 5 phút


def get_graph() -> CodeGraph:
    """Lấy singleton graph, tự scan nếu cần."""
    global _graph, _last_scan
    if _graph is None:
        _graph = CodeGraph(db_path=os.environ.get(
            "CODEGRAPH_DB", "/tmp/yui_code_graph.db"
        ))
    return _graph


def ensure_indexed(root: str = "/app", force: bool = False):
    """Đảm bảo codebase đã được index. Gọi lúc startup."""
    global _last_scan
    g = get_graph()
    now = time.time()
    if force or (now - _last_scan) > _SCAN_INTERVAL:
        stats = g.scan_directory(root)
        _last_scan = now
        return stats
    return g.stats()


# ─── PUBLIC QUERY API ─────────────────────────────────────────────
def query_graph(user_query: str, max_tokens: int = 1500) -> str:
    """
    Entry point cho ural_code_engine.py.
    Trả về context string để inject vào code prompt.
    """
    g = get_graph()
    if g.stats()["nodes"] == 0:
        return ""
    return g.get_context_for_query(user_query, max_tokens)


def graph_stats_string() -> str:
    """Stats string ngắn gọn cho /status command."""
    g = get_graph()
    s = g.stats()
    return (
        f"📊 CodeGraph: {s['nodes']} nodes | {s['edges']} edges | "
        f"{s['files']} files | {s['db_size_kb']}KB SQLite"
    )


# ─── STARTUP INDEX ───────────────────────────────────────────────
import threading

def _background_index():
    """Index codebase trong background khi startup."""
    time.sleep(8)  # Đợi app khởi động xong
    try:
        root = os.environ.get("CODEGRAPH_ROOT", "/app")
        if not os.path.exists(root):
            root = os.path.dirname(os.path.abspath(__file__))
        stats = ensure_indexed(root)
        print(f"[CODEGRAPH] Background index xong: {stats}")
    except Exception as e:
        print(f"[CODEGRAPH] Background index lỗi: {e}")


_bg_thread = threading.Thread(target=_background_index, daemon=True)
_bg_thread.start()

print("✅ [CODEGRAPH] Knowledge Graph v1.0 loaded!")
print("   Storage: SQLite (local, free)")
print("   API: query_graph(query) → context string")
print("   Background indexing started...")