"""
extractor.py — Parses LLM output text using regex to extract code references:
functions, classes, file paths, import statements, and method calls.
"""

from __future__ import annotations

import re
from dataclasses import dataclass
from enum import Enum
from typing import List, Optional


class ReferenceKind(str, Enum):
    FUNCTION = "function"
    CLASS = "class"
    FILE = "file"
    IMPORT = "import"
    METHOD = "method"
    VARIABLE = "variable"


@dataclass
class CodeReference:
    kind: ReferenceKind
    name: str
    line_number: int
    raw_text: str
    # For imports: the module and the specific name imported
    import_module: Optional[str] = None
    import_name: Optional[str] = None


# ---------------------------------------------------------------------------
# Regex patterns
# ---------------------------------------------------------------------------

# Inline code spans (backtick-wrapped): `some_func()` or `SomeClass`
_BACKTICK = r"`([^`]+)`"

# Function call patterns — match name( in prose or code blocks
# e.g. get_user_data(), fetchResults(), some_func(arg)
_FUNC_CALL_PATTERN = re.compile(
    r"\b([a-z_][a-zA-Z0-9_]*)\s*\(", re.MULTILINE
)

# PascalCase class-like identifiers in prose
# e.g. UserManager, DataProcessor, MyClass
_CLASS_PATTERN = re.compile(
    r"\b([A-Z][a-zA-Z0-9]*(?:[A-Z][a-zA-Z0-9]*)+)\b"
)

# File path patterns — relative paths with extensions
# e.g. src/utils.py, config/settings.json, ./helpers/format.py, ./manage.py
_FILE_PATH_PATTERN = re.compile(
    r"(?:^|[\s`\"'(])(\./|\.\./)?((?:[a-zA-Z0-9_\-]+/)+[a-zA-Z0-9_\-]+\.[a-zA-Z]{1,6})"
    r"(?:[\s`\"'),]|$)",
    re.MULTILINE,
)

# Also match bare ./filename.ext pattern (single file without directory)
_DOT_SLASH_FILE_PATTERN = re.compile(
    r"(?:^|[\s`\"'(])(\./)([a-zA-Z0-9_\-]+\.[a-zA-Z]{1,6})(?:[\s`\"'),]|$)",
    re.MULTILINE,
)

# Bare filename with extension in backticks or quotes
_BARE_FILE_PATTERN = re.compile(
    r'[`"\']([a-zA-Z0-9_\-]+\.(?:py|js|ts|json|yaml|yml|toml|txt|md|sh|env|cfg|ini))[`"\']'
)

# Import statements  (from X import Y  or  import X)
_FROM_IMPORT_PATTERN = re.compile(
    r"from\s+([\w.]+)\s+import\s+([\w,\s*]+)"
)
_IMPORT_PATTERN = re.compile(r"^import\s+([\w.,\s]+)", re.MULTILINE)

# Method call: .method_name() in text
_METHOD_CALL_PATTERN = re.compile(r"\.(save|load|delete|update|process|run|execute|start|stop|get|set|create|fetch|send|receive|parse|format|validate|check|test|build|init|setup|teardown|connect|disconnect|open|close|read|write|flush|clear|reset|refresh|sync|push|pull|merge|commit|rollback|query|filter|order|group|join|select|insert|update|delete|find|search|sort|map|reduce|apply|transform|convert|encode|decode|serialize|deserialize|render|display|show|hide|toggle|enable|disable|add|remove|append|extend|insert|pop|push|shift|unshift|slice|splice|concat|split|join|trim|strip|replace|match|test|search|index|count|sum|min|max|avg|mean|std|var|log|warn|error|info|debug|trace|assert|raise|throw|catch|handle|wrap|unwrap|pack|unpack|zip|unzip|compress|decompress|encrypt|decrypt|sign|verify|hash|digest|random|seed|generate|sample|shuffle|choose|pick|select|drop|keep|skip|take|limit|offset|paginate|chunk|batch|retry|timeout|cancel|abort|resume|pause|wait|sleep|delay|schedule|queue|dequeue|enqueue|peek|poll|subscribe|publish|emit|broadcast|notify|observe|watch|listen|trigger|fire|dispatch|redirect|navigate|route|resolve|reject|fulfill|then|catch|finally|next|prev|first|last|head|tail|keys|values|items|entries|pairs|clone|copy|deepcopy|merge|diff|patch|apply|revert|undo|redo|save|load|dump|restore|backup|archive|extract|compress|decompress|upload|download|import_data|export_data|to_json|from_json|to_dict|from_dict|to_list|from_list|to_string|from_string|to_bytes|from_bytes)\s*\(", re.IGNORECASE)  # noqa: E501

# Generic method call on an object: obj.method()
_GENERIC_METHOD_PATTERN = re.compile(
    r"\b[a-z_][a-zA-Z0-9_]*\.([a-z_][a-zA-Z0-9_]*)\s*\("
)

