"""AI-assisted editor for an EXISTING world.

The player chats with the AI; each turn the AI returns:
  - a short player-facing message describing what it changed / will change,
  - an updated `definition` (the full new WorldDefinition).

The caller (API endpoint) decides whether to persist the new definition
to the World row. The editor itself is stateless aside from the in-memory
dialogue cache (keyed by world_id), so the player can iterate.
"""
from __future__ import annotations

import json
import uuid
from typing import Any, Dict, List, Optional, Tuple

from sqlalchemy.ext.asyncio import AsyncSession

from app.core.llm import LlmClient
from app.core.settings_service import cast_setting, get_all_settings
from app.logging_setup import get_logger
from app.models import User, World
from app.prompts.templates import get_prompt
from app.schemas import WorldDefinition
from app.engine.tools.tools import WORLD_BUILDER_TOOL_SCHEMAS

log = get_logger("world_editor")


# In-memory dialogue cache: world_id -> list of messages.
_DIALOGUES: Dict[uuid.UUID, Dict[str, Any]] = {}


async def edit_world_via_chat(
    db: AsyncSession,
    user: User,
    world: World,
    message: str,
) -> Tuple[str, Optional[Dict[str, Any]], bool]:
    """Run one turn of AI-assisted world editing.

    Returns (ai_message, new_definition_dict_or_None, changed).
      - ai_message: short prose reply for the player (in world.language).
      - new_definition_dict: the full updated definition if the AI proposed
        changes this turn, else None.
      - changed: True if new_definition_dict is not None and differs from
        the current world.definition.
    """
    settings_map = await get_all_settings(db)
    llm = LlmClient(settings_map)

    # Get / init dialogue state for this world.
    dialogue = _DIALOGUES.get(world.id)
    if not dialogue:
        system_prompt = _build_system_prompt(world)
        dialogue = {
            "user_id": user.id,
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": _build_seed_message(world)},
            ],
        }
        _DIALOGUES[world.id] = dialogue

    # Authorization: only the owner (or admin) may continue an existing dialogue.
    if dialogue["user_id"] != user.id and not user.is_admin:
        raise ValueError("forbidden")

    dialogue["messages"].append({"role": "user", "content": message})

    response = await llm.chat(
        messages=dialogue["messages"],
        tools=WORLD_BUILDER_TOOL_SCHEMAS,
        temperature=float(cast_setting("llm.temperature", settings_map.get("llm.temperature", 0.7))),
        purpose="world_editor",
        user_id=user.id,
        db=db,
    )

    ai_message, new_defn, _is_final = _extract_world_definition(
        response.text, response.tool_calls,
    )

    # If the model did not return a tool call but did emit JSON in text,
    # _extract_world_definition handles it. If still None, just return the
    # conversational message without changes.
    if new_defn is None:
        dialogue["messages"].append({
            "role": "assistant",
            "content": response.text or "",
            "tool_calls": response.tool_calls or None,
        })
        return ai_message, None, False

    new_defn_dict = new_defn.model_dump()
    changed = new_defn_dict != (world.definition or {})

    dialogue["messages"].append({
        "role": "assistant",
        "content": response.text or "",
        "tool_calls": response.tool_calls or None,
    })
    return ai_message, new_defn_dict, changed


def reset_editor_dialogue(world_id: uuid.UUID) -> None:
    """Drop the cached editor dialogue for a world (e.g. after manual save)."""
    _DIALOGUES.pop(world_id, None)


def _build_system_prompt(world: World) -> str:
    """System prompt for the world editor.

    Reuses the world-builder prompt but overrides the workflow: instead of
    designing from scratch, the AI is told to MODIFY the existing definition.
    """
    base = get_prompt("world_builder", world.language)
    override = (
        "\n\nADDITIONAL CONTEXT — YOU ARE EDITING AN EXISTING WORLD:\n"
        "The world already exists with the definition provided in the first "
        "user message. The player will give you edit instructions in their "
        "language ({world_language}). For EACH instruction:\n"
        "1. Call `submit_world_definition` with the FULL updated definition "
        "(not just the changed fields — the entire object, all required keys).\n"
        "2. Your text response should briefly summarize what you changed in "
        "the player's language. 2-4 sentences max.\n"
        "3. NEVER set `is_final=true` — the player will commit changes "
        "manually via the Save button.\n"
        "4. Preserve `initial_state` consistency with `world_schema`. If you "
        "change the schema, update the state accordingly.\n"
        "5. Preserve `initial_time` and `calendar` unless the player asks to "
        "change them.\n"
    ).format(world_language=world.language or "en")
    return base + override


