"""
Chat workflow tools extracted from ChatWorkflowService.
"""
from __future__ import annotations

import asyncio
import time
import uuid
from typing import Any, Dict, List

try:
    from backend.models.agent import TurnContext
    from backend.models.internal import CharacterCard, TokenUsageStatus
    from backend.models.regex_rules import RegexPlacement
    from backend.services.character_service import CharacterService
    from backend.services.regex_service import regex_service
    from backend.services.task_queue_manager import TaskType, task_queue_manager
    from backend.services.token_usage_service import token_usage_service
    from backend.core.config import settings
except ImportError:
    from models.agent import TurnContext
    from models.internal import CharacterCard, TokenUsageStatus
    from models.regex_rules import RegexPlacement
    from services.character_service import CharacterService
    from services.regex_service import regex_service
    from services.task_queue_manager import TaskType, task_queue_manager
    from services.token_usage_service import token_usage_service
    from core.config import settings


_character_service = CharacterService()
_workflow_service = None


def _get_workflow_service():
    """Lazy init to avoid circular import with chat_workflow_service."""
    global _workflow_service
    if _workflow_service is None:
        try:
            from backend.services.chat_workflow_service import ChatWorkflowService
        except ImportError:
            from services.chat_workflow_service import ChatWorkflowService
        _workflow_service = ChatWorkflowService()
    return _workflow_service


async def regex_apply_user_input(ctx: TurnContext) -> None:
    processed = regex_service.apply_rules_by_placement(
        text=ctx.user_message,
        placement=RegexPlacement.USER_INPUT.value,
        character_name=ctx.current_role,
        preset_name=ctx.preset_name,
        message_depth=0,
        is_for_llm=True,
        is_markdown_rendered=False,
    )
    if processed != ctx.user_message:
        print("[WorkflowTool] Applied user-input regex rules")
    ctx.user_message = processed


async def load_character(ctx: TurnContext) -> None:
    character_data = ctx.request_data.get("characterData")
    if not character_data:
        character = _character_service.get_character_by_name(ctx.current_role)
        if not character:
            raise ValueError(f"角色 '{ctx.current_role}' 不存在")
    else:
        character = CharacterCard(**character_data)
    ctx.character = character
    print(f"[WorkflowTool] Loaded character: {character.name}")


async def activate_worldbook(ctx: TurnContext) -> None:
    svc = _get_workflow_service()
    active_entries = await svc._collect_and_activate_worldbooks(
        ctx.request_data,
        ctx.character,
    )
    ctx.active_entries = active_entries
    print(f"[WorkflowTool] Activated {len(active_entries)} worldbook entries")

    if ctx.callbacks and ctx.callbacks.on_worldbook_active:
        entries_payload = [
            entry.model_dump() if hasattr(entry, "model_dump") else entry
            for entry in active_entries
        ]
        await ctx.callbacks.on_worldbook_active(entries_payload)


async def load_chat_history(ctx: TurnContext) -> None:
    svc = _get_workflow_service()
    chat_history = await svc._load_chat_history(
        ctx.current_role,
        ctx.current_chat,
    )
    ctx.chat_history = chat_history
    print(f"[WorkflowTool] Loaded {len(chat_history)} history messages")


async def build_prompt_messages(ctx: TurnContext) -> None:
    svc = _get_workflow_service()
    prompt_messages = svc._assemble_prompt(
        ctx.character,
        ctx.chat_history,
        ctx.user_message,
        ctx.active_entries,
        ctx.request_data,
    )
    ctx.prompt_messages = prompt_messages
    print(f"[WorkflowTool] Built {len(prompt_messages)} prompt messages")


