# src/agent.py
import os
from pathlib import Path

from google.adk.agents import LlmAgent
from google.adk.agents.callback_context import CallbackContext
from google.adk.models import LlmRequest
from google.adk.models.google_llm import Gemini
from google.adk.models.lite_llm import LiteLlm
from google.adk.runners import Runner
from google.adk.sessions import DatabaseSessionService, InMemorySessionService
from google.genai import types
from typing import Awaitable, Callable

from src.config import settings
from src.memory.core import get_block
from src.memory.recall import SqliteRecallMemoryService
from src.memory.archival import ArchivalStore
from src.tools.send import send_message
from src.tools.core_tools import core_memory_append, core_memory_replace
from src.tools.search_tools import recall_search, archival_insert, archival_search


APP_NAME = "memgpt_lite"
DEFAULT_USER_ID = "default"
_SYSTEM_PROMPT_PATH = Path(__file__).parent / "prompts" / "system.md"

# Storage mode: "disk" (default, dev + demo recording) uses on-disk SQLite
# + on-disk FAISS + DatabaseSessionService. "memory" (set on HF Space) uses
# InMemorySessionService + :memory: SQLite + non-persisted FAISS — container
# sleep naturally resets state, no cleanup code needed. See config.py for
# MEMGPT_DATA_DIR (disk mode's data root).
_STORAGE_MODE = os.environ.get("MEMGPT_STORAGE_MODE", "disk")
if _STORAGE_MODE not in ("disk", "memory"):
    raise RuntimeError(
        f"MEMGPT_STORAGE_MODE must be 'disk' or 'memory', got {_STORAGE_MODE!r}"
    )

# Demo-abuse cap: bounds cost per LLM call. Default 1024 out; override with
# MEMGPT_MAX_OUTPUT_TOKENS env var (2048 for demo recording; higher for
# self-hosters). Documented in README §Deployment.
_MAX_OUTPUT_TOKENS = int(os.environ.get("MEMGPT_MAX_OUTPUT_TOKENS", "1024"))

if _STORAGE_MODE == "memory":
    session_service = InMemorySessionService()
    memory_service = SqliteRecallMemoryService(db_path=":memory:")
    archival_store = ArchivalStore(
        db_path=":memory:",
        index_path=None,
        dim=settings.embeddings.dim,
    )
else:
    session_service = DatabaseSessionService(
        db_url=f"sqlite+aiosqlite:///{Path(settings.storage.sessions_db).as_posix()}"
    )
    memory_service = SqliteRecallMemoryService(db_path=settings.storage.recall_db)
    archival_store = ArchivalStore(
        db_path=settings.storage.archival_db,
        index_path=settings.storage.archival_index,
        dim=settings.embeddings.dim,
    )


def _build_model():
    if settings.model.provider == "gemini":
        retry_config = types.HttpRetryOptions(
            attempts=settings.model.retry.attempts,
            exp_base=settings.model.retry.exp_base,
            initial_delay=settings.model.retry.initial_delay,
            http_status_codes=settings.model.retry.http_status_codes,
        )
        return Gemini(model=settings.model.name, retry_options=retry_config)

    if settings.model.provider == "openrouter":
        return LiteLlm(
            model=settings.model.name,
            num_retries=settings.model.retry.attempts,
            base_url="https://openrouter.ai/api/v1",
        )

    raise ValueError(f"Unknown model provider: {settings.model.provider}")


def _inject_core_memory(
    callback_context: CallbackContext, llm_request: LlmRequest
) -> None:
    """Append <persona> and <human> blocks to system_instruction each turn."""
    state = callback_context._invocation_context.session.state
    persona = get_block(state, "persona")
    human = get_block(state, "human")

    base = llm_request.config.system_instruction or ""
    if persona:
        base += f"\n\n<persona>\n{persona}\n</persona>"
    if human:
        base += f"\n\n<human>\n{human}\n</human>"
    llm_request.config.system_instruction = base


async def _auto_save_to_memory(callback_context: CallbackContext) -> None:
    """Mirror every completed turn into recall memory."""
    await callback_context._invocation_context.memory_service.add_session_to_memory(
        callback_context._invocation_context.session
    )


def build_agent() -> LlmAgent:
    return LlmAgent(
        model=_build_model(),
        name="memgpt_lite",
        description="Chatbot with tiered persistent memory (MemGPT-lite).",
        instruction=_SYSTEM_PROMPT_PATH.read_text(encoding="utf-8"),
        tools=[
            send_message,
            core_memory_append,
            core_memory_replace,
            recall_search,
            archival_insert,
            archival_search,
        ],
        before_model_callback=_inject_core_memory,
        after_agent_callback=_auto_save_to_memory,
        # Demo-abuse cap: bounds cost per LLM call. Default 1024 out; override with
        # MEMGPT_MAX_OUTPUT_TOKENS env var (2048 for demo recording;
        # self-hosters can raise/remove). Documented in README §Deployment.
        generate_content_config=types.GenerateContentConfig(
            max_output_tokens=_MAX_OUTPUT_TOKENS,
        ),
    )


