# retoor <retoor@molodetz.nl>

from __future__ import annotations

import json
import logging
from typing import Any, Optional

from ..config import Settings
from ..errors import ToolInputError
from .lessons import LessonStore
from .loop import TraceCallback, react_loop
from .state import (
    MAX_EVAL_DEPTH,
    AgentState,
    get_eval_depth,
    get_state,
    reset_eval_depth,
    set_eval_depth,
)

logger = logging.getLogger("devii.agentic.controller")

SUB_AGENT_SYSTEM_PROMPT = (
    "You are a focused sub-agent spawned to complete a single scoped task on this "
    "DevPlace platform. Begin with a plan() call. Investigate, then act with the most specific "
    "tools available. If you change anything, confirm it and call verify() before returning. "
    "Return a concise factual result string under 1500 characters."
)
NO_DELEGATE = "delegate"


class AgenticController:
    def __init__(self, lessons: LessonStore, settings: Settings) -> None:
        self._lessons = lessons
        self._settings = settings
        self._llm: Any = None
        self._dispatcher: Any = None
        self._tools: list[dict[str, Any]] = []
        self._on_trace: Optional[TraceCallback] = None
        self._cost_tracker: Any = None
        self._chunk_store: Any = None

    def bind(
        self,
        llm: Any,
        dispatcher: Any,
        tools: list[dict[str, Any]],
        on_trace: Optional[TraceCallback],
        cost_tracker: Any = None,
        chunk_store: Any = None,
    ) -> None:
        self._llm = llm
        self._dispatcher = dispatcher
        self._tools = tools
        self._on_trace = on_trace
        self._cost_tracker = cost_tracker
        self._chunk_store = chunk_store

    async def dispatch(self, name: str, arguments: dict[str, Any]) -> str:
        handlers = {
            "plan": self._plan,
            "reflect": self._reflect,
            "recall": self._recall,
            "forget_lessons": self._forget,
            "lesson_rate": self._lesson_rate,
            "lesson_count": self._lesson_count,
            "verify": self._verify,
            "delegate": self._delegate,
            "eval": self._eval,
        }
        handler = handlers.get(name)
        if handler is None:
            raise ToolInputError(f"Unknown agentic tool: {name}")
        return await handler(arguments)

    async def _plan(self, arguments: dict[str, Any]) -> str:
        goal = str(arguments.get("goal", "")).strip()
        steps = arguments.get("steps") or []
        if not goal:
            raise ToolInputError("plan requires a goal.")
        if not isinstance(steps, list) or not steps:
            raise ToolInputError("plan requires a non-empty steps list.")
        confidence = float(arguments.get("confidence", 0.8) or 0.8)
        state = get_state()
        if state is not None:
            state.plan = {
                "goal": goal,
                "steps": steps,
                "success_criteria": arguments.get("success_criteria", ""),
                "confidence": confidence,
            }
        advice = ""
        if confidence < 0.6:
            advice = "Confidence is below 0.6 - gather more context or recall() past lessons before executing."
        return json.dumps(
            {
                "status": "success",
                "plan_recorded": True,
                "step_count": len(steps),
                "advice": advice,
            },
            ensure_ascii=False,
        )

    async def _reflect(self, arguments: dict[str, Any]) -> str:
        observation = str(arguments.get("observation", "")).strip()
        conclusion = str(arguments.get("conclusion", "")).strip()
        next_action = str(arguments.get("next_action", "")).strip()
        if not (observation and conclusion and next_action):
            raise ToolInputError(
                "reflect requires observation, conclusion, and next_action."
            )
        tags = str(arguments.get("tags", "") or "").strip()
        record = self._lessons.add(observation, conclusion, next_action, tags)
        state = get_state()
        if state is not None:
            state.reflections.append(
                {
                    "observation": observation,
                    "conclusion": conclusion,
                    "next_action": next_action,
                }
            )
        return json.dumps(
            {
                "status": "success",
                "lesson_uid": record["uid"],
                "total_lessons": self._lessons.count(),
            },
            ensure_ascii=False,
        )

    async def _recall(self, arguments: dict[str, Any]) -> str:
        query = str(arguments.get("query", "")).strip()
        if not query:
            raise ToolInputError("recall requires a query.")
        k = int(
            arguments.get("k", self._settings.recall_top_k)
            or self._settings.recall_top_k
        )
        hits = self._lessons.search(query, k=k)
        return json.dumps(
            {"status": "success", "count": len(hits), "lessons": hits},
            ensure_ascii=False,
        )

    async def _forget(self, arguments: dict[str, Any]) -> str:
        query = str(arguments.get("query", "") or "").strip()
        if query:
            removed = 0
            for hit in self._lessons.search(query, k=50):
                uid = hit.get("uid")
                if uid and self._lessons.delete(uid):
                    removed += 1
        else:
            removed = self._lessons.clear()
        return json.dumps(
            {
                "status": "success",
                "forgotten": removed,
                "remaining": self._lessons.count(),
            },
            ensure_ascii=False,
        )

    async def _lesson_rate(self, arguments: dict[str, Any]) -> str:
        uid = str(arguments.get("uid", "")).strip()
        if not uid:
            raise ToolInputError("lesson_rate requires a lesson uid.")
        raw_value = arguments.get("value")
        if raw_value is None:
            raise ToolInputError("lesson_rate requires a value (1 or -1).")
        try:
            value = int(raw_value)
        except (ValueError, TypeError):
            raise ToolInputError("lesson_rate value must be an integer (1 or -1).") from None
        if value not in (1, -1):
            raise ToolInputError("lesson_rate value must be 1 (useful) or -1 (unhelpful).")
        ok = self._lessons.rate(uid, value)
        if not ok:
            raise ToolInputError(f"No lesson found with uid '{uid}'.")
        return json.dumps({"status": "success", "lesson_uid": uid, "rated": value}, ensure_ascii=False)

    async def _lesson_count(self, arguments: dict[str, Any]) -> str:
        count = self._lessons.count()
        return json.dumps(
            {"status": "success", "active_lessons": count},
            ensure_ascii=False,
        )

    async def _verify(self, arguments: dict[str, Any]) -> str:
        summary = str(arguments.get("summary", "")).strip()
        if not summary:
            raise ToolInputError("verify requires a summary of what was confirmed.")
        confirmed = arguments.get("confirmed", True)
        if isinstance(confirmed, str):
            confirmed = confirmed.strip().lower() not in ("", "false", "no", "0")
        state = get_state()
        if state is not None and confirmed:
            state.verified = True
        return json.dumps(
            {"status": "success", "verified": bool(confirmed), "summary": summary},
            ensure_ascii=False,
        )

    async def _spawn(
        self,
        prompt: str,
        tools: list[dict[str, Any]],
        system_prompt: str = SUB_AGENT_SYSTEM_PROMPT,
    ) -> tuple[str, AgentState]:
        if self._llm is None or self._dispatcher is None:
            raise ToolInputError(
                "Sub-agent execution is not available in this context."
            )
        depth = get_eval_depth()
        if depth >= MAX_EVAL_DEPTH:
            raise ToolInputError(
                "Nested self-evaluation limit reached; a tool or eval cannot keep calling itself."
            )
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": prompt},
        ]
        sub_state = AgentState()
        token = set_eval_depth(depth + 1)
        try:
            result = await react_loop(
                llm=self._llm,
                dispatcher=self._dispatcher,
                messages=messages,
                tools=tools,
                state=sub_state,
                settings=self._settings,
                max_iterations=self._settings.delegate_max_iterations,
                plan_required=self._settings.plan_required,
                verify_required=self._settings.verify_required,
                on_trace=self._on_trace,
                cost_tracker=self._cost_tracker,
                chunk_store=self._chunk_store,
            )
        finally:
            reset_eval_depth(token)
        return result, sub_state

    async def run_subagent(self, prompt: str) -> str:
        prompt = str(prompt or "").strip()
        if not prompt:
            raise ToolInputError("A non-empty prompt is required.")
        tools = [
            tool for tool in self._tools if tool["function"]["name"] != NO_DELEGATE
        ]
        result, _ = await self._spawn(prompt, tools)
        return result

    async def _delegate(self, arguments: dict[str, Any]) -> str:
        task = str(arguments.get("task", "")).strip()
        if not task:
            raise ToolInputError("delegate requires a task description.")
        allowed = arguments.get("allowed_tools")
        if allowed:
            allowed_set = {str(name) for name in allowed}
            tools = [
                tool
                for tool in self._tools
                if tool["function"]["name"] in allowed_set
                and tool["function"]["name"] != NO_DELEGATE
            ]
        else:
            tools = [
                tool for tool in self._tools if tool["function"]["name"] != NO_DELEGATE
            ]
        result, sub_state = await self._spawn(task, tools)
        return json.dumps(
            {
                "status": "success",
                "iterations": sub_state.iteration,
                "verified": sub_state.verified,
                "reflections": len(sub_state.reflections),
                "result": result,
            },
            ensure_ascii=False,
        )

    async def _eval(self, arguments: dict[str, Any]) -> str:
        prompt = str(arguments.get("prompt", "")).strip()
        if not prompt:
            raise ToolInputError("eval requires a non-empty prompt.")
        result = await self.run_subagent(prompt)
        return json.dumps({"status": "success", "result": result}, ensure_ascii=False)
