"""Core RLM (Recursive Language Model) loop."""

from __future__ import annotations

import re
import sys
from dataclasses import dataclass, field
from typing import Any

from .sandbox import Sandbox, SandboxError
from .modelrelay import chat_completion, ModelRelayError
from .prompts import SYSTEM_PROMPT, USER_PROMPT_TEMPLATE


@dataclass
class ExecutionContext:
    """Context for tracking recursive LLM calls within sandbox execution.

    This enables the true RLM pattern where the LLM can spawn sub-queries
    to process chunks of data recursively.
    """
    depth: int = 0              # 0 = root query
    max_depth: int = 1          # Limit recursion depth (paper uses 1)
    max_calls: int = 10         # Limit sub-calls per execution
    calls_made: int = 0         # Track sub-calls made
    total_tokens: int = 0       # Track token usage (for cost monitoring)
    verbose: bool = False       # Whether to print debug output


@dataclass
class RLMResult:
    """Result of an RLM query.

    Attributes:
        answer: The final answer text (may contain error message if error occurred)
        iterations: Number of iterations (always 1 for true RLM)
        execution_log: List of code/output pairs from execution
        sub_calls_made: Total recursive sub-calls made via llm_query()
        retrieved_guids: GUIDs retrieved during execution (for benchmarks)
        error: Error information if execution failed, None on success
        total_tokens: Total tokens used (root call + all sub-calls)
    """
    answer: str
    iterations: int
    execution_log: list[dict[str, str]]  # List of {"code": ..., "output": ...}
    sub_calls_made: int = 0
    retrieved_guids: list[str] = field(default_factory=list)
    error: dict[str, str] | None = None  # {"type": ..., "message": ..., "code": ...}
    total_tokens: int = 0

    @property
    def success(self) -> bool:
        """Returns True if execution completed without errors."""
        return self.error is None


def extract_code_block(response: str, block_type: str) -> str | None:
    """Extract a fenced code block from the response."""
    pattern = rf"```{block_type}\n(.*?)```"
    match = re.search(pattern, response, re.DOTALL)
    if match:
        return match.group(1).strip()
    return None


def run_rlm(
    question: str,
    db_path: str,
    conversation_context: str | None = None,
    max_depth: int = 1,
    max_calls: int = 10,
    verbose: bool = False,
) -> RLMResult:
    """Run true RLM - single code generation with recursive sub-calls.

    This implements the RLM paper's approach: the LLM generates code ONCE,
    and recursion happens INSIDE the sandbox via llm_query() sub-calls.
    There is no outer iteration loop.

    Args:
        question: The user's question
        db_path: Path to the shadow.db SQLite database
        conversation_context: Optional prior turn context for follow-up questions
        max_depth: Maximum recursion depth for llm_query() sub-calls
        max_calls: Maximum number of llm_query() sub-calls
        verbose: Whether to print debug output

    Returns:
        RLMResult with the answer, execution details, and error info if failed
    """
    context = ExecutionContext(
        depth=0,
        max_depth=max_depth,
        max_calls=max_calls,
        verbose=verbose,
    )

    sandbox = Sandbox(db_path, context=context)

    try:
        user_content = USER_PROMPT_TEMPLATE.format(question=question)
        if conversation_context:
            user_content = f"{user_content}\n\nConversation Context:\n{conversation_context}"
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_content},
        ]

        if verbose:
            print(f"\n{'='*60}", file=sys.stderr)
            print("Generating code...", file=sys.stderr)
            print(f"{'='*60}", file=sys.stderr)

        # Single LLM call to generate code
        try:
            response, root_usage = chat_completion(messages, return_usage=True)
            context.total_tokens += root_usage.total_tokens
        except ModelRelayError as e:
            return RLMResult(
                answer="",
                iterations=0,
                execution_log=[],
                sub_calls_made=0,
                retrieved_guids=[],
                error={"type": "ModelRelayError", "message": str(e)},
                total_tokens=context.total_tokens,
            )

        if verbose:
            print(f"\nLLM Response:\n{response}", file=sys.stderr)

        # Extract code block
        code = extract_code_block(response, "python")
        if not code:
            # No code block - return response as answer
            if verbose:
                print("\nNo code block found, returning response as answer", file=sys.stderr)
            return RLMResult(
                answer=response,
                iterations=1,
                execution_log=[],
                sub_calls_made=0,
                retrieved_guids=[],
                total_tokens=context.total_tokens,
            )

        # Execute the code - recursion happens here via llm_query()
        if verbose:
            print(f"\nExecuting code:\n{code}", file=sys.stderr)

        execution_log: list[dict[str, str]] = []

        try:
            output = sandbox.execute(code)
            error = None
            execution_log.append({"code": code, "output": output})
        except SandboxError as e:
            # Fail-fast: don't silently convert to output string
            # Return structured error so caller can handle appropriately
            output = ""
            error = {"type": "SandboxError", "message": str(e), "code": code}
            execution_log.append({"code": code, "output": f"ERROR: {e}"})
            if verbose:
                print(f"\nExecution error: {e}", file=sys.stderr)

            # One-shot repair pass for forbidden imports
            if "Imports are not allowed" in str(e) or "ImportError" in str(e) or "__import__" in str(e):
                if verbose:
                    print("\nRetrying with no-import constraint...", file=sys.stderr)

                repair_messages = [
                    {"role": "system", "content": SYSTEM_PROMPT + "\n\nIMPORTANT: Do NOT import modules. Avoid datetime/timedelta usage. Use integer timestamp math (nanoseconds) and SQL MAX(timestamp) if you need cutoffs."},
                    {"role": "user", "content": user_content + "\n\nThe previous code failed due to imports. Regenerate code with no import statements and no datetime module usage."},
                ]
                try:
                    repair_response, repair_usage = chat_completion(repair_messages, return_usage=True)
                    context.total_tokens += repair_usage.total_tokens
                except ModelRelayError as repair_error:
                    return RLMResult(
                        answer="",
                        iterations=1,
                        execution_log=execution_log,
                        sub_calls_made=context.calls_made,
                        retrieved_guids=sandbox.get_retrieved_guids(),
                        error={"type": "ModelRelayError", "message": str(repair_error)},
                        total_tokens=context.total_tokens,
                    )

                if verbose:
                    print(f"\nLLM Repair Response:\n{repair_response}", file=sys.stderr)

                repaired_code = extract_code_block(repair_response, "python")
                if repaired_code:
                    if verbose:
                        print(f"\nExecuting repaired code:\n{repaired_code}", file=sys.stderr)
                    try:
                        output = sandbox.execute(repaired_code)
                        error = None
                        execution_log.append({"code": repaired_code, "output": output})
                    except SandboxError as repair_exec_error:
                        output = ""
                        error = {"type": "SandboxError", "message": str(repair_exec_error), "code": repaired_code}
                        execution_log.append({"code": repaired_code, "output": f"ERROR: {repair_exec_error}"})
                        if verbose:
                            print(f"\nRepair execution error: {repair_exec_error}", file=sys.stderr)

        if verbose:
            if output:
                print(f"\nOutput:\n{output}", file=sys.stderr)
            print(f"\nSub-calls made: {context.calls_made}", file=sys.stderr)

        return RLMResult(
            answer=output if output.strip() else ("Error: " + error["message"] if error else "No output produced."),
            iterations=1,
            execution_log=execution_log,
            sub_calls_made=context.calls_made,
            retrieved_guids=sandbox.get_retrieved_guids(),
            error=error,
            total_tokens=context.total_tokens,
        )

    finally:
        sandbox.close()
