"""Core RLM (Recursive Language Model) loop wrapper using rlm-core."""

from __future__ import annotations

import os
from dataclasses import dataclass, field
from typing import Callable

from rlm_core import DataSourceRegistry, PolicyHints
from rlm_core import RLMConfig as CoreRLMConfig
from rlm_core import RLMResult as CoreRLMResult
from rlm_core import run_rlm as run_core_rlm
from rlm_core.loop.code_extractor import extract_code_block

from .environment import RecallEnvironment
from .execution_context import (
    DEFAULT_MAX_CALLS,
    DEFAULT_MAX_DEPTH,
    DEFAULT_MAX_ITERATIONS,
    DEFAULT_MAX_OUTPUT_CHARS,
    create_execution_context,
)
from .llm_client import LLMClient, create_default_llm_client
from .llm_orchestrator import LlmOrchestrator
from .message_database import MessageDatabase
from .modelrelay import ModelRelayError, resolve_root_model, resolve_subcall_model
from .prompts import build_system_prompt, USER_PROMPT_TEMPLATE, REFINEMENT_PROMPT_TEMPLATE
from .restricted_executor import RestrictedExecutor
from .sources import MessageDataSource
from .subcalls import RclSubcallClient


@dataclass
class RLMResult:
    """Result of an RLM query."""
    answer: str
    iterations: int
    execution_log: list[dict[str, str]]
    sub_calls_made: int = 0
    retrieved_guids: list[str] = field(default_factory=list)
    error: dict[str, str] | None = None
    total_tokens: int = 0

    @property
    def success(self) -> bool:
        return self.error is None


ProgressCallback = Callable[[str], None]


def _core_result_to_rlm_result(
    core: CoreRLMResult,
    *,
    sub_calls_made: int,
    retrieved_guids: list[str],
) -> RLMResult:
    execution_log = [
        {
            "code": record.code_executed,
            "output": record.execution_result,
            "iteration": record.iteration,
        }
        for record in core.trajectory
    ]
    error = None
    if core.error is not None:
        error = {
            "type": core.error.type,
            "message": core.error.message,
        }
        if core.error.code is not None:
            error["code"] = core.error.code
    return RLMResult(
        answer=core.answer,
        iterations=core.iterations,
        execution_log=execution_log,
        sub_calls_made=sub_calls_made,
        retrieved_guids=retrieved_guids,
        error=error,
        total_tokens=core.tokens_used,
    )



def run_rlm(
    question: str,
    db_path: str,
    contacts_db_path: str | None = None,
    conversation_context: str | None = None,
    max_depth: int = DEFAULT_MAX_DEPTH,
    max_calls: int = DEFAULT_MAX_CALLS,
    max_iterations: int = DEFAULT_MAX_ITERATIONS,
    max_output_chars: int = DEFAULT_MAX_OUTPUT_CHARS,
    verbose: bool = False,
    root_model: str | None = None,
    subcall_model: str | None = None,
    on_progress: ProgressCallback | None = None,
    llm_client: LLMClient | None = None,
    tips_profile: str | None = None,
    policy_hints: PolicyHints | None = None,
) -> RLMResult:
    """Run RLM with iterative refinement using rlm-core."""
    client = llm_client or create_default_llm_client()

    context = create_execution_context(
        max_depth=max_depth,
        max_calls=max_calls,
        max_iterations=max_iterations,
        max_output_chars=max_output_chars,
        verbose=verbose,
        on_progress=on_progress,
    )

    has_contacts = (
        isinstance(contacts_db_path, str)
        and os.path.exists(contacts_db_path)
        and os.path.getsize(contacts_db_path) > 0
    )

    resolved_tips_profile = tips_profile or os.environ.get("RECALL_RLM_TIPS_PROFILE", "full")
    system_prompt = build_system_prompt(has_contacts=has_contacts, tips_profile=resolved_tips_profile)
    resolved_root_model = resolve_root_model(root_model)
    resolved_subcall_model = resolve_subcall_model(subcall_model, root_fallback=root_model)

    db = MessageDatabase(db_path, contacts_db_path=contacts_db_path)
    data_source = MessageDataSource(db, has_contacts=has_contacts)
    registry = DataSourceRegistry([data_source], default_source_name="messages")

    orchestrator = LlmOrchestrator(context, client, subcall_model=resolved_subcall_model)
    subcalls = RclSubcallClient(orchestrator)

    executor = RestrictedExecutor(db, max_output_chars=max_output_chars)
    environment = RecallEnvironment(
        executor=executor,
        registry=registry,
        subcalls=subcalls,
        max_output_chars=max_output_chars,
    )

    core_config = CoreRLMConfig(
        max_iterations=max_iterations,
        max_depth=max_depth,
        max_calls=max_calls,
        max_output_chars=max_output_chars,
        tips_profile=resolved_tips_profile,
        verbose=verbose,
        root_model=resolved_root_model,
        on_progress=on_progress,
        policy_hints=policy_hints,
    )

    try:
        core_result = run_core_rlm(
            question=question,
            environment=environment,
            root_llm=client,
            subcalls=subcalls,
            config=core_config,
            system_prompt=system_prompt,
            conversation_context=conversation_context,
            user_prompt_template=USER_PROMPT_TEMPLATE,
            refinement_prompt_template=REFINEMENT_PROMPT_TEMPLATE,
            context=context,
        )
    except ModelRelayError as exc:
        return RLMResult(
            answer="",
            iterations=0,
            execution_log=[],
            sub_calls_made=context.stats.calls_made,
            retrieved_guids=data_source.get_retrieved_guids(),
            error={"type": "ModelRelayError", "message": str(exc)},
            total_tokens=context.stats.total_tokens,
        )
    finally:
        db.close()

    return _core_result_to_rlm_result(
        core_result,
        sub_calls_made=context.stats.calls_made,
        retrieved_guids=data_source.get_retrieved_guids(),
    )
