#!/usr/bin/env python3
"""
strategies.py — Generation strategies for The Chaos engine.

Contains the three-phase conversational approach and the single-call
tool-based approach. All functions are standalone — no dependency on
GameEngine (config values and callbacks are passed explicitly).
"""

from __future__ import annotations

import json
import random
import re
from datetime import datetime
from typing import Iterator

from .models import GenerationResult, TurnResult
from .prompts import PROSE_PROMPT
from .llm import set_llm_env, call_llm
from .tools_handler import (
    execute_tool, describe_tool_action, describe_change,
    parse_changes_block, extract_tool_calls,
)
from .context import build_system_prompt, build_user_message, build_prose_prompt
from .validation import auto_prompt, validate_narrative
from .parsing import parse_response, log_turn_details
from . import state


# ── Synchronous (legacy) ───────────────────────────────────────────────────

def generate(
    player_action: str | None = None,
    last_narrative: str | None = None,
    *,
    model: str,
    temperature: float,
    timeout: int,
    max_tokens: int,
    api_key: str | None = None,
    api_base: str | None = None,
) -> GenerationResult:
    """
    Synchronous generation. Calls the LLM, parses the response,
    and returns a GenerationResult.
    """
    system = build_system_prompt()
    user = build_user_message(
        player_action=player_action, last_prompt=last_narrative
    )

    messages = [
        {"role": "system", "content": system},
        {"role": "user", "content": user},
    ]

    try:
        import litellm
    except ImportError:
        return GenerationResult(
            narrative="",
            error="litellm is not installed. Run: pip install litellm",
        )

    set_llm_env(model, api_key, api_base)

    try:
        response = litellm.completion(
            model=model,
            messages=messages,
            temperature=temperature,
            stream=False,
            timeout=timeout,
        )
        text = response.choices[0].message.content or ""
    except Exception as e:
        return GenerationResult(
            narrative="",
            error=f"LLM call failed: {e}",
        )

    return parse_response(text)


# ── Streaming (legacy) ─────────────────────────────────────────────────────

def generate_stream(
    player_action: str | None = None,
    last_narrative: str | None = None,
    *,
    model: str,
    temperature: float,
    timeout: int,
    max_tokens: int,
    api_key: str | None = None,
    api_base: str | None = None,
) -> Iterator[str]:
    """
    Streaming generator. Yields text chunks as they arrive from the LLM.
    On completion, the final yield is the FULL text (for parsing).
    """
    system = build_system_prompt()
    user = build_user_message(
        player_action=player_action, last_prompt=last_narrative
    )

    messages = [
        {"role": "system", "content": system},
        {"role": "user", "content": user},
    ]

    try:
        import litellm
    except ImportError:
        yield json.dumps({
            "error": "litellm is not installed. Run: pip install litellm"
        })
        return

    set_llm_env(model, api_key, api_base)

    try:
        response = litellm.completion(
            model=model,
            messages=messages,
            temperature=temperature,
            stream=True,
            timeout=timeout,
        )
        full_text = ""
        for chunk in response:
            delta = chunk.choices[0].delta.content or ""
            if delta:
                full_text += delta
                yield full_text
        yield full_text
    except Exception as e:
        yield json.dumps({"error": f"LLM call failed: {e}"})


# ── Three-phase generation ─────────────────────────────────────────────────

