"""
LangChain Skills Agent 主体
"""

import os
from pathlib import Path
from typing import Optional, Iterator

from dotenv import load_dotenv
from langchain.agents import create_agent
from langchain.chat_models import init_chat_model
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.messages import AIMessage, AIMessageChunk
from langgraph.checkpoint.memory import InMemorySaver

from app.config import settings
from .skill_loader import SkillLoader
from .context import SkillAgentContext
from .tools import ALL_TOOLS
from .stream.emitter import StreamEventEmitter
from .stream.tracker import ToolCallTracker
from .stream.utils import is_success, DisplayLimits

load_dotenv(override=True)

DEFAULT_MODEL = "claude-sonnet-4-5-20250929"
DEFAULT_MAX_TOKENS = 16000
DEFAULT_TEMPERATURE = 1.0
DEFAULT_THINKING_BUDGET = 10000


def get_anthropic_credentials() -> tuple[str | None, str | None]:
    api_key = os.getenv("ANTHROPIC_API_KEY") or os.getenv("ANTHROPIC_AUTH_TOKEN")
    base_url = os.getenv("ANTHROPIC_BASE_URL")
    return api_key, base_url


class LangChainSkillsAgent:
    """基于 LangChain 1.0 的 Skills Agent"""

    def __init__(
        self,
        model: Optional[str] = None,
        skill_paths: Optional[list[Path]] = None,
        working_directory: Optional[Path] = None,
        max_tokens: Optional[int] = None,
        temperature: Optional[float] = None,
        enable_thinking: bool = True,
        thinking_budget: int = DEFAULT_THINKING_BUDGET,
    ):
        self.enable_thinking = enable_thinking
        self.thinking_budget = thinking_budget

        self.model_name = model or os.getenv("CLAUDE_MODEL", DEFAULT_MODEL)
        self.max_tokens = max_tokens or int(os.getenv("MAX_TOKENS", str(DEFAULT_MAX_TOKENS)))
        if enable_thinking:
            self.temperature = 1.0
        else:
            self.temperature = temperature or float(os.getenv("MODEL_TEMPERATURE", str(DEFAULT_TEMPERATURE)))
        self.working_directory = working_directory or Path.cwd()

        self.skill_loader = SkillLoader(skill_paths)
        self.context = SkillAgentContext(
            skill_loader=self.skill_loader,
            working_directory=self.working_directory,
        )
        self.checkpointer = InMemorySaver()
        self.system_prompt = self._build_system_prompt()
        self.agent = self._create_agent()

    def refresh_agent(self):
        """重新构建 LangChain Agent 以应用新的 System Prompt"""
        self.agent = self._create_agent()

    def _build_system_prompt(self) -> str:
        base_prompt = """You are a helpful coding assistant with access to specialized skills.

Your capabilities include:
- Loading and using specialized skills for specific tasks
- Loading reference documents from skills when needed (via load_reference)
- Executing bash commands and scripts
- Reading and writing files
- Following skill instructions to complete complex tasks

When a user request matches a skill's description:
1. Use `load_skill(skill_name)` to get detailed instructions
2. If the skill has references and the task requires detailed information, use `load_reference(skill_name, reference_name)` to load specific reference material
3. Follow the skill's instructions to complete the task

**Important**: Reference files in the references/ directory are NOT automatically loaded. Only load them when the skill instructions indicate they are needed for the current task."""
        return self.skill_loader.build_system_prompt(base_prompt)

    def _create_agent(self):
        # GLM Support
        if "glm" in self.model_name.lower():
            model = ChatZhipuAI(
                model=self.model_name,
                api_key=settings.zai_api_key,
                temperature=self.temperature,
            )
            
            agent = create_agent(
                model=model,
                tools=ALL_TOOLS,
                system_prompt=self.system_prompt,
                context_schema=SkillAgentContext,
                checkpointer=self.checkpointer,
            )
            return agent

        api_key, base_url = get_anthropic_credentials()
        init_kwargs = {
            "temperature": self.temperature,
            "max_tokens": self.max_tokens,
        }
        if api_key:
            init_kwargs["api_key"] = api_key
        if base_url:
            init_kwargs["base_url"] = base_url

        if self.enable_thinking:
            init_kwargs["thinking"] = {
                "type": "enabled",
                "budget_tokens": self.thinking_budget,
            }

        model = init_chat_model(self.model_name, **init_kwargs)

        agent = create_agent(
            model=model,
            tools=ALL_TOOLS,
            system_prompt=self.system_prompt,
            context_schema=SkillAgentContext,
            checkpointer=self.checkpointer,
        )
        return agent

    def stream_events(self, message: str, thread_id: str = "default") -> Iterator[dict]:
        config = {"configurable": {"thread_id": thread_id}}
        emitter = StreamEventEmitter()
        tracker = ToolCallTracker()
        full_response = ""

        try:
            for event in self.agent.stream(
                {"messages": [{"role": "user", "content": message}]},
                config=config,
                context=self.context,
                stream_mode="messages",
            ):
                if isinstance(event, tuple) and len(event) >= 2:
                    chunk = event[0]
                else:
                    chunk = event

                if isinstance(chunk, (AIMessageChunk, AIMessage)):
                    for ev in self._process_chunk_content(chunk, emitter, tracker):
                        if ev.type == "text":
                            full_response += ev.data.get("content", "")
                        yield ev.data

                    if hasattr(chunk, "tool_calls") and chunk.tool_calls:
                        for ev in self._process_tool_calls(chunk.tool_calls, emitter, tracker):
                            yield ev.data

                elif hasattr(chunk, "type") and chunk.type == "tool":
                    for ev in self._process_tool_result(chunk, emitter, tracker):
                        yield ev.data

        except Exception as e:
            yield emitter.error(str(e)).data
            raise

        yield emitter.done(full_response).data

    def _process_chunk_content(self, chunk, emitter, tracker):
        content = chunk.content
        if isinstance(content, str):
            if content:
                yield emitter.text(content)
                return

        blocks = None
        if hasattr(chunk, "content_blocks"):
            blocks = chunk.content_blocks
        
        if blocks is None:
            if isinstance(content, dict):
                blocks = [content]
            elif isinstance(content, list):
                blocks = content
            else:
                return

        for block in blocks:
            if not isinstance(block, dict):
                if hasattr(block, "model_dump"):
                    block = block.model_dump()
                elif hasattr(block, "dict"):
                    block = block.dict()
                else:
                    continue

            block_type = block.get("type")

            if block_type in ("thinking", "reasoning"):
                thinking = block.get("thinking") or block.get("reasoning") or ""
                if thinking:
                    yield emitter.thinking(thinking)

            elif block_type == "text":
                text = block.get("text") or block.get("content") or ""
                if text:
                    yield emitter.text(text)

            elif block_type in ("tool_use", "tool_call"):
                tool_id = block.get("id", "")
                name = block.get("name", "")
                args = block.get("input") if block_type == "tool_use" else block.get("args")
                args_payload = args if isinstance(args, dict) else {}

                if tool_id:
                    tracker.update(tool_id, name=name, args=args_payload)
                    if tracker.is_ready(tool_id):
                        tracker.mark_emitted(tool_id)
                        yield emitter.tool_call(name, args_payload, tool_id)

            elif block_type == "input_json_delta":
                partial_json = block.get("partial_json", "")
                if partial_json:
                    tracker.append_json_delta(partial_json)

            elif block_type == "tool_call_chunk":
                tool_id = block.get("id", "")
                name = block.get("name", "")
                if tool_id:
                    tracker.update(tool_id, name=name)
                partial_args = block.get("args", "")
                if isinstance(partial_args, str) and partial_args:
                    tracker.append_json_delta(partial_args)

    def _process_tool_calls(self, tool_calls, emitter, tracker):
        for tc in tool_calls:
            tool_id = tc.get("id", "")
            if tool_id:
                name = tc.get("name", "")
                args = tc.get("args", {})
                args_payload = args if isinstance(args, dict) else {}
                tracker.update(tool_id, name=name, args=args_payload)
                if tracker.is_ready(tool_id):
                    tracker.mark_emitted(tool_id)
                    yield emitter.tool_call(name, args_payload, tool_id)

    def _process_tool_result(self, chunk, emitter, tracker):
        tracker.finalize_all()
        for info in tracker.get_all():
            yield emitter.tool_call(info.name, info.args, info.id)

        name = getattr(chunk, "name", "unknown")
        raw_content = str(getattr(chunk, "content", ""))
        content = raw_content[:DisplayLimits.TOOL_RESULT_MAX]
        if len(raw_content) > DisplayLimits.TOOL_RESULT_MAX:
            content += "\n... (truncated)"
        
        success = is_success(content)
        yield emitter.tool_result(name, content, success)