# Variables / constants (UPPER_CASE or snake_case mentioned in backticks)
_VARIABLE_IN_BACKTICK = re.compile(r"`([A-Z_][A-Z0-9_]{2,})`")

# Code fences — we want to extract references from both prose and code blocks
_CODE_FENCE = re.compile(r"```[\w]*\n(.*?)```", re.DOTALL)

# Words that are almost certainly not real code references
_STOP_WORDS: set = {
    # Python builtins
    "print", "len", "range", "list", "dict", "set", "tuple", "str", "int",
    "float", "bool", "type", "isinstance", "hasattr", "getattr", "setattr",
    "delattr", "open", "close", "read", "write", "input", "super", "object",
    "property", "classmethod", "staticmethod", "abstractmethod",
    # Common keywords that look like function calls
    "if", "for", "while", "with", "return", "yield", "raise", "assert",
    "import", "from", "as", "pass", "break", "continue", "lambda",
    "not", "and", "or", "in", "is", "del", "global", "nonlocal",
    # Very generic method names that appear everywhere
    "format", "encode", "decode", "strip", "split", "join", "lower",
    "upper", "replace", "find", "index", "count", "append", "extend",
    "insert", "remove", "pop", "clear", "copy", "update", "get",
    "items", "keys", "values", "sort", "sorted", "reversed", "map",
    "filter", "zip", "enumerate", "any", "all", "min", "max", "sum",
    # Testing helpers
    "assertEqual", "assertTrue", "assertFalse", "assertRaises",
    "setUp", "tearDown", "describe", "it", "expect", "test",
}

# PascalCase words that are stdlib/built-in types, not user classes
_BUILTIN_CLASSES: set = {
    "True", "False", "None", "Exception", "ValueError", "TypeError",
    "KeyError", "IndexError", "AttributeError", "ImportError",
    "RuntimeError", "StopIteration", "NotImplementedError", "OSError",
    "IOError", "FileNotFoundError", "PermissionError", "TimeoutError",
    "ConnectionError", "OverflowError", "MemoryError", "RecursionError",
    "SystemExit", "KeyboardInterrupt", "GeneratorExit", "BaseException",
    "ArithmeticError", "LookupError", "EnvironmentError", "UnicodeError",
    "UnicodeDecodeError", "UnicodeEncodeError", "UnicodeTranslateError",
    "BufferError", "EOFError", "FloatingPointError", "ZeroDivisionError",
    "NameError", "UnboundLocalError", "ReferenceError", "SyntaxError",
    "IndentationError", "TabError", "SystemError", "AssertionError",
    "NotImplemented",
    # Common stdlib classes
    "Path", "PurePath", "PosixPath", "WindowsPath",
    "Thread", "Process", "Queue", "Lock", "Event",
    "datetime", "date", "time", "timedelta", "timezone",
    "Decimal", "Fraction", "Counter", "OrderedDict", "defaultdict",
    "deque", "namedtuple", "dataclass",
    "Enum", "IntEnum", "Flag", "IntFlag",
    "ABC", "ABCMeta",
    "Optional", "Union", "List", "Dict", "Set", "Tuple", "Any",
    "Type", "Callable", "Iterator", "Generator", "Iterable",
    "ClassVar", "Final", "Literal",
    "Request", "Response", "Session", "Client",
    "DataFrame", "Series", "Index", "MultiIndex",
    "Model", "Field", "Schema", "Base",
    "TestCase", "TestSuite", "TestRunner", "TestResult",
    "Logger", "Handler", "Formatter", "Filter",
    "Socket", "SSLSocket", "SSLContext",
    "StringIO", "BytesIO", "FileIO", "BufferedReader", "BufferedWriter",
    "Json", "Yaml",
    "App", "Blueprint", "Router", "APIRouter",
}