def build_runner() -> Runner:
    return Runner(
        agent=build_agent(),
        app_name=APP_NAME,
        session_service=session_service,
        memory_service=memory_service,
    )


async def _ensure_session(session_id: str, user_id: str):
    """Get an existing session or create a new one.

    Follows the notebook's try-create-except-get pattern: create_session
    raises if the id is taken; get_session returns the existing one.
    """
    try:
        return await session_service.create_session(
            app_name=APP_NAME, user_id=user_id, session_id=session_id
        )
    except Exception:
        return await session_service.get_session(
            app_name=APP_NAME, user_id=user_id, session_id=session_id
        )


ConfirmCallback = Callable[[str, dict], Awaitable[bool]]


def _extract_reply(events) -> str:
    """Return the last send_message emission across the given events (may be '')."""
    reply = ""
    for event in events:
        if not event.content or not event.content.parts:
            continue
        for part in event.content.parts:
            fc = getattr(part, "function_call", None)
            if fc is not None and fc.name == "send_message":
                reply = fc.args.get("message", "")
    return reply


async def start_turn(
    runner: Runner,
    user_text: str,
    session_id: str = "default",
    user_id: str = DEFAULT_USER_ID,
) -> dict:
    """Run phase 1 of a user turn. Returns a dict with reply + pending approval.

    Return shape:
        {
            "reply":   str,   # from send_message in phase-1 events (may be "")
            "pending": None | {
                "approval_id":   str,
                "invocation_id": str,
                "hint":          str,
                "payload":       dict,
                "session_id":    str,
                "user_id":       str,
            }
        }

    If ``pending`` is ``None``, the turn completed. Otherwise the invocation
    is paused; the caller must invoke ``resume_turn(runner, pending, approved)``
    to finish it.
    """
    session = await _ensure_session(session_id, user_id)
    query = types.Content(role="user", parts=[types.Part(text=user_text)])

    events = []
    async for event in runner.run_async(
        user_id=user_id, session_id=session.id, new_message=query
    ):
        events.append(event)

    reply = _extract_reply(events)
    approval = _check_for_approval(events)

    if approval is None:
        return {"reply": reply, "pending": None}

    return {
        "reply": reply,
        "pending": {
            **approval,
            "session_id": session.id,
            "user_id": user_id,
        },
    }


async def resume_turn(
    runner: Runner,
    pending: dict,
    approved: bool,
) -> str:
    """Resume a paused invocation with an approval decision. Returns reply text."""
    resume_msg = _build_approval_response(pending["approval_id"], approved)
    events = []
    async for event in runner.run_async(
        user_id=pending["user_id"],
        session_id=pending["session_id"],
        new_message=resume_msg,
        invocation_id=pending["invocation_id"],
    ):
        events.append(event)
    return _extract_reply(events)


async def run_turn(
    runner: Runner,
    user_text: str,
    session_id: str = "default",
    user_id: str = DEFAULT_USER_ID,
    on_confirm: ConfirmCallback | None = None,
) -> str:
    """Convenience wrapper: start_turn, auto-resume if paused. For REPL/tests.

    The Gradio UI in Step 18c uses start_turn + resume_turn directly instead,
    because it needs to hand control back to the user for the yes/no reply.
    """
    result = await start_turn(runner, user_text, session_id, user_id)
    if result["pending"] is None:
        return result["reply"]

    approved = (
        await on_confirm(result["pending"]["hint"], result["pending"]["payload"])
        if on_confirm is not None
        else False
    )
    resumed_reply = await resume_turn(runner, result["pending"], approved)
    return resumed_reply if resumed_reply else result["reply"]


def _check_for_approval(events) -> dict | None:
    """Scan collected events for a confirmation request. Returns None if not paused.

    ADK nests the confirmation data under a ``toolConfirmation`` object inside
    ``fc.args`` (alongside ``originalFunctionCall``); the hint and payload are
    NOT at the top level. Inspected via scratch/inspect_confirmation_args.py —
    banking the shape here so a future ADK bump has a diff to grep for.
    """
    for event in events:
        if not event.content or not event.content.parts:
            continue
        for part in event.content.parts:
            fc = getattr(part, "function_call", None)
            if fc is not None and fc.name == "adk_request_confirmation":
                confirmation = fc.args.get("toolConfirmation", {}) or {}
                return {
                    "approval_id": fc.id,
                    "invocation_id": event.invocation_id,
                    "hint": confirmation.get("hint", ""),
                    "payload": confirmation.get("payload", {}),
                }
    return


def _build_approval_response(approval_id: str, approved: bool) -> types.Content:
    """Build the FunctionResponse Content that resumes a paused invocation."""
    response = types.FunctionResponse(
        id=approval_id,
        name="adk_request_confirmation",
        response={"confirmed": approved},
    )
    return types.Content(role="user", parts=[types.Part(function_response=response)])