"""
REPL Environment - Recursive Language Model Execution Engine

Implements Section 2 of the research paper:
"The Prompt-as-Environment Abstraction" and "The Recursive Loop"

Key Concepts:
1. DOC_VAR: The document is loaded as a persistent variable, NOT passed to LLM context
2. Code Execution: LLM generates Python code to interact with DOC_VAR
3. Variable Stitching: Intermediate results frozen into variables to prevent Context Rot
4. Recursive Calls: Agent can invoke sub-LLM instances for sub-tasks

Per Zhang et al. (2025):
"The RLM initializes a computational environment where the massive input context 
(the 'Long Prompt') is loaded into a variable, typically denoted as P or DOC_VAR.
The Neural Network itself does not ingest P. Instead, it ingests a system instruction 
that informs it of P's existence and provides a set of tools to interact with it."

This decouples Memory (REPL state) from Processing (LLM inference), enabling
effectively unbounded context lengths.
"""

from __future__ import annotations

import asyncio
import hashlib
import re
import traceback
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Callable

import structlog

from rnsr.agent.variable_store import VariableStore, generate_pointer_name
from rnsr.indexing.kv_store import KVStore
from rnsr.models import SkeletonNode

logger = structlog.get_logger(__name__)


# =============================================================================
# RLM System Prompt - Informs LLM of DOC_VAR and available tools
# =============================================================================

RLM_SYSTEM_PROMPT = """You are a Recursive Language Model (RLM) operating in a Python REPL environment.

CRITICAL: You do NOT have the document in your context. The document is stored in the variable DOC_VAR.
You interact with the document by writing Python code that will be executed in the REPL.

## Available Variables:
- DOC_VAR: The full document text (string). Use slicing to access portions.
- DOC_TREE: Hierarchical tree structure of the document (dict).
- VARIABLES: Dictionary storing your findings (use store_variable() to add).

## Available Functions:
- len(DOC_VAR): Get document length in characters
- DOC_VAR[i:j]: Slice document to get characters from i to j
- search_text(pattern): Search DOC_VAR for regex pattern, returns list of (start, end, match)
- split_by_regex(pattern): Split DOC_VAR by pattern, returns list of segments
- list_children(node_id): Get summaries of child nodes in DOC_TREE
- read_node(node_id): Get full text content of a node
- store_variable(name, content): Store content as $NAME for later synthesis
- get_variable(name): Retrieve stored variable content
- sub_llm(prompt, context): Invoke sub-LLM to process context with prompt (for decomposition)
- batch_sub_llm(prompts, contexts): Batch process multiple sub-tasks in parallel

## Workflow for Complex Queries:
1. First, explore the document structure: `list_children('root')`
2. Navigate to relevant sections based on query
3. For multi-part queries, decompose and use sub_llm():
   ```python
   # Example: Compare clauses across documents
   clause_2023 = sub_llm("Extract the indemnification clause", read_node('2023_legal'))
   clause_2024 = sub_llm("Extract the indemnification clause", read_node('2024_legal'))
   store_variable("CLAUSE_2023", clause_2023)
   store_variable("CLAUSE_2024", clause_2024)
   ```
4. For bulk processing, use batch_sub_llm():
   ```python
   # Process 50 contracts in parallel batches
   clauses = batch_sub_llm(
       prompts=["Extract liability clause"] * 50,
       contexts=[read_node(f'contract_{i}') for i in range(50)]
   )
   ```

## Output:
Write Python code blocks to execute. The REPL will run them and return results.
When ready to synthesize, call: synthesize_answer(query, variable_names)

IMPORTANT: 
- Always use store_variable() to save findings - this prevents Context Rot
- Prefer batch_sub_llm() over sequential sub_llm() calls for efficiency
- Target ~200k characters per batch for optimal token density
"""


# =============================================================================
# REPL Environment Class
# =============================================================================