class ReferenceExtractor:
    """Extracts code references from LLM-generated prose."""

    def __init__(self, checks: Optional[List[str]] = None) -> None:
        """
        Args:
            checks: Optional list of check types to include.
                    Values: 'functions', 'classes', 'files', 'imports', 'methods'
                    If None, all checks are run.
        """
        self._checks = set(checks) if checks else None

    def _check_enabled(self, kind: str) -> bool:
        if self._checks is None:
            return True
        return kind in self._checks

    def extract(self, text: str) -> List[CodeReference]:
        """Extract all code references from text."""
        refs: List[CodeReference] = []
        lines = text.splitlines()

        for lineno, line in enumerate(lines, start=1):
            if self._check_enabled("imports"):
                refs.extend(self._extract_imports(line, lineno))
            if self._check_enabled("files"):
                refs.extend(self._extract_files(line, lineno))
            if self._check_enabled("functions"):
                refs.extend(self._extract_functions(line, lineno))
            if self._check_enabled("classes"):
                refs.extend(self._extract_classes(line, lineno))
            if self._check_enabled("methods"):
                refs.extend(self._extract_methods(line, lineno))
            if self._check_enabled("variables"):
                refs.extend(self._extract_variables(line, lineno))

        # Deduplicate: keep first occurrence per (kind, name)
        seen: set = set()
        unique: List[CodeReference] = []
        for ref in refs:
            key = (ref.kind, ref.name)
            if key not in seen:
                seen.add(key)
                unique.append(ref)

        return unique

    # ------------------------------------------------------------------
    # Per-kind extractors
    # ------------------------------------------------------------------

    def _extract_imports(self, line: str, lineno: int) -> List[CodeReference]:
        refs = []

        for m in _FROM_IMPORT_PATTERN.finditer(line):
            module = m.group(1).strip()
            names_raw = m.group(2).strip()

            # Handle "from x import a, b, c" and "from x import *"
            if names_raw == "*":
                names = ["*"]
            else:
                names = [n.strip() for n in names_raw.split(",") if n.strip()]

            for name in names:
                if name and name not in _STOP_WORDS:
                    refs.append(
                        CodeReference(
                            kind=ReferenceKind.IMPORT,
                            name=f"{module}.{name}" if name != "*" else module,
                            line_number=lineno,
                            raw_text=m.group(0).strip(),
                            import_module=module,
                            import_name=name if name != "*" else None,
                        )
                    )

        for m in _IMPORT_PATTERN.finditer(line):
            modules_raw = m.group(1).strip()
            for mod in modules_raw.split(","):
                mod = mod.strip()
                if mod:
                    refs.append(
                        CodeReference(
                            kind=ReferenceKind.IMPORT,
                            name=mod,
                            line_number=lineno,
                            raw_text=m.group(0).strip(),
                            import_module=mod,
                            import_name=None,
                        )
                    )

        return refs

    def _extract_files(self, line: str, lineno: int) -> List[CodeReference]:
        refs = []

        for m in _FILE_PATH_PATTERN.finditer(line):
            prefix = m.group(1) or ""
            path = prefix + m.group(2)
            refs.append(
                CodeReference(
                    kind=ReferenceKind.FILE,
                    name=path,
                    line_number=lineno,
                    raw_text=path,
                )
            )

        for m in _BARE_FILE_PATTERN.finditer(line):
            refs.append(
                CodeReference(
                    kind=ReferenceKind.FILE,
                    name=m.group(1),
                    line_number=lineno,
                    raw_text=m.group(0),
                )
            )

        for m in _DOT_SLASH_FILE_PATTERN.finditer(line):
            path = m.group(1) + m.group(2)  # e.g. "./manage.py"
            refs.append(
                CodeReference(
                    kind=ReferenceKind.FILE,
                    name=path,
                    line_number=lineno,
                    raw_text=path,
                )
            )

        return refs

    def _extract_functions(self, line: str, lineno: int) -> List[CodeReference]:
        refs = []

        for m in _FUNC_CALL_PATTERN.finditer(line):
            name = m.group(1)
            if name in _STOP_WORDS:
                continue
            if len(name) < 3:
                continue
            # Skip if it's a builtin class being called as a constructor
            if name[0].isupper():
                continue
            refs.append(
                CodeReference(
                    kind=ReferenceKind.FUNCTION,
                    name=name,
                    line_number=lineno,
                    raw_text=m.group(0).rstrip("("),
                )
            )

        return refs

    def _extract_classes(self, line: str, lineno: int) -> List[CodeReference]:
        refs = []

        for m in _CLASS_PATTERN.finditer(line):
            name = m.group(1)
            if name in _BUILTIN_CLASSES:
                continue
            if name in _STOP_WORDS:
                continue
            # Must have at least one lowercase letter after the first uppercase
            # to avoid matching acronyms like HTTP, URL, API
            has_lower = any(c.islower() for c in name[1:])
            if not has_lower:
                continue
            refs.append(
                CodeReference(
                    kind=ReferenceKind.CLASS,
                    name=name,
                    line_number=lineno,
                    raw_text=name,
                )
            )

        return refs

    def _extract_methods(self, line: str, lineno: int) -> List[CodeReference]:
        refs = []

        for m in _GENERIC_METHOD_PATTERN.finditer(line):
            name = m.group(1)
            if name in _STOP_WORDS:
                continue
            if len(name) < 3:
                continue
            refs.append(
                CodeReference(
                    kind=ReferenceKind.METHOD,
                    name=name,
                    line_number=lineno,
                    raw_text=f".{name}()",
                )
            )

        return refs

    def _extract_variables(self, line: str, lineno: int) -> List[CodeReference]:
        refs = []

        for m in _VARIABLE_IN_BACKTICK.finditer(line):
            name = m.group(1)
            if name in _STOP_WORDS:
                continue
            refs.append(
                CodeReference(
                    kind=ReferenceKind.VARIABLE,
                    name=name,
                    line_number=lineno,
                    raw_text=m.group(0),
                )
            )

        return refs
