import ast
import datetime
import io
import re
import time
import traceback
from contextlib import ExitStack
from types import TracebackType
from typing import Any

import fused
from fused._options import options as OPTIONS
from fused._udf.compile_v2 import compile_udf_and_run_with_kwargs
from fused._udf.state import (
    decorator_src_override_context,
)
from fused.core.threading_utils import (
    capture_current_thread_output,
    isolate_streams_per_thread,
)
from fused.models.udf import EMPTY_UDF, AnyBaseUdf
from fused.models.udf._eval_result import UdfEvaluationResult

USER_CODE_DETECT_STRING = 'File "<udf '


def _detect_start_tb_line(tb: TracebackType | None) -> int | None:
    """Find the first line of the traceback that is in user code"""
    for line_no, line in enumerate(traceback.format_tb(tb)):
        if USER_CODE_DETECT_STRING in line:
            return line_no
    return None


def _extract_lineno(traceback_str: str) -> int | None:
    """Extract the line number from the first line of user code"""
    # Extract line number from error statement
    for line in traceback_str.splitlines():
        if USER_CODE_DETECT_STRING in line:
            match = re.search(r"line (\d+)", line)
            if match:
                return int(match.group(1))
    return None


def write_traceback_error(
    e: Exception, err_buf: io.StringIO, traceback_buf: io.StringIO
) -> None:
    """Write traceback info to two buffers to be used when setting metadata."""
    # First, make sure the traceback is provided to the user in stderr:
    # print_tb takes care of only printing from the current stack frame,
    # which minimizes the amount of Fused overhead in the logs.
    detected_start_tb_line = _detect_start_tb_line(e.__traceback__)
    if detected_start_tb_line is None:
        traceback.print_tb(e.__traceback__, file=err_buf)
        # Make an additional copy of the traceback for detecting the line number
        # in the UDF.
        traceback.print_tb(e.__traceback__, file=traceback_buf)
    else:
        traceback.print_list(
            traceback.extract_tb(e.__traceback__)[detected_start_tb_line:],
            file=err_buf,
        )
        traceback.print_list(
            traceback.extract_tb(e.__traceback__)[detected_start_tb_line:],
            file=traceback_buf,
        )
    err_buf.write(f"{traceback.format_exception_only(e)[-1].strip()}\n")


def execute_against_sample(
    udf: AnyBaseUdf,
    input: list[Any],
    validate_imports: bool | None = None,
    cache_max_age: int | None = None,
    **kwargs,
) -> UdfEvaluationResult:
    if udf is EMPTY_UDF:
        raise ValueError("Empty UDF cannot be evaluated. Use `set_udf` to set a UDF.")

    # Validate import statements correspond to valid modules
    validate_imports_whitelist(udf, validate_imports=validate_imports)

    if not OPTIONS.local_engine_cache:
        cache_max_age = 0

    if cache_max_age is None:
        if udf.cache_max_age is not None:
            cache_max_age = udf.cache_max_age
        else:
            cache_max_age = int(datetime.timedelta(days=90).total_seconds())

    # Run UDF with stdout/stderr capture
    _output = None
    has_exception = False
    errormsg = None
    exception_class = None
    lineno = None
    traceback_str = None
    stdout_content = None

    time_start = time.perf_counter()

    # Set up thread-local stream routing (reuses existing routers if already installed)
    with (
        isolate_streams_per_thread(),
        capture_current_thread_output() as (
            out_buf,
            err_buf,
        ),
    ):
        original_exception = None
        try:
            if cache_max_age == 0:
                _output = compile_udf_and_run_with_kwargs(udf, *input, **kwargs)
            else:
                fn = compile_udf_and_run_with_kwargs
                fn.__name__ = f"{udf.entrypoint}"
                wrapped = fused.cache(
                    fn,
                    cache_max_age=cache_max_age,
                    cache_verbose=False,
                )
                _output = wrapped(udf, *input, **kwargs)
        except Exception as exc:
            original_exception = exc

        # Capture logs and traceback after execution
        stdout_content = out_buf.getvalue()

        if original_exception is not None:
            has_exception = True
            exception_class = original_exception.__class__.__name__

            with io.StringIO() as traceback_buf:
                write_traceback_error(original_exception, err_buf, traceback_buf)
                # Get the full stderr content including the traceback
                traceback_str = err_buf.getvalue()

                # Extract line number from traceback
                if (
                    isinstance(original_exception, SyntaxError)
                    and original_exception.lineno is not None
                ):
                    lineno = original_exception.lineno
                else:
                    lineno = _extract_lineno(traceback_buf.getvalue())

            errormsg = f"{exception_class}: {str(original_exception)}"
        else:
            # No exception - capture stderr from normal execution
            traceback_str = err_buf.getvalue() or None

    time_end = time.perf_counter()
    time_taken_seconds = time_end - time_start

    # Always return UdfEvaluationResult - let the caller decide how to process it
    return UdfEvaluationResult(
        data=_output,
        udf=udf,
        time_taken_seconds=time_taken_seconds,
        stdout=stdout_content,
        stderr=traceback_str,
        error_message=errormsg,
        has_exception=has_exception,
        exception_class=exception_class,
        error_lineno=lineno,
    )


def execute_for_decorator(udf: AnyBaseUdf) -> AnyBaseUdf:
    """Evaluate a UDF for the purpose of getting the UDF object out of it."""
    # Define custom function in environment

    # This is a stripped-down version of execute_against_sample, above.

    src = udf.code

    exec_globals_locals = {"fused": fused}

    with ExitStack() as stack:
        stack.enter_context(decorator_src_override_context(src))

        # Add headers to sys.meta_path
        if udf.headers is not None:
            for header in udf.headers:
                stack.enter_context(header._register_custom_finder())

        exec(src, exec_globals_locals, exec_globals_locals)

        if udf.entrypoint not in exec_globals_locals:
            raise NameError(
                f"Could not find {udf.entrypoint}. You need to define a UDF with `def {udf.entrypoint}()`."
            )

        _fn = exec_globals_locals[udf.entrypoint]

        return _fn


def validate_imports_whitelist(udf: AnyBaseUdf, validate_imports: bool | None = None):
    # Skip import validation if the option is set
    if not fused.options.default_validate_imports and validate_imports is not True:
        return

    # Skip import validation if not logged in
    if not fused.api.AUTHORIZATION.is_configured():
        return

    from fused._global_api import get_api

    # Get the dependency whitelist from the cached API endpoint
    api = get_api()
    package_dependencies = api.dependency_whitelist()

    # Initialize a list to store the import statements
    import_statements = []

    # Parse the source code into an AST
    tree = ast.parse(udf.code)

    # Traverse the AST to find import statements
    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for alias in node.names:
                import_statements.append(alias.name)
        elif isinstance(node, ast.ImportFrom):
            module_name = node.module
            import_statements.append(module_name)

    # Check for unavailable modules
    header_modules = [header.module_name for header in udf.headers]
    fused_modules = ["fused"]  # assume fused is always available
    available_modules = (
        list(package_dependencies["dependency_whitelist"].keys())
        + header_modules
        + fused_modules
    )
    unavailable_modules = []
    for import_statement in import_statements:
        if import_statement.split(".", 1)[0] not in available_modules:
            unavailable_modules.append(import_statement)

    if unavailable_modules:
        raise ValueError(
            f"The following imports in the UDF might not be available: {repr(unavailable_modules)}. Please check the UDF headers and imports and try again."
        )

    # TODO: check major versions for some packages