@dataclass
class REPLEnvironment:
    """
    Python REPL Environment for Recursive Language Model execution.
    
    Implements the "Prompt-as-Environment" abstraction from Section 2.1:
    - DOC_VAR: Document loaded as persistent variable
    - Code execution sandbox for LLM-generated code
    - Variable stitching for intermediate results
    - Recursive sub-LLM invocation
    
    Example:
        env = REPLEnvironment(
            document_text=full_text,
            skeleton=skeleton_index,
            kv_store=kv_store,
        )
        
        # LLM generates code, REPL executes it
        result = env.execute("len(DOC_VAR)")
        # -> 152347
        
        result = env.execute("list_children('root')")
        # -> [{'id': 'sec_001', 'header': 'Introduction', ...}, ...]
    """
    
    # The document as persistent variable (NOT in LLM context)
    document_text: str
    skeleton: dict[str, SkeletonNode]
    kv_store: KVStore
    
    # Variable store for findings
    variable_store: VariableStore = field(default_factory=VariableStore)
    
    # Execution state
    execution_history: list[dict[str, Any]] = field(default_factory=list)
    max_output_length: int = 10000  # Truncate long outputs
    
    # Batching configuration
    batch_size: int = 5  # Process N items per batch
    max_parallel_batches: int = 4  # Max concurrent batch calls
    optimal_chars_per_call: int = 200_000  # Target ~200k chars per LLM call
    
    # LLM function for sub-calls (injected)
    _llm_fn: Callable[[str], str] | None = None
    _async_llm_fn: Callable[[str], Any] | None = None
    
    def __post_init__(self):
        """Build the execution namespace."""
        self._namespace = self._build_namespace()
        logger.info(
            "repl_initialized",
            doc_length=len(self.document_text),
            num_nodes=len(self.skeleton),
        )
    
    def _build_namespace(self) -> dict[str, Any]:
        """Build the Python namespace for code execution."""
        return {
            # Core variables (the "environment")
            "DOC_VAR": self.document_text,
            "DOC_TREE": self._build_tree_dict(),
            "VARIABLES": {},
            
            # Built-in functions available in namespace
            "len": len,
            "print": print,
            "str": str,
            "int": int,
            "float": float,
            "list": list,
            "dict": dict,
            "range": range,
            "enumerate": enumerate,
            "zip": zip,
            "sorted": sorted,
            "min": min,
            "max": max,
            "sum": sum,
            "abs": abs,
            "round": round,
            "re": re,  # For regex operations
            
            # Navigator functions
            "search_text": self._search_text,
            "split_by_regex": self._split_by_regex,
            "list_children": self._list_children,
            "read_node": self._read_node,
            "get_node_path": self._get_node_path,
            
            # Variable stitching functions
            "store_variable": self._store_variable,
            "get_variable": self._get_variable,
            "list_variables": self._list_variables,
            
            # Recursive LLM functions
            "sub_llm": self._sub_llm,
            "batch_sub_llm": self._batch_sub_llm,
            
            # Synthesis
            "synthesize_answer": self._synthesize_answer,
        }
    
    def _build_tree_dict(self) -> dict[str, Any]:
        """Convert skeleton to navigable dict structure."""
        tree = {}
        for node_id, node in self.skeleton.items():
            tree[node_id] = {
                "id": node.node_id,
                "header": node.header,
                "summary": node.summary,
                "level": node.level,
                "children": node.child_ids,
                "parent": node.parent_id,
            }
        return tree
    
    # =========================================================================
    # Core Execution
    # =========================================================================
    
    def execute(self, code: str) -> dict[str, Any]:
        """
        Execute Python code in the REPL environment.
        
        Args:
            code: Python code to execute.
            
        Returns:
            Dict with 'success', 'output', 'error', and 'variables'.
        """
        start_time = datetime.now(timezone.utc)
        
        # Clean code (remove markdown code blocks if present)
        code = self._clean_code(code)
        
        result = {
            "success": False,
            "output": None,
            "error": None,
            "variables": list(self._namespace.get("VARIABLES", {}).keys()),
            "execution_time_ms": 0,
        }
        
        try:
            # Capture output
            import io
            import sys
            
            old_stdout = sys.stdout
            sys.stdout = captured_output = io.StringIO()
            
            try:
                # Try as expression first (returns value)
                compiled = compile(code, "<repl>", "eval")
                output = eval(compiled, self._namespace)
                result["output"] = self._format_output(output)
                result["success"] = True
            except SyntaxError:
                # Execute as statements
                exec(code, self._namespace)
                printed = captured_output.getvalue()
                result["output"] = printed if printed else "Executed successfully"
                result["success"] = True
            finally:
                sys.stdout = old_stdout
            
            # Update VARIABLES reference
            result["variables"] = list(self._namespace.get("VARIABLES", {}).keys())
            
        except Exception as e:
            result["error"] = f"{type(e).__name__}: {str(e)}\n{traceback.format_exc()}"
            logger.warning("repl_execution_error", error=str(e), code=code[:200])
        
        # Record timing
        elapsed = (datetime.now(timezone.utc) - start_time).total_seconds() * 1000
        result["execution_time_ms"] = round(elapsed, 2)
        
        # Record in history
        self.execution_history.append({
            "code": code,
            "result": result,
            "timestamp": start_time.isoformat(),
        })
        
        logger.debug(
            "repl_executed",
            success=result["success"],
            output_length=len(str(result["output"])) if result["output"] else 0,
        )
        
        return result
    
    def _clean_code(self, code: str) -> str:
        """Remove markdown code blocks and clean whitespace."""
        # Remove ```python ... ``` blocks
        code = re.sub(r"```python\s*", "", code)
        code = re.sub(r"```\s*", "", code)
        return code.strip()
    
    def _format_output(self, output: Any) -> str:
        """Format output for display, truncating if needed."""
        if output is None:
            return "None"
        
        output_str = str(output)
        
        if len(output_str) > self.max_output_length:
            truncated = output_str[:self.max_output_length]
            return f"{truncated}\n... [truncated, {len(output_str)} total chars]"
        
        return output_str
    
    # =========================================================================
    # Navigator Functions (exposed to REPL)
    # =========================================================================
    
    def _search_text(self, pattern: str) -> list[tuple[int, int, str]]:
        """
        Search DOC_VAR for regex pattern.
        
        Returns list of (start_pos, end_pos, matched_text).
        """
        matches = []
        for match in re.finditer(pattern, self.document_text, re.IGNORECASE):
            matches.append((match.start(), match.end(), match.group()))
        return matches[:100]  # Limit results
    
    def _split_by_regex(self, pattern: str) -> list[str]:
        """Split DOC_VAR by regex pattern."""
        return re.split(pattern, self.document_text)
    
    def _list_children(self, node_id: str) -> list[dict[str, Any]]:
        """Get summaries of child nodes."""
        node = self.skeleton.get(node_id)
        if not node:
            return []
        
        children = []
        for child_id in node.child_ids:
            child = self.skeleton.get(child_id)
            if child:
                children.append({
                    "id": child.node_id,
                    "header": child.header,
                    "summary": child.summary[:200] if child.summary else "",
                    "has_children": len(child.child_ids) > 0,
                })
        return children
    
    def _read_node(self, node_id: str) -> str:
        """Get full text content of a node."""
        content = self.kv_store.get(node_id)
        if content:
            return content
        
        # Fallback to skeleton summary
        node = self.skeleton.get(node_id)
        return node.summary if node else ""
    
    def _get_node_path(self, node_id: str) -> list[str]:
        """Get path from root to node."""
        path = []
        current = node_id
        while current:
            node = self.skeleton.get(current)
            if node:
                path.insert(0, node.header)
                current = node.parent_id
            else:
                break
        return path
    
    # =========================================================================
    # Variable Stitching Functions (Section 2.2)
    # =========================================================================
    
    def _store_variable(self, name: str, content: str) -> str:
        """
        Store content as a named variable.
        
        This implements Variable Stitching to prevent Context Rot.
        Intermediate results are "frozen" into variables.
        """
        if not name.startswith("$"):
            name = "$" + name.upper().replace(" ", "_")
        
        self.variable_store.assign(name, content, source_node_id="repl")
        self._namespace["VARIABLES"][name] = content
        
        logger.info("variable_stored", name=name, length=len(content))
        return name
    
    def _get_variable(self, name: str) -> str | None:
        """Retrieve a stored variable."""
        if not name.startswith("$"):
            name = "$" + name.upper()
        
        return self.variable_store.resolve(name)
    
    def _list_variables(self) -> list[str]:
        """List all stored variables."""
        return [v.pointer for v in self.variable_store.list_variables()]
    
    # =========================================================================
    # Recursive LLM Functions (Section 2.2)
    # =========================================================================
    
    def _sub_llm(self, prompt: str, context: str = "") -> str:
        """
        Invoke a sub-LLM call for a specific sub-task.
        
        This is the recursive mechanism: the model can call itself
        with a specific sub-prompt to solve sub-problems.
        
        Example:
            liability_1 = sub_llm("Extract the liability clause", contract_1_text)
        """
        if self._llm_fn is None:
            return "[ERROR: LLM function not configured]"
        
        full_prompt = f"{prompt}\n\nContext:\n{context}" if context else prompt
        
        try:
            result = self._llm_fn(full_prompt)
            logger.debug("sub_llm_called", prompt_length=len(prompt), result_length=len(result))
            return result
        except Exception as e:
            logger.error("sub_llm_error", error=str(e))
            return f"[ERROR: {str(e)}]"
    
    def _batch_sub_llm(
        self,
        prompts: list[str],
        contexts: list[str],
    ) -> list[str]:
        """
        Batch process multiple sub-tasks in parallel.
        
        Implements Section 2.3 Optimization via Batching:
        "Instead of making 1,000 individual calls to summarize 1,000 paragraphs,
        the RLM writes code to group paragraphs into chunks of 5 and processes 
        them in parallel threads (or batched API calls)."
        
        Args:
            prompts: List of prompts for each sub-task.
            contexts: List of context strings for each sub-task.
            
        Returns:
            List of results from sub-LLM calls.
        """
        if len(prompts) != len(contexts):
            raise ValueError("prompts and contexts must have same length")
        
        if not prompts:
            return []
        
        logger.info(
            "batch_sub_llm_start",
            num_tasks=len(prompts),
            batch_size=self.batch_size,
        )
        
        # Group into batches for optimal token density
        batches = self._create_optimal_batches(prompts, contexts)
        
        # Process batches
        results = []
        for batch in batches:
            batch_results = self._process_batch(batch)
            results.extend(batch_results)
        
        logger.info("batch_sub_llm_complete", num_results=len(results))
        return results
    
    def _create_optimal_batches(
        self,
        prompts: list[str],
        contexts: list[str],
    ) -> list[list[tuple[str, str]]]:
        """
        Create batches targeting optimal token density (~200k chars).
        """
        batches = []
        current_batch = []
        current_chars = 0
        
        for prompt, context in zip(prompts, contexts):
            item_chars = len(prompt) + len(context)
            
            if current_chars + item_chars > self.optimal_chars_per_call and current_batch:
                batches.append(current_batch)
                current_batch = []
                current_chars = 0
            
            current_batch.append((prompt, context))
            current_chars += item_chars
        
        if current_batch:
            batches.append(current_batch)
        
        return batches
    
    def _process_batch(self, batch: list[tuple[str, str]]) -> list[str]:
        """Process a batch of (prompt, context) pairs."""
        if self._llm_fn is None:
            return ["[ERROR: LLM not configured]"] * len(batch)
        
        results = []
        
        # Try parallel execution with ThreadPoolExecutor
        try:
            with ThreadPoolExecutor(max_workers=self.max_parallel_batches) as executor:
                futures = []
                for prompt, context in batch:
                    full_prompt = f"{prompt}\n\nContext:\n{context}" if context else prompt
                    futures.append(executor.submit(self._llm_fn, full_prompt))
                
                for future in futures:
                    try:
                        results.append(future.result(timeout=60))
                    except Exception as e:
                        results.append(f"[ERROR: {str(e)}]")
        except Exception as e:
            logger.error("batch_processing_error", error=str(e))
            # Fallback to sequential
            for prompt, context in batch:
                results.append(self._sub_llm(prompt, context))
        
        return results
    
    # =========================================================================
    # Synthesis
    # =========================================================================
    
    def _synthesize_answer(
        self,
        query: str,
        variable_names: list[str] | None = None,
    ) -> str:
        """
        Synthesize final answer from stored variables.
        
        This is the final step where all accumulated variables
        are resolved and compared to generate the answer.
        """
        if variable_names is None:
            variable_names = self._list_variables()
        
        # Resolve all variables
        context_parts = []
        for name in variable_names:
            content = self._get_variable(name)
            if content:
                context_parts.append(f"=== {name} ===\n{content}")
        
        if not context_parts:
            return "No relevant information found."
        
        context = "\n\n".join(context_parts)
        
        # Use LLM to synthesize
        if self._llm_fn:
            prompt = f"""Based on the following collected information, answer the query.

Query: {query}

Collected Information:
{context}

Provide a comprehensive answer based on the above information."""
            
            return self._llm_fn(prompt)
        
        return f"Found {len(variable_names)} relevant sections. Variables: {variable_names}"
    
    # =========================================================================
    # Configuration
    # =========================================================================
    
    def set_llm_function(self, llm_fn: Callable[[str], str]) -> None:
        """Set the LLM function for sub-calls."""
        self._llm_fn = llm_fn
        logger.info("llm_function_configured")
    
    def get_system_prompt(self) -> str:
        """Get the RLM system prompt for the LLM."""
        return RLM_SYSTEM_PROMPT
    
    def get_state_summary(self) -> dict[str, Any]:
        """Get current REPL state summary."""
        return {
            "doc_length": len(self.document_text),
            "num_nodes": len(self.skeleton),
            "variables": self._list_variables(),
            "execution_count": len(self.execution_history),
        }


