import copy
import io
import json
import os
import shutil
import sys
import tempfile
import threading
import time
import uuid
from collections import OrderedDict
from contextlib import contextmanager
from functools import lru_cache
from typing import TYPE_CHECKING, Any, Optional

from groknroll.core.comms_utils import LMRequest, send_lm_request, send_lm_request_batched
from groknroll.core.types import REPLResult, RLMChatCompletion
from groknroll.environments.base_env import NonIsolatedEnv

if TYPE_CHECKING:
    from groknroll.storage.database import Database

# Maximum number of local variables to retain in memory (prevents memory leaks)
# When database persistence is enabled, this is relaxed since context is stored in DB
MAX_LOCALS = 1000
MAX_LOCALS_WITH_DB = 10000  # Higher limit when DB persistence is enabled

# =============================================================================
# Safe Builtins
# =============================================================================

# Safe builtins - blocks dangerous operations like eval/exec/input
_SAFE_BUILTINS = {
    # Core types and functions
    "print": print,
    "len": len,
    "str": str,
    "int": int,
    "float": float,
    "list": list,
    "dict": dict,
    "set": set,
    "tuple": tuple,
    "bool": bool,
    "type": type,
    "isinstance": isinstance,
    "issubclass": issubclass,
    "enumerate": enumerate,
    "zip": zip,
    "map": map,
    "filter": filter,
    "sorted": sorted,
    "reversed": reversed,
    "range": range,
    "min": min,
    "max": max,
    "sum": sum,
    "abs": abs,
    "round": round,
    "any": any,
    "all": all,
    "pow": pow,
    "divmod": divmod,
    "chr": chr,
    "ord": ord,
    "hex": hex,
    "bin": bin,
    "oct": oct,
    "repr": repr,
    "ascii": ascii,
    "format": format,
    "hash": hash,
    "id": id,
    "iter": iter,
    "next": next,
    "slice": slice,
    "callable": callable,
    "hasattr": hasattr,
    "getattr": getattr,
    "setattr": setattr,
    "delattr": delattr,
    "dir": dir,
    "vars": vars,
    "bytes": bytes,
    "bytearray": bytearray,
    "memoryview": memoryview,
    "complex": complex,
    "object": object,
    "super": super,
    "property": property,
    "staticmethod": staticmethod,
    "classmethod": classmethod,
    "__import__": __import__,
    "open": open,
    # Exceptions
    "Exception": Exception,
    "BaseException": BaseException,
    "ValueError": ValueError,
    "TypeError": TypeError,
    "KeyError": KeyError,
    "IndexError": IndexError,
    "AttributeError": AttributeError,
    "FileNotFoundError": FileNotFoundError,
    "OSError": OSError,
    "IOError": IOError,
    "RuntimeError": RuntimeError,
    "NameError": NameError,
    "ImportError": ImportError,
    "StopIteration": StopIteration,
    "AssertionError": AssertionError,
    "NotImplementedError": NotImplementedError,
    "ArithmeticError": ArithmeticError,
    "LookupError": LookupError,
    "Warning": Warning,
    # Blocked
    "input": None,
    "eval": None,
    "exec": None,
    "compile": None,
    "globals": None,
    "locals": None,
}


