#!/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 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, validate_action
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,
) -> 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",
        )

    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,
) -> 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

    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,
) -> 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.
    """
    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)")

    # ── Pre-generation validation ────────────────────────────────────
    if player_action:
        valid, reason = validate_action(
            player_action,
            model=model,
            timeout=timeout,
            on_debug=on_debug,
        )
        if not valid:
            state.append_llm_log(f"\n[VALIDATION REJECTED] {reason}")
            fail_narrative = f"You can't do that — {reason}."
            return TurnResult(
                book_log=fail_narrative,
                log_entry=fail_narrative,
                user_prompt=auto_prompt(""),
            )

    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,
) -> 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.

You are the sole authority over the game state. The player's action is a **proposal**, not a fact. If their action contradicts the character sheet (e.g. using an item they don't have, spending cash they don't have, claiming stats they don't have), narrate the failure with the narrative tool and do NOT call any state-changing tools.

**Inventory rule**: If the player wants to use an item, verify it's on the character sheet first. If it is, you MUST call `remove_from_inventory` for that item AND apply effects (e.g. `modify_vitals`). If it's not on the sheet, narrate the failure — do not let them use items they don't have.

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).*"

    # ── Pre-generation validation ────────────────────────────────────
    if player_action:
        valid, reason = validate_action(
            player_action,
            model=model,
            timeout=timeout,
            on_debug=on_debug,
        )
        if not valid:
            state.append_llm_log(f"\n[VALIDATION REJECTED] {reason}")
            fail_narrative = f"You can't do that — {reason}."
            return TurnResult(
                book_log=fail_narrative,
                log_entry=fail_narrative,
                user_prompt=auto_prompt(""),
            )

    start_time = datetime.now()
    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,
    )