# =============================================================================
# Factory Function
# =============================================================================


def create_repl_environment(
    document_text: str | None = None,
    skeleton: dict[str, SkeletonNode] | None = None,
    kv_store: KVStore | None = None,
    llm_provider: str | None = None,
) -> REPLEnvironment:
    """
    Create a configured REPL environment.
    
    Args:
        document_text: Full document text (DOC_VAR).
        skeleton: Skeleton index for tree navigation.
        kv_store: KV store for full node content.
        llm_provider: LLM provider for sub-calls ('openai', 'anthropic', 'gemini').
        
    Returns:
        Configured REPLEnvironment.
        
    Example:
        from rnsr import ingest_document, build_skeleton_index
        from rnsr.agent.repl_env import create_repl_environment
        
        result = ingest_document("contract.pdf")
        skeleton, kv_store = build_skeleton_index(result.tree)
        
        # Get full text
        full_text = kv_store.get("root") or ""
        
        env = create_repl_environment(
            document_text=full_text,
            skeleton=skeleton,
            kv_store=kv_store,
            llm_provider="gemini",
        )
        
        # Execute code
        env.execute("len(DOC_VAR)")  # -> 152347
        env.execute("list_children('root')")  # -> [...]
    """
    from rnsr.indexing.kv_store import InMemoryKVStore
    
    env = REPLEnvironment(
        document_text=document_text or "",
        skeleton=skeleton or {},
        kv_store=kv_store or InMemoryKVStore(),
    )
    
    # Configure LLM if provider specified
    if llm_provider:
        try:
            from rnsr.llm import get_llm, LLMProvider
            
            # Convert string to LLMProvider enum if needed
            provider_enum = LLMProvider(llm_provider) if isinstance(llm_provider, str) else llm_provider
            llm = get_llm(provider=provider_enum)
            
            def llm_fn(prompt: str) -> str:
                response = llm.complete(prompt)
                return str(response)
            
            env.set_llm_function(llm_fn)
        except Exception as e:
            logger.warning("llm_config_failed", error=str(e))
    
    return env


# =============================================================================
# Async Batch Processing
# =============================================================================


async def batch_process_async(
    env: REPLEnvironment,
    prompts: list[str],
    contexts: list[str],
) -> list[str]:
    """
    Async batch processing for maximum throughput.
    
    Uses asyncio for concurrent LLM calls when available.
    """
    if env._async_llm_fn is None:
        # Fallback to sync batch
        return env._batch_sub_llm(prompts, contexts)
    
    async_fn = env._async_llm_fn  # Capture for type narrowing
    
    async def process_one(prompt: str, context: str) -> str:
        full_prompt = f"{prompt}\n\nContext:\n{context}" if context else prompt
        try:
            return await async_fn(full_prompt)
        except Exception as e:
            return f"[ERROR: {str(e)}]"
    
    # Process all concurrently
    tasks = [
        process_one(p, c) for p, c in zip(prompts, contexts)
    ]
    
    return await asyncio.gather(*tasks)