def generate_with_tools(
    player_action: str | None = None,
    last_prompt: str | None = None,
    on_thought: callable = None,
    on_action: callable = None,
    on_player_roll: callable = None,
    on_debug: callable = None,
    *,
    model: str,
    temperature: float,
    timeout: int,
    max_tokens: int,
    api_key: str | None = None,
    api_base: str | None = None,
) -> TurnResult:
    """
    Three-phase generation:

    1. **Prose** — LLM writes the full book_log from context + player action.
    2. **Summarize** — LLM condenses the book_log into one log line.
    3. **Extract** — LLM reads the book_log and outputs tool calls for state changes.
    """
    set_llm_env(model, api_key, api_base)
    datetime_now = datetime.now()
    state.append_llm_log(f"\n{'='*60}")
    state.append_llm_log(f"=== Turn — {datetime_now.strftime('%Y-%m-%d %H:%M:%S')} ===")
    state.append_llm_log(f"{'='*60}")
    if player_action:
        state.append_llm_log(f"Player: {player_action}")
    elif last_prompt:
        state.append_llm_log(f"Resume from: {last_prompt[:120]}")

    die_roll = random.randint(1, 6)
    state.append_llm_log(f"Dice: {die_roll} (1d6)")

    book_log = None
    changes_block = ""
    log_entry = None
    user_prompt = auto_prompt("")
    ambience = None
    debug_info = ""
    changes: list[str] = []

    for outer_attempt in range(3):
        # ── Phase 1: Prose ────────────────────────────────────────────────
        if on_action:
            on_action(f"Phase 1/3: writing story (dice={die_roll})")
        if on_debug:
            on_debug("phase", {"phase": 1, "name": "prose", "status": "start", "dice": die_roll, "outer_attempt": outer_attempt + 1})

        system = build_prose_prompt()

        user = build_user_message(
            player_action=player_action,
            last_prompt=last_prompt,
        )
        user += f"\n\n*A die is cast: **{die_roll}** (1d6).*"

        text = call_llm([
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ], model=model, temperature=temperature, timeout=timeout,
           max_tokens=1024, label=f"Prose attempt {outer_attempt + 1}", on_debug=on_debug)

        if not text or not text.strip():
            if on_debug:
                on_debug("phase", {"phase": 1, "status": "empty", "attempt": outer_attempt + 1})
            continue

        raw = text.strip()
        changes_block = ""
        if "### Changes" in raw:
            parts = raw.split("### Changes", 1)
            book_log = parts[0].strip()
            changes_block = "### Changes" + parts[1]
        else:
            book_log = raw
        if on_debug:
            preview = book_log[:150].replace("\n", "\\n")
            on_debug("phase", {"phase": 1, "status": "done", "chars": len(book_log), "changes": bool(changes_block), "preview": preview})

        # ── Validation ────────────────────────────────────────────────────
        if on_debug:
            on_debug("phase", {"phase": 1, "name": "validation", "status": "start"})
        valid, reason = validate_narrative(book_log, model=model, temperature=temperature, timeout=timeout, on_debug=on_debug)
        if not valid:
            if on_debug:
                on_debug("phase", {"phase": 1, "status": "validation_failed", "reason": reason, "outer_attempt": outer_attempt + 1})
            book_log = None
            continue

        # ── Phase 2: Summarize ────────────────────────────────────────────
        if on_action:
            on_action("Phase 2/3: summarizing story")
        if on_debug:
            on_debug("phase", {"phase": 2, "name": "summarize", "status": "start"})

        log_context = state.read_recent_log()
        log_entry = None
        for p2_attempt in range(2):
            context = book_log
            if changes_block:
                context += f"\n\n{changes_block}"
            text = call_llm([
                {"role": "user", "content":
                    f"Given the session log so far, summarize the new story in one line. "
                    f"Focus on who was involved (character and NPC names):\n\n"
                    f"## Session Log\n{log_context}\n\n"
                    f"## New Story\n{context}"}
            ], model=model, temperature=temperature, timeout=timeout,
               max_tokens=max_tokens, label=f"Summarize attempt {p2_attempt + 1}", on_debug=on_debug)
            if text and text.strip():
                log_entry = text.strip().split("\n")[0][:300]
                if on_debug:
                    on_debug("phase", {"phase": 2, "status": "done", "summary": log_entry})
                break

        if not log_entry:
            log_entry = book_log.split("\n")[0][:120]
            if on_debug:
                on_debug("phase", {"phase": 2, "status": "fallback", "summary": log_entry})

        # ── Phase 3: Extract state changes ────────────────────────────────
        if on_action:
            on_action("Phase 3/3: extracting state changes")
        if on_debug:
            on_debug("phase", {"phase": 3, "name": "extract", "status": "start"})

        user_prompt = auto_prompt(book_log)
        ambience = None
        phase3_errors = []
        changes = []

        # Step 1: Parse ### Changes block directly
        if changes_block.strip():
            for tc in parse_changes_block(changes_block):
                name = tc["tool"]
                args = tc.get("args", {})
                if name == "finalize_turn":
                    continue
                result = execute_tool(name, args)
                if result.startswith("**Error:") or result.startswith("Tool error") or result.startswith("Unknown"):
                    phase3_errors.append(f"{name}: {result}")
                else:
                    desc = describe_change(name, args)
                    if desc:
                        changes.append(desc)

        # Step 2: LLM Phase 3 for finalize_turn + any extra changes
        previous_attempt = None
        phase3_ok = False
        for p3_attempt in range(5):
            from paths import CHAR_PATH, WORLD_PATH
            current_char = state.read_file(CHAR_PATH) or "*No character.*"
            current_world = state.truncate_world(state.read_file(WORLD_PATH) or "") or "*No world.*"

            phase3_prompt = (
                f"## Current Character\n{current_char}\n\n"
                f"## Current World\n{current_world}\n\n"
                f"## Story\n{book_log}\n\n"
            )
            if changes_block.strip():
                phase3_prompt += (
                    f"## Changes already applied\n{changes_block}\n\n"
                    f"Output the finalize_turn tool to end the turn. "
                    f"Add extra tool calls if you spot changes the list above missed.\n\n"
                )
            else:
                phase3_prompt += (
                    f"Read the story and compare with current state. Output tool calls for any changes:\n\n"
                )
            phase3_prompt += (
                f"Output ```tool blocks for changes only. Examples:\n\n"
            )

            if previous_attempt:
                phase3_prompt += (
                    f"--- PREVIOUS ATTEMPT (had errors) ---\n"
                    f"{previous_attempt['output']}\n\n"
                    f"--- FEEDBACK ---\n"
                    f"{previous_attempt['feedback']}\n\n"
                    f"Fix the issues above. Output corrected tool calls only.\n\n"
                )

            text = call_llm([
                {"role": "user", "content": phase3_prompt +
                    f"```tool\n{{\"tool\": \"modify_vitals\", \"args\": {{\"current_hp\": 5, \"cash\": 45}}}}\n```\n"
                    f"```tool\n{{\"tool\": \"modify_traits\", \"args\": {{\"dex\": 15}}}}\n```\n"
                    f"```tool\n{{\"tool\": \"add_to_inventory\", \"args\": {{\"item\": \"Silver key\"}}}}\n```\n"
                    f"```tool\n{{\"tool\": \"remove_from_inventory\", \"args\": {{\"item\": \"Torches (10)\"}}}}\n```\n"
                    f"```tool\n{{\"tool\": \"replace_gear\", \"args\": {{\"before\": \"Mace (1d6+1)\", \"after\": \"Mace (1d6+2, sharpened)\"}}}}\n```\n"
                    f"```tool\n{{\"tool\": \"add_note\", \"args\": {{\"note\": \"Found a hidden passage under the temple\"}}}}\n```\n"
                    f"```tool\n{{\"tool\": \"replace_note\", \"args\": {{\"before\": \"Old note text\", \"after\": \"New note text\"}}}}\n```\n"
                    f"```tool\n{{\"tool\": \"world_update\", \"args\": {{\"content\": \"# The World\\n\\n...full new world state...\"}}}}\n```\n"
                    f"```tool\n{{\"tool\": \"journal_update\", \"args\": {{\"add\": [\"Investigate the mine\"], \"done\": [\"Defeat the demon\"]}}}}\n```\n"
                    f"```tool\n{{\"tool\": \"finalize_turn\", \"args\": {{\"user_prompt\": \"What do you do?\", \"ambience\": \"dungeon\"}}}}\n```\n\n"
                    f"Only output tools for things that actually changed. Omit unchanged fields."}
            ], model=model, temperature=temperature, timeout=timeout,
               max_tokens=max_tokens, label=f"Extract attempt {p3_attempt + 1}", on_debug=on_debug)

            if not text or not text.strip():
                if on_debug:
                    on_debug("phase", {"phase": 3, "status": "empty", "attempt": p3_attempt + 1})
                continue

            tool_calls_list = extract_tool_calls(
                text, round_num=p3_attempt + 1, on_debug=on_debug
            )
            if on_debug and tool_calls_list:
                names = [tc.get("tool", "?") for tc in tool_calls_list if tc.get("tool") != "finalize_turn"]
                fin = any(tc.get("tool") == "finalize_turn" for tc in tool_calls_list)
                on_debug("phase", {"phase": 3, "status": "tools_found", "tools": names, "has_finalize": fin})

            errors: list[str] = []
            attempt_changes: list[str] = []
            for tc in tool_calls_list:
                name = tc.get("tool", "?")
                args = tc.get("args", {})
                if name == "finalize_turn":
                    if args.get("user_prompt"):
                        user_prompt = args["user_prompt"]
                    if args.get("ambience"):
                        ambience = args["ambience"]
                    continue
                if on_action:
                    on_action(f"State: {describe_tool_action(name, args)}")
                if on_debug:
                    on_debug("tool_call", {"round": p3_attempt + 1, "tool": name, "args": args})

                if name == "player_roll" and on_player_roll:
                    dice = args.get("dice", "1d6")
                    reason = args.get("reason", "a check")
                    roll_val = on_player_roll(dice, reason)
                    result = f"Player rolled {dice} for '{reason}': {roll_val}"
                else:
                    result = execute_tool(name, args)

                if result.startswith("**Error:") or result.startswith("Tool error") or result.startswith("Unknown"):
                    errors.append(f"{name}: {result}")
                else:
                    desc = describe_change(name, args)
                    if desc:
                        attempt_changes.append(desc)
                if on_debug:
                    on_debug("tool_result", {"round": p3_attempt + 1, "tool": name, "result": result})

            if not errors:
                phase3_ok = True
                debug_info = ""
                changes.extend(attempt_changes)
                if on_debug:
                    on_debug("phase", {"phase": 3, "status": "done", "applied": len([tc for tc in tool_calls_list if tc.get("tool") != "finalize_turn"])})
                break

            phase3_errors = errors
            debug_info = "; ".join(errors)
            if on_debug:
                on_debug("phase", {"phase": 3, "status": "errors", "errors": errors, "attempt": p3_attempt + 1})

            feedback_lines = ["The previous tool calls had errors:"]
            for e in errors:
                feedback_lines.append(f"- {e}")
            feedback_lines.append("")
            feedback_lines.append("Fix ALL issues above. Use correct tool names, valid JSON, and reasonable values.")
            previous_attempt = {"output": text, "feedback": "\n".join(feedback_lines)}

        if phase3_ok:
            break

        if on_debug:
            on_debug("phase", {"phase": 3, "status": "exhausted", "errors": phase3_errors})
            on_debug("phase", {"phase": 1, "status": "retry_after_phase3_failure", "outer_attempt": outer_attempt + 1})
        book_log = None

    if not book_log:
        return TurnResult(error="Generation failed after exhausting all retries")

    if on_action:
        on_action("Turn complete")
    if on_debug:
        on_debug("phase_done", {
            "book_log_chars": len(book_log),
            "log_entry": log_entry,
            "user_prompt": user_prompt,
            "ambience": ambience,
            "extract_errors": debug_info or None,
        })

    state.append_llm_log(
        f"\n--- FINAL ---\n"
        f"book_log: {book_log[:200]}\n"
        f"log_entry: {log_entry}\n"
        f"user_prompt: {user_prompt}\n"
        f"ambience: {ambience}\n"
    )
    return TurnResult(
        book_log=book_log,
        log_entry=log_entry,
        user_prompt=user_prompt,
        ambience=ambience,
        debug_info=debug_info,
        changes=changes,
    )