async def llm_main_reply(ctx: TurnContext) -> None:
    svc = _get_workflow_service()
    api_config = ctx.request_data.get("apiConfig", {})
    preset_config = ctx.request_data.get("presetConfig", {})

    if ctx.stream:
        if not api_config.get("api_key"):
            raise ValueError("API Key 未配置，请先在 API 配置页面保存密钥")

        generated_content = ""
        chunk_count = 0
        start_time = time.time()

        async for chunk_dict in svc.llm_client.stream_chat(
            messages=ctx.prompt_messages,
            api_url=api_config.get("api_url", ""),
            api_key=api_config.get("api_key", ""),
            model=api_config.get("model", ""),
            temperature=preset_config.get("parameters", {}).get("temperature", 1.0),
            max_tokens=preset_config.get("parameters", {}).get("max_tokens", 30000),
            request_timeout=preset_config.get("parameters", {}).get("request_timeout", 60),
        ):
            if isinstance(chunk_dict, dict):
                if chunk_dict.get("type") == "chunk":
                    chunk_content = chunk_dict.get("content", "")
                elif chunk_dict.get("type") == "usage":
                    continue
                else:
                    chunk_content = chunk_dict.get("content", str(chunk_dict))
            else:
                chunk_content = str(chunk_dict)

            generated_content += chunk_content
            chunk_count += 1

            if ctx.callbacks and ctx.callbacks.on_chunk:
                await ctx.callbacks.on_chunk(chunk_content)

        ctx.duration = time.time() - start_time
        ctx.generated_content = generated_content
        ctx.token_usage = {
            "prompt_tokens": len(str(ctx.prompt_messages)) // 4,
            "completion_tokens": len(generated_content) // 4,
            "total_tokens": (len(str(ctx.prompt_messages)) // 4)
            + (len(generated_content) // 4),
        }
        print(
            f"[WorkflowTool] Stream LLM complete: {chunk_count} chunks, "
            f"{len(generated_content)} chars"
        )
    else:
        result = await svc._generate_response(
            ctx.prompt_messages,
            api_config,
            preset_config,
            stream=False,
        )
        ctx.generated_content = result["content"]
        ctx.token_usage = result.get("usage", {})
        ctx.duration = result.get("duration", 0.0)
        print(f"[WorkflowTool] LLM complete: {len(ctx.generated_content)} chars")


async def regex_apply_ai_output(ctx: TurnContext) -> None:
    processed = regex_service.apply_rules_by_placement(
        text=ctx.generated_content,
        placement=RegexPlacement.AI_OUTPUT.value,
        character_name=ctx.current_role,
        preset_name=ctx.preset_name,
        message_depth=0,
        is_for_llm=False,
        is_markdown_rendered=False,
    )
    if processed != ctx.generated_content:
        print("[WorkflowTool] Applied AI-output regex rules")
    ctx.generated_content = processed


async def record_token_usage(ctx: TurnContext) -> None:
    chat_id = f"{ctx.current_role}/{ctx.current_chat}"
    floor = ctx.request_data.get("floor", 0)
    api_config = ctx.request_data.get("apiConfig", {})

    try:
        await token_usage_service.record_usage(
            chat_id=chat_id,
            role_name=ctx.current_role,
            chat_name=ctx.current_chat,
            prompt_tokens=ctx.token_usage.get("prompt_tokens", 0),
            completion_tokens=ctx.token_usage.get("completion_tokens", 0),
            total_tokens=ctx.token_usage.get("total_tokens", 0),
            status=TokenUsageStatus.COMPLETED,
            floor=floor + 1,
            duration=ctx.duration,
            model=api_config.get("model"),
            api_provider="openai",
            api_url=api_config.get("api_url"),
        )
    except Exception as exc:
        print(f"[WorkflowTool] Token usage recording failed: {exc}")


async def enqueue_parallel_tasks(ctx: TurnContext) -> None:
    chat_id = f"{ctx.current_role}/{ctx.current_chat}"
    options = ctx.request_data.get("options", {})
    image_task_id = None
    table_task_id = None

    if options.get("imageWorkflow", False):
        image_task_id = f"img_{uuid.uuid4().hex[:8]}"
        await task_queue_manager.add_task(image_task_id, TaskType.IMAGE_WORKFLOW, chat_id)

    if options.get("dynamicTable", False):
        table_task_id = f"tbl_{uuid.uuid4().hex[:8]}"
        await task_queue_manager.add_task(table_task_id, TaskType.DYNAMIC_TABLE, chat_id)

    ctx.task_ids = {
        "imageWorkflow": image_task_id,
        "dynamicTable": table_task_id,
    }

    if ctx.callbacks and ctx.callbacks.on_tasks_created:
        if image_task_id or table_task_id:
            await ctx.callbacks.on_tasks_created(ctx.task_ids)

    # Fire-and-forget parallel workers (same as legacy service)
    svc = _get_workflow_service()
    asyncio.create_task(
        svc._start_parallel_tasks(
            ctx.request_data,
            ctx.generated_content,
            image_task_id,
            table_task_id,
        )
    )


def register_chat_tools(registry) -> None:
    """Register all chat workflow tools on the given registry."""
    registry.register("regex_apply_user_input", regex_apply_user_input, description="Apply user-input regex")
    registry.register("load_character", load_character, description="Load character card")
    registry.register("activate_worldbook", activate_worldbook, description="Activate worldbook entries")
    registry.register("load_chat_history", load_chat_history, description="Load chat history")
    registry.register("build_prompt_messages", build_prompt_messages, description="Assemble LLM prompt")
    registry.register("llm_main_reply", llm_main_reply, description="Call main LLM (supports stream)")
    registry.register("regex_apply_ai_output", regex_apply_ai_output, description="Apply AI-output regex")
    registry.register("record_token_usage", record_token_usage, description="Persist token usage")
    registry.register("enqueue_parallel_tasks", enqueue_parallel_tasks, description="Enqueue parallel tasks")
