import re
import ast
import sqlglot
from typing import Dict, List, Any
from pathlib import Path

def extract_sql_schema(sql_code: str) -> Dict[str, List[str]]:
    """
    Extracts table -> columns mapping from SQL code using sqlglot.
    Returns a canonical schema.
    """
    schema: Dict[str, List[str]] = {}
    try:
        statements = sqlglot.parse(sql_code, read="postgres")
    except Exception:
        return schema

    for stmt in statements:
        if stmt.key and stmt.key.upper() == "CREATE":
            for table in stmt.find_all(sqlglot.exp.Create):
                try:
                    tname = table.this.this
                    cols = []
                    for coldef in table.find_all(sqlglot.exp.ColumnDef):
                        cname = getattr(coldef.this, "name", None)
                        if cname:
                            cols.append(cname)
                    if tname and cols:
                        schema[tname] = cols
                except Exception:
                    continue
    return schema


def extract_python_structure(code: str) -> Dict[str, List[str]]:
    """Return module structure: functions and classes."""
    try:
        tree = ast.parse(code)
    except SyntaxError:
        return {}

    funcs = [n.name for n in ast.walk(tree) if isinstance(n, ast.FunctionDef)]
    classes = [n.name for n in ast.walk(tree) if isinstance(n, ast.ClassDef)]
    return {"functions": funcs, "classes": classes}


def extract_go_structure(code: str) -> Dict[str, List[str]]:
    """Simple regex-based Go structure detection (lightweight)."""
    funcs = re.findall(r"func\s+([A-Z]\w+)", code)
    structs = re.findall(r"type\s+(\w+)\s+struct", code)
    return {"functions": funcs, "structs": structs}


def extract_rust_structure(code: str) -> Dict[str, List[str]]:
    """Heuristic Rust index."""
    structs = re.findall(r"struct\s+(\w+)", code)
    traits = re.findall(r"trait\s+(\w+)", code)
    funcs = re.findall(r"fn\s+(\w+)", code)
    return {"functions": funcs, "structs": structs, "traits": traits}


def extract_svelte_structure(code: str) -> Dict[str, List[str]]:
    """Minimal Svelte export/prop finder."""
    exports = re.findall(r"export\s+let\s+(\w+)", code)
    funcs = re.findall(r"function\s+(\w+)", code)
    return {"props": exports, "functions": funcs}


def build_quick_index(language: str, code: str, filepath: Path) -> Dict[str, Any]:
    """Dispatch to the appropriate structure extractor."""
    data = {}
    if language == "sql":
        data = extract_sql_schema(code)
    elif language == "python":
        data = extract_python_structure(code)
    elif language == "go":
        data = extract_go_structure(code)
    elif language == "rust":
        data = extract_rust_structure(code)
    elif language == "svelte":
        data = extract_svelte_structure(code)

    return {
        "text": f"Quick index for {filepath.name}:\n{data}",
        "metadata": {"language": language, "file": str(filepath), "type": "index"}
    }