# ── Single-call generation ─────────────────────────────────────────────────

def generate_with_tools_single(
    player_action: str | None = None,
    last_prompt: str | None = None,
    on_thought: callable = None,
    on_action: callable = None,
    on_player_roll: callable = None,
    on_debug: callable = None,
    *,
    model: str,
    temperature: float,
    timeout: int,
    max_tokens: int,
    api_key: str | None = None,
    api_base: str | None = None,
) -> TurnResult:
    """
    Single-call generation using tools.

    Uses a single LLM call with all tools available — LLM outputs
    narrative + tool blocks in one go. No retry loop.
    """
    datetime_now = datetime.now()
    state.append_llm_log(f"\n{'='*60}")
    state.append_llm_log(f"=== Turn — {datetime_now.strftime('%Y-%m-%d %H:%M:%S')} ===")
    state.append_llm_log(f"{'='*60}")
    if player_action:
        state.append_llm_log(f"Player: {player_action}")
    elif last_prompt:
        state.append_llm_log(f"Resume from: {last_prompt[:120]}")

    strategy_name = "tools"
    if on_action:
        on_action(f"LLM: {model} | temp={temperature} | tokens={max_tokens} | strategy={strategy_name}")
    if on_debug:
        on_debug("config", {"model": model, "temperature": temperature, "max_tokens": max_tokens, "strategy": strategy_name})

    die_roll = random.randint(1, 6)
    state.append_llm_log(f"Dice: {die_roll} (1d6)")

    system = """You are an RPG dungeon master. The player just took an action.

Output ONLY ```tool blocks — no prose, no reasoning, no explanation outside tool blocks. Every piece of output must be in a tool block.

Use these tools to perform every action. Wrap each in its own ```tool block:
```tool
{"tool": "narrative", "args": {"text": "The full vivid narrative prose goes here."}}
```
```tool
{"tool": "modify_vitals", "args": {"current_hp": 5, "cash": 45}}
```
```tool
{"tool": "modify_traits", "args": {"dex": 15}}
```
```tool
{"tool": "add_to_inventory", "args": {"item": "Silver key"}}
```
```tool
{"tool": "remove_from_inventory", "args": {"item": "Torches (10)"}}
```
```tool
{"tool": "replace_gear", "args": {"before": "Mace (1d6+1)", "after": "Mace (1d6+2, sharpened)"}}
```
```tool
{"tool": "add_note", "args": {"note": "Found a hidden passage under the temple"}}
```
```tool
{"tool": "replace_note", "args": {"before": "Old note text", "after": "New note text"}}
```
```tool
{"tool": "world_update", "args": {"content": "# The World\n\n...full new world state..."}}
```
```tool
{"tool": "journal_update", "args": {"add": ["Investigate the mine"], "done": ["Defeat the demon"]}}
```
```tool
{"tool": "finalize_turn", "args": {"user_prompt": "What do you do?", "ambience": "dungeon"}}
```
"""
    system += build_prose_prompt()

    user = build_user_message(
        player_action=player_action,
        last_prompt=last_prompt,
    )
    user += f"\n\n*A die is cast: **{die_roll}** (1d6).*"

    start_time = datetime.now()
    set_llm_env(model, api_key, api_base)
    state.append_llm_log(f"\n[TOOL] Single call — {len(system)} chars system, {len(user)} chars user")
    state.append_llm_log(f"System preview: {system.split(chr(10))[0][:80]}...")
    state.append_llm_log(f"User preview: {user.split(chr(10))[0][:80]}...")

    text = call_llm(
        [{"role": "system", "content": system},
         {"role": "user", "content": user}],
        model=model, temperature=temperature, timeout=timeout,
        max_tokens=4096, label="Single tool call", on_debug=on_debug,
    )

    total_elapsed = (datetime.now() - start_time).total_seconds() * 1000
    if text:
        state.append_llm_log(f"\n[TOOL] got {len(text)} chars in {total_elapsed:.1f}ms")

    if not text or not text.strip():
        return TurnResult(error="Single tool call returned empty response")

    raw = text.strip()
    book_log = ""
    log_entry = None
    user_prompt = auto_prompt("")
    ambience = None
    tool_calls = []
    changes: list[str] = []
    phase3_errors: list[str] = []

    tool_pattern = r"```tool\s*\n?(.*?)\n?```"
    matches = re.findall(tool_pattern, text, re.DOTALL)
    if matches:
        for block in matches:
            block = block.strip()
            try:
                tc = json.loads(block)
                tool_calls.append(tc)
                name = tc.get("tool", "unknown")
                args = tc.get("args", {})
                state.append_llm_log(f"\n[EXTRACT] {name}: {json.dumps(args)[:100]}")

                if name == "narrative":
                    book_log = args.get("text", book_log)
                elif name == "finalize_turn":
                    if args.get("user_prompt"):
                        user_prompt = args["user_prompt"]
                    if args.get("ambience"):
                        ambience = args["ambience"]
            except json.JSONDecodeError as e:
                state.append_llm_log(f"\n[EXTRACT] bad JSON: {e}")
                continue

        log_entry = None
        if book_log:
            clean = re.sub(r'\s+', ' ', book_log).strip()
            first_sentence = re.split(r'(?<=[.!?])\s+', clean)
            if first_sentence:
                log_entry = first_sentence[0].strip()[:200]
            else:
                log_entry = clean[:200]
            state.append_llm_log(f"\n[SUMMARY] \"{log_entry}\"")

        extr_start = datetime.now()
        for tc in tool_calls:
            name = tc.get("tool", "unknown")
            args = tc.get("args", {})
            if name in ("finalize_turn", "narrative"):
                continue
            result = execute_tool(name, args)
            if result.startswith("**Error:") or result.startswith("Tool error") or result.startswith("Unknown"):
                phase3_errors.append(f"{name}: {result}")
            else:
                desc = describe_change(name, args)
                if desc:
                    changes.append(desc)

        apply_elapsed = (datetime.now() - extr_start).total_seconds() * 1000
        state.append_llm_log(f"\n[APPLY] {len(changes)} changes in {apply_elapsed:.1f}ms")
    else:
        state.append_llm_log(f"\n[TOOL] no tool blocks found")

    elapsed = (datetime.now() - start_time).total_seconds() * 1000

    if on_action:
        on_action("Turn complete")
    if on_debug:
        applied = len([tc for tc in tool_calls if tc.get("tool") != "finalize_turn"])
        on_debug("phase_done", {
            "book_log_chars": len(book_log),
            "log_entry": log_entry,
            "user_prompt": user_prompt,
            "ambience": ambience,
            "extract_errors": phase3_errors or None,
            "total_elapsed_ms": elapsed,
            "tool_calls_count": len(tool_calls),
            "applied_changes_count": applied,
            "tool_call_results": tool_calls,
        })

    log_turn_details(
        player_action=player_action or last_prompt or "",
        last_prompt=last_prompt or "",
        strategy_name=strategy_name,
        die_roll=die_roll,
        model=model,
        temperature=temperature,
        max_tokens=max_tokens,
        book_log=book_log,
        log_entry=log_entry or "",
        ambience=ambience,
        tool_calls=tool_calls,
        on_debug=on_debug,
    )
    return TurnResult(
        book_log=book_log,
        log_entry=log_entry,
        user_prompt=user_prompt,
        ambience=ambience,
        debug_info="; ".join(phase3_errors) if phase3_errors else "",
        changes=changes,
    )