def _build_seed_message(world: World) -> str:
    """First user message: dumps the current world definition as context."""
    defn = world.definition or {}
    parts = [
        "=== CURRENT WORLD DEFINITION ===",
        f"Name: {world.name}",
        f"Language: {world.language}",
        f"Current time: {world.current_time or '(none)'}",
        f"Definition JSON:\n```json\n{json.dumps(defn, ensure_ascii=False, indent=2)}\n```",
        f"Live state JSON:\n```json\n{json.dumps(world.state or {}, ensure_ascii=False, indent=2)[:2000]}\n```",
        "",
        "The player will now give you edit instructions. Apply each one by "
        "calling submit_world_definition with the FULL updated definition.",
    ]
    return "\n".join(parts)


# === Output extraction (mirrors world_builder._extract_world_definition) ===
def _extract_world_definition(
    text: str,
    tool_calls: Optional[List[Dict[str, Any]]],
) -> Tuple[str, Optional[WorldDefinition], bool]:
    import re as _re

    proposed: Optional[WorldDefinition] = None
    is_final = False
    ai_text = _sanitize_model_text(text or "")

    if tool_calls:
        for tc in tool_calls:
            if tc.get("function", {}).get("name") == "submit_world_definition":
                args_str = tc.get("function", {}).get("arguments", "{}")
                data = _safe_json_loads(args_str, {})
                if data:
                    proposed = _try_build_definition(data)
                    is_final = bool(data.get("is_final", False))
                    break

    if proposed is None:
        json_str = _extract_json_block(ai_text)
        if json_str:
            data = _safe_json_loads(json_str, None)
            if isinstance(data, dict):
                target = data
                if "proposed_definition" in data and isinstance(data["proposed_definition"], dict):
                    target = data["proposed_definition"]
                    if "is_final" in data:
                        is_final = bool(data["is_final"])
                    if "ai_message" in data and isinstance(data["ai_message"], str):
                        ai_text = data["ai_message"]
                else:
                    if "is_final" in data:
                        is_final = bool(data["is_final"])
                proposed = _try_build_definition(target)

    if proposed is not None:
        ai_text = _strip_json_blocks(ai_text).strip()
        if not ai_text:
            ai_text = "(definition updated)"

    return ai_text, proposed, is_final


def _sanitize_model_text(text: str) -> str:
    if not text:
        return ""
    import re as _re
    cleaned = _re.sub(r"<\s*/?\s*[a-zA-Z]+\s*>", "", text)
    cleaned = _re.sub(r"\n{3,}", "\n\n", cleaned)
    return cleaned.strip()


def _strip_json_blocks(text: str) -> str:
    if not text:
        return ""
    import re as _re
    cleaned = _re.sub(r"```[a-zA-Z]*\s*[\s\S]*?```", "", text)
    m = _re.search(r"\n\{[\s\S]*\}\s*$", cleaned)
    if m:
        cleaned = cleaned[: m.start()] + cleaned[m.end():]
    return cleaned


def _safe_json_loads(s: str, default: Any) -> Any:
    if not s:
        return default
    try:
        return json.loads(s)
    except json.JSONDecodeError:
        pass
    import re as _re
    repaired = _re.sub(r",\s*([}\]])", r"\1", s)
    try:
        return json.loads(repaired)
    except json.JSONDecodeError:
        pass
    try:
        repaired2 = repaired.replace("'", '"')
        return json.loads(repaired2)
    except json.JSONDecodeError:
        return default


def _try_build_definition(data: Dict[str, Any]) -> Optional[WorldDefinition]:
    try:
        return WorldDefinition.model_validate(data)
    except Exception as e:
        log.warning("world_editor_definition_invalid", error=str(e), keys=list(data.keys()))
        return None


def _extract_json_block(text: str) -> Optional[str]:
    if not text:
        return None
    import re as _re
    m = _re.search(r"```(?:json)?\s*(\{[\s\S]*?\})\s*```", text)
    if m:
        return m.group(1)
    start = text.find("{")
    if start == -1:
        return None
    depth = 0
    in_str = False
    esc = False
    for i in range(start, len(text)):
        c = text[i]
        if in_str:
            if esc:
                esc = False
            elif c == "\\":
                esc = True
            elif c == '"':
                in_str = False
        else:
            if c == '"':
                in_str = True
            elif c == "{":
                depth += 1
            elif c == "}":
                depth -= 1
                if depth == 0:
                    return text[start:i + 1]
    return None