class LocalREPL(NonIsolatedEnv):
    """
    Local REPL environment with persistent Python namespace.
    Executes code in a sandboxed namespace with access to context data.

    When database is provided, context is persisted to SQLite for
    100% contextual awareness across sessions. Variables are never
    evicted when DB persistence is enabled.
    """

    def __init__(
        self,
        lm_handler_address: tuple[str, int] | None = None,
        context_payload: dict | list | str | None = None,
        setup_code: str | None = None,
        persistent: bool = False,
        depth: int = 1,
        database: Optional["Database"] = None,
        project_id: Optional[int] = None,
        session_id: Optional[int] = None,
        **kwargs,
    ):
        super().__init__(persistent=persistent, depth=depth, **kwargs)

        self.lm_handler_address = lm_handler_address
        self.original_cwd = os.getcwd()
        self.temp_dir = tempfile.mkdtemp(prefix=f"repl_env_{uuid.uuid4()}_")
        self._lock = threading.Lock()
        self._context_count: int = 0
        self._history_count: int = 0

        # Database persistence
        self._database = database
        self._project_id = project_id
        self._session_id = session_id
        self._db_persistence_enabled = database is not None and project_id is not None

        # Lazy loading support - track known keys without loading content
        self._known_context_keys: set[str] = set()
        self._context_cache: dict[str, Any] = {}  # LRU-style cache for lazy loaded content
        self._pending_persists: list[dict] = []  # Batch persistence queue

        # Setup globals, locals, and modules in environment.
        self.setup()

        # Load persistent context from database if available
        if self._db_persistence_enabled:
            self._load_persistent_context()

        # Load context if provided
        if context_payload is not None:
            self.load_context(context_payload)

        # Run setup code if provided
        if setup_code:
            self.execute_code(setup_code)

    def setup(self):
        """Setup the environment."""
        # Create sandboxed globals
        self.globals: dict[str, Any] = {
            "__builtins__": _SAFE_BUILTINS.copy(),
            "__name__": "__main__",
        }
        # Use OrderedDict with bounded size to prevent memory leaks
        self.locals: OrderedDict[str, Any] = OrderedDict()

        # Track LLM calls made during code execution
        self._pending_llm_calls: list[RLMChatCompletion] = []

        # Add helper functions
        self.globals["FINAL_VAR"] = self._final_var
        self.globals["llm_query"] = self._llm_query
        self.globals["llm_query_batched"] = self._llm_query_batched

    def _final_var(self, variable_name: str) -> str:
        """Return the value of a variable as a final answer."""
        variable_name = variable_name.strip().strip("\"'")
        if variable_name in self.locals:
            return str(self.locals[variable_name])
        return f"Error: Variable '{variable_name}' not found"

    def _llm_query(self, prompt: str, model: str | None = None) -> str:
        """Query the LM via socket connection to the handler.

        Args:
            prompt: The prompt to send to the LM.
            model: Optional model name to use (if handler has multiple clients).
        """
        if not self.lm_handler_address:
            return "Error: No LM handler configured"

        try:
            request = LMRequest(prompt=prompt, model=model, depth=self.depth)
            response = send_lm_request(self.lm_handler_address, request)

            if not response.success:
                return f"Error: {response.error}"

            # Track this LLM call
            self._pending_llm_calls.append(
                response.chat_completion,
            )

            return response.chat_completion.response
        except Exception as e:
            return f"Error: LM query failed - {e}"

    def _llm_query_batched(self, prompts: list[str], model: str | None = None) -> list[str]:
        """Query the LM with multiple prompts concurrently.

        Args:
            prompts: List of prompts to send to the LM.
            model: Optional model name to use (if handler has multiple clients).

        Returns:
            List of responses in the same order as input prompts.
        """
        if not self.lm_handler_address:
            return ["Error: No LM handler configured"] * len(prompts)

        try:
            responses = send_lm_request_batched(
                self.lm_handler_address, prompts, model=model, depth=self.depth
            )

            results = []
            for response in responses:
                if not response.success:
                    results.append(f"Error: {response.error}")
                else:
                    # Track this LLM call in list of all calls -- we may want to do this hierarchically
                    self._pending_llm_calls.append(response.chat_completion)
                    results.append(response.chat_completion.response)

            return results
        except Exception as e:
            return [f"Error: LM query failed - {e}"] * len(prompts)

    def _load_persistent_context(self) -> None:
        """
        Load persistent context metadata from database (lazy loading).

        Only loads context keys at startup, not the full content.
        Content is loaded on-demand when accessed via _get_lazy_context().
        This dramatically improves startup time with large context stores.
        """
        if not self._db_persistence_enabled:
            return

        try:
            # Load only metadata (keys), not full content - search_contexts returns entries
            entries = self._database.search_contexts(self._project_id, limit=10000)

            for entry in entries:
                key = entry.context_key
                self._known_context_keys.add(key)

                # Track context/history counts from keys
                if key.startswith("context_"):
                    try:
                        idx = int(key.split("_")[1])
                        self._context_count = max(self._context_count, idx + 1)
                    except (ValueError, IndexError):
                        pass
                elif key.startswith("history_"):
                    try:
                        idx = int(key.split("_")[1])
                        self._history_count = max(self._history_count, idx + 1)
                    except (ValueError, IndexError):
                        pass

            # Load commonly used variables eagerly for quick access
            # Only load context_0, history_0 and their aliases immediately
            eager_keys = ["context_0", "history_0"]
            for key in eager_keys:
                if key in self._known_context_keys:
                    value = self._get_lazy_context(key)
                    if value is not None:
                        self.locals[key] = value

            # Set up aliases
            if "context_0" in self.locals and "context" not in self.locals:
                self.locals["context"] = self.locals["context_0"]
            if "history_0" in self.locals and "history" not in self.locals:
                self.locals["history"] = self.locals["history_0"]

        except Exception:
            # Silently fail if database loading fails - don't break REPL
            pass

    def _get_lazy_context(self, context_key: str) -> Optional[Any]:
        """
        Lazy load a context entry from database with caching.

        Args:
            context_key: The key to retrieve

        Returns:
            The deserialized content or None if not found
        """
        if not self._db_persistence_enabled:
            return None

        # Check cache first
        if context_key in self._context_cache:
            return self._context_cache[context_key]

        # Load from database
        value = self._database.get_context(self._project_id, context_key)
        if value is not None:
            # Cache for future access (limit cache size to 500 entries)
            if len(self._context_cache) >= 500:
                # Remove oldest entry (first key)
                oldest_key = next(iter(self._context_cache))
                del self._context_cache[oldest_key]
            self._context_cache[context_key] = value

        return value

    def _determine_context_type(self, var_name: str) -> str:
        """Determine the context type for a variable name."""
        if var_name.startswith("context_"):
            return "context"
        elif var_name.startswith("history_"):
            return "history"
        else:
            return "variable"

    def _persist_variable(self, var_name: str, value: Any) -> None:
        """Persist a variable to the database (immediate, single variable)."""
        if not self._db_persistence_enabled:
            return

        try:
            self._database.store_context(
                project_id=self._project_id,
                context_key=var_name,
                context_type=self._determine_context_type(var_name),
                content=value,
                variable_name=var_name,
                session_id=self._session_id,
            )
        except Exception:
            # Silently fail - don't break execution
            pass

    def _persist_variables_batch(self, variables: dict[str, Any]) -> None:
        """
        Persist multiple variables to the database in a single transaction.

        This is significantly faster than persisting variables individually,
        especially when executing code that creates many new variables.
        """
        if not self._db_persistence_enabled or not variables:
            return

        try:
            entries = [
                {
                    "context_key": var_name,
                    "context_type": self._determine_context_type(var_name),
                    "content": value,
                    "variable_name": var_name,
                }
                for var_name, value in variables.items()
            ]
            self._database.store_contexts_batch(
                project_id=self._project_id,
                entries=entries,
                session_id=self._session_id,
            )
        except Exception:
            # Silently fail - don't break execution
            pass

    def load_context(self, context_payload: dict | list | str):
        """Load context into the environment as context_0 (and 'context' alias)."""
        self.add_context(context_payload, 0)

    def add_context(
        self, context_payload: dict | list | str, context_index: int | None = None
    ) -> int:
        """
        Add a context with versioned variable name.

        Args:
            context_payload: The context data to add
            context_index: Optional explicit index. If None, auto-increments.

        Returns:
            The context index used.
        """
        if context_index is None:
            context_index = self._context_count

        var_name = f"context_{context_index}"

        if isinstance(context_payload, str):
            context_path = os.path.join(self.temp_dir, f"context_{context_index}.txt")
            with open(context_path, "w") as f:
                f.write(context_payload)
            self.execute_code(f"with open(r'{context_path}', 'r') as f:\n    {var_name} = f.read()")
        else:
            context_path = os.path.join(self.temp_dir, f"context_{context_index}.json")
            with open(context_path, "w") as f:
                json.dump(context_payload, f)
            self.execute_code(
                f"import json\nwith open(r'{context_path}', 'r') as f:\n    {var_name} = json.load(f)"
            )

        # Alias context_0 as 'context' for backward compatibility
        if context_index == 0:
            self.execute_code(f"context = {var_name}")

        # Persist context to database
        if self._db_persistence_enabled:
            self._persist_variable(var_name, context_payload)

        self._context_count = max(self._context_count, context_index + 1)
        return context_index

    def update_handler_address(self, address: tuple[str, int]) -> None:
        """Update the LM handler address for a new completion call."""
        self.lm_handler_address = address

    def get_context_count(self) -> int:
        """Return the number of contexts loaded."""
        return self._context_count

    def add_history(
        self, message_history: list[dict[str, Any]], history_index: int | None = None
    ) -> int:
        """
        Store a conversation's message history as a versioned variable.

        Args:
            message_history: The list of message dicts from a completion call
            history_index: Optional explicit index. If None, auto-increments.

        Returns:
            The history index used.
        """
        if history_index is None:
            history_index = self._history_count

        var_name = f"history_{history_index}"

        # Store deep copy to avoid reference issues with nested dicts
        history_copy = copy.deepcopy(message_history)
        self.locals[var_name] = history_copy

        # Alias history_0 as 'history' for convenience
        if history_index == 0:
            self.locals["history"] = self.locals[var_name]

        # Persist history to database
        if self._db_persistence_enabled:
            self._persist_variable(var_name, history_copy)

        self._history_count = max(self._history_count, history_index + 1)
        return history_index

    def get_history_count(self) -> int:
        """Return the number of conversation histories stored."""
        return self._history_count

    @contextmanager
    def _capture_output(self):
        """Thread-safe context manager to capture stdout/stderr."""
        with self._lock:
            old_stdout, old_stderr = sys.stdout, sys.stderr
            stdout_buf, stderr_buf = io.StringIO(), io.StringIO()
            try:
                sys.stdout, sys.stderr = stdout_buf, stderr_buf
                yield stdout_buf, stderr_buf
            finally:
                sys.stdout, sys.stderr = old_stdout, old_stderr

    @contextmanager
    def _temp_cwd(self):
        """Temporarily change to temp directory for execution."""
        old_cwd = os.getcwd()
        try:
            os.chdir(self.temp_dir)
            yield
        finally:
            os.chdir(old_cwd)

    def execute_code(self, code: str) -> REPLResult:
        """Execute code in the persistent namespace and return result."""
        start_time = time.perf_counter()

        # Clear pending LLM calls from previous execution
        self._pending_llm_calls = []

        # Use higher limit when DB persistence is enabled
        max_locals = MAX_LOCALS_WITH_DB if self._db_persistence_enabled else MAX_LOCALS

        with self._capture_output() as (stdout_buf, stderr_buf), self._temp_cwd():
            try:
                combined = {**self.globals, **self.locals}
                exec(code, combined, combined)

                # Collect new/updated variables for batch persistence
                variables_to_persist: dict[str, Any] = {}

                # Update locals with new variables (bounded to prevent memory leaks)
                for key, value in combined.items():
                    if key not in self.globals and not key.startswith("_"):
                        # Move to end if exists (LRU behavior)
                        if key in self.locals:
                            self.locals.move_to_end(key)
                        self.locals[key] = value

                        # Queue for batch persistence
                        if self._db_persistence_enabled:
                            variables_to_persist[key] = value

                # Batch persist all new/updated variables in single transaction
                if variables_to_persist:
                    self._persist_variables_batch(variables_to_persist)

                # Trim oldest entries if over limit
                while len(self.locals) > max_locals:
                    self.locals.popitem(last=False)  # Remove oldest

                stdout = stdout_buf.getvalue()
                stderr = stderr_buf.getvalue()
            except Exception as e:
                stdout = stdout_buf.getvalue()
                stderr = stderr_buf.getvalue() + f"\n{type(e).__name__}: {e}"

        return REPLResult(
            stdout=stdout,
            stderr=stderr,
            locals=self.locals.copy(),
            execution_time=time.perf_counter() - start_time,
            rlm_calls=self._pending_llm_calls.copy(),
        )

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.cleanup()
        return False

    def cleanup(self):
        """Clean up temp directory and reset state."""
        try:
            shutil.rmtree(self.temp_dir)
        except Exception:
            pass
        self.globals.clear()
        self.locals.clear()

    def __del__(self):
        self.cleanup()
