from pathlib import Path
from typing import AsyncGenerator, Callable, Literal, Unpack
from uuid import uuid4

from opentelemetry import trace
from opentelemetry.context import Context
from opentelemetry.trace import SpanKind, set_span_in_context

from ..llm import ILLMService, LLMResponse
from ..llm.errors import AceAIConfigurationError, AceAIRuntimeError
from ..llm.models import (
    LLMHostedToolSpec,
    LLMMessage,
    LLMRequestMeta,
    LLMToolCall,
    LLMToolCallDelta,
    LLMToolSpec,
)
from .models import AgentStep, ToolExecutionResult
from .skills import (
    SkillLoader,
    SkillRegistry,
    format_skills_for_prompt,
)
from aceai.llm.tracing import get_trace_ctx, set_trace_ctx
from .context_manager import ContextManager
from .events import (
    AgentEvent,
    AgentEventBuilder,
    LLMCompletedEvent,
    RunCompletedEvent,
    RunFailedEvent,
    StepFailedEvent,
    ToolCompletedEvent,
    ToolFailedEvent,
)
from .executor import RunState, IExecutor, ToolExecutor


class ToolExecutionFailure(AceAIRuntimeError):
    """Raised when a tool invocation fails during a reasoning step."""

    def __init__(self, *, tool_call: LLMToolCall, error: Exception):
        message = str(error)
        super().__init__(message)
        self.tool_call = tool_call
        self.original_error = error


class AgentBase:
    """Base class for agents using an LLM provider."""

    def __init__(
        self,
        prompt: str = "",
        *,
        default_model: str,
        llm_service: ILLMService,
        max_steps: int = 5,
        executor: IExecutor | None = None,
        tracer: trace.Tracer | None = None,
        skill_path: str | Path | Literal["auto", "disable"] = "auto",
        skill_loader_factory: Callable[[str], SkillLoader] = SkillLoader,
        hosted_tools: list[LLMHostedToolSpec] | None = None,
    ):
        if max_steps < 1:
            raise AceAIConfigurationError("max_steps must be at least 1")
        self._skill_registry = self._load_skill_registry(
            skill_path=skill_path,
            skill_loader_factory=skill_loader_factory,
        )
        skill_prompt = format_skills_for_prompt(self._skill_registry)
        self._default_model = default_model
        self._llm_service = llm_service
        self._prompt = prompt
        self._ctx_mgr: ContextManager = ContextManager(prompt + skill_prompt)
        self._executor = executor
        self._hosted_tools = hosted_tools if hosted_tools is not None else []
        if isinstance(executor, ToolExecutor) and self._skill_registry.get_skills():
            executor.register_tools(*self._skill_registry.as_tools())
        self._max_steps = max_steps
        self._tracer = tracer or trace.get_tracer("aceai.core")

    @property
    def default_model(self) -> str:
        return self._default_model

    @property
    def llm_service(self) -> ILLMService:
        return self._llm_service

    @property
    def max_steps(self) -> int:
        return self._max_steps

    @property
    def skill_registry(self) -> SkillRegistry:
        return self._skill_registry

    @property
    def system_message(self) -> LLMMessage:
        return self._ctx_mgr.system_message

    def add_instruction(self, instruction: str) -> None:
        """Add an instruction into the agent's context manager."""
        if instruction == "":
            raise ValueError("Empty Instruction")
        self._ctx_mgr.add_instruction(instruction)

    def _load_skill_registry(
        self,
        *,
        skill_path: str | Path | Literal["auto", "disable"],
        skill_loader_factory: Callable[[str], SkillLoader],
    ) -> SkillRegistry:
        registry = SkillRegistry()
        for path in self._resolve_skill_paths(skill_path):
            registry.register(*skill_loader_factory(str(path)).load_skills().get_skills())
        return registry

    def _resolve_skill_paths(
        self, skill_path: str | Path | Literal["auto", "disable"]
    ) -> list[Path]:
        global_skills = Path.home() / ".aceai" / "skills"
        if skill_path == "auto":
            return [global_skills, Path.cwd() / ".agent" / "skills"]
        if skill_path == "disable":
            return []
        return [global_skills, Path(skill_path).expanduser()]

    async def _call_llm(
        self,
        request_meta: LLMRequestMeta,
        messages: list[LLMMessage],
        event_builder: AgentEventBuilder,
    ):
        executor = self._executor

        tools: list[LLMToolSpec] = []
        if executor:
            tools.extend(executor.select_tools())
        tools.extend(self._hosted_tools)

        if tools:
            stream = self._llm_service.stream(
                messages=messages,
                tools=tools,
                metadata=request_meta,
            )
        else:
            stream = self._llm_service.stream(
                messages=messages,
                metadata=request_meta,
            )

        try:
            async for stream_event in stream:
                match stream_event.event_type:
                    case "response.output_text.delta":
                        chunk = stream_event.text_delta
                        if isinstance(chunk, str):
                            yield event_builder.llm_text_delta(text_delta=chunk)
                    case "response.media":
                        yield event_builder.llm_media(segments=stream_event.segments)
                    case "response.function_call_arguments.delta":
                        tool_call_delta = stream_event.tool_call_delta
                        if isinstance(tool_call_delta, LLMToolCallDelta):
                            yield event_builder.llm_tool_call_delta(
                                tool_call_delta=tool_call_delta,
                            )
                    case "response.error":
                        if isinstance(stream_event.error, str):
                            raise AceAIRuntimeError(stream_event.error)
                        raise AceAIRuntimeError("LLM streaming error")
                    case "response.completed":
                        response = stream_event.response
                        if not isinstance(response, LLMResponse):
                            raise AceAIRuntimeError(
                                "LLM stream completed without a response payload"
                            )
                        for segment in response.segments:
                            if segment.type == "reasoning":
                                yield event_builder.llm_reasoning(segment=segment)
                        yield AgentStep(
                            step_id=event_builder.step_id,
                            llm_response=response,
                        )
                        return
                    case _:
                        raise AceAIRuntimeError(
                            f"Unsupported LLM stream event: {stream_event.event_type}"
                        )
        finally:
            await stream.aclose()

    async def _make_toolcalls(
        self,
        current_step: AgentStep,
        event_builder: AgentEventBuilder,
        run_state: RunState,
    ):
        executor = self._executor

        for call in current_step.llm_response.tool_calls:
            yield event_builder.tool_started(tool_call=call)
            if executor is None:
                raise AceAIConfigurationError(
                    "executor must be provided when tool calls are enabled"
                )
            try:
                tool_output = await executor.execute_tool(
                    call,
                    run_state=run_state,
                )
            except Exception as exc:
                raise ToolExecutionFailure(tool_call=call, error=exc) from exc

            tool_result = ToolExecutionResult(call=call, output=tool_output)
            current_step.tool_results.append(tool_result)
            yield event_builder.tool_completed(
                tool_call=call,
                tool_result=tool_result,
            )

    async def _run_step(
        self,
        *,
        event_builder: AgentEventBuilder,
        run_state: RunState,
        **request_meta: Unpack[LLMRequestMeta],
    ) -> AsyncGenerator[AgentEvent, None]:
        run_context = get_trace_ctx()
        if run_context is None:
            raise AceAIRuntimeError("trace context is not set for agent step")
        step_span = self._tracer.start_span(
            "agent.step",
            kind=SpanKind.INTERNAL,
            context=run_context,
            attributes={
                "agent.step_id": event_builder.step_id,
                "agent.max_steps": self._max_steps,
            },
        )
        step_context = set_span_in_context(step_span, run_context)
        set_trace_ctx(step_context)

        try:
            yield event_builder.llm_started()
            llm_gen = self._call_llm(
                request_meta,
                messages=self._ctx_mgr.context,
                event_builder=event_builder,
            )
            try:
                async for event in llm_gen:
                    if not isinstance(event, AgentStep):
                        yield event
                        continue

                    current_step = event
                    yield event_builder.llm_completed(step=current_step)

                    if not current_step.llm_response.tool_calls:
                        yield event_builder.step_completed(step=current_step)
                        if final_answer := current_step.llm_response.text:
                            yield event_builder.run_completed(
                                step=current_step,
                                final_answer=final_answer,
                            )
                        return
                    self._ctx_mgr.add_tool_call(current_step.llm_response)
                    tool_gen = self._make_toolcalls(
                        current_step,
                        event_builder,
                        run_state,
                    )
                    try:
                        async for tool_event in tool_gen:
                            if isinstance(tool_event, ToolCompletedEvent):
                                self._ctx_mgr.add_tool_use(tool_event)
                            yield tool_event
                    finally:
                        await tool_gen.aclose()
            finally:
                await llm_gen.aclose()

        finally:
            set_trace_ctx(run_context)
            if step_span.is_recording():
                step_span.end()

    def _handle_failed_step(
        self,
        step: AgentStep,
        step_index: int,
        exc: Exception | None = None,
        error_msg: str | None = None,
    ):
        step_id = step.step_id
        error_msg = error_msg or str(exc)

        if isinstance(exc, ToolExecutionFailure):
            tool_call = exc.tool_call
            error_msg = str(exc.original_error)
            failed_result = ToolExecutionResult(
                call=tool_call,
                error=error_msg,
            )
            step.tool_results.append(failed_result)
            yield ToolFailedEvent(
                step_index=step_index,
                step_id=step_id,
                tool_call=tool_call,
                tool_name=tool_call.name,
                tool_result=failed_result,
                error=error_msg,
            )

        yield StepFailedEvent(
            step_index=step_index,
            step_id=step_id,
            step=step,
            error=error_msg,
        )
        yield RunFailedEvent(
            step_index=step_index,
            step_id=step_id,
            step=step,
            error=error_msg,
        )

    async def run(
        self,
        question: str,
        trace_ctx: Context | None = None,
        **request_meta: Unpack[LLMRequestMeta],
    ) -> AsyncGenerator[AgentEvent, None]:
        """Yield AgentEvent entries as the agent reasons."""
        messages = [LLMMessage.build(role="user", content=question)]
        async for event in self._run_with_messages(
            question,
            messages=messages,
            trace_ctx=trace_ctx,
            **request_meta,
        ):
            yield event

    async def resume(
        self,
        question: str,
        history: list[LLMMessage],
        trace_ctx: Context | None = None,
        **request_meta: Unpack[LLMRequestMeta],
    ) -> AsyncGenerator[AgentEvent, None]:
        """Yield AgentEvent entries with existing conversation history."""
        messages = list(history) + [LLMMessage.build(role="user", content=question)]
        async for event in self._run_with_messages(
            question,
            messages=messages,
            trace_ctx=trace_ctx,
            **request_meta,
        ):
            yield event

    async def _run_with_messages(
        self,
        question: str,
        *,
        messages: list[LLMMessage],
        trace_ctx: Context | None,
        **request_meta: Unpack[LLMRequestMeta],
    ) -> AsyncGenerator[AgentEvent, None]:
        run_span = self._tracer.start_span(
            "agent.run",
            kind=SpanKind.INTERNAL,
            context=trace_ctx,
            attributes={
                "agent.max_steps": self._max_steps,
                "agent.run.input": question,
                "agent.run.output": "",
                "langfuse.trace.name": "aceai.run",
                "langfuse.trace.input": question,
                "langfuse.trace.output": "",
            },
        )
        run_context = set_span_in_context(run_span, trace_ctx or Context())
        set_trace_ctx(run_context)
        self._ctx_mgr.init_context(messages)

        steps: list[AgentStep] = []
        run_state = RunState()
        error_msg: str | None = None
        try:
            for _ in range(self._max_steps):
                step_id = str(uuid4())
                event_builder = AgentEventBuilder(
                    step_index=len(steps),
                    step_id=step_id,
                )
                step_gen = self._run_step(
                    event_builder=event_builder,
                    run_state=run_state,
                    **request_meta,
                )
                try:
                    async for event in step_gen:
                        if isinstance(event, LLMCompletedEvent):
                            steps.append(event.step)
                        elif isinstance(event, RunCompletedEvent):
                            run_span.set_attribute(
                                "agent.run.output", event.final_answer
                            )
                            run_span.set_attribute(
                                "langfuse.trace.output", event.final_answer
                            )
                            yield event
                            return
                        yield event
                except Exception as exc:
                    if not steps:
                        raise
                    error_msg = str(exc)
                    last_step = steps[-1]
                    last_index = len(steps) - 1
                    for event in self._handle_failed_step(
                        last_step, last_index, exc=exc
                    ):
                        yield event
                    raise
                finally:
                    await step_gen.aclose()
            else:
                error_msg = (
                    f"Agent exceeded maximum steps: {self._max_steps} without answering"
                )
                last_step = steps[-1]
                last_index = len(steps) - 1
                for event in self._handle_failed_step(
                    last_step, last_index, error_msg=error_msg
                ):
                    yield event
                raise AceAIRuntimeError(error_msg)
        finally:
            if error_msg is not None:
                run_span.set_attribute("agent.run.output", error_msg)
                run_span.set_attribute("langfuse.trace.output", error_msg)
            set_trace_ctx(None)
            if run_span.is_recording():
                run_span.end()

    async def ask(
        self,
        question: str,
        trace_ctx: Context | None = None,
        **request_meta: Unpack[LLMRequestMeta],
    ) -> str:
        """Run the agent to completion and return the final answer in plain text."""
        async for event in self.run(
            question,
            trace_ctx=trace_ctx,
            **request_meta,
        ):
            if isinstance(event, RunCompletedEvent):
                return event.final_answer

        raise AceAIRuntimeError("Agent run did not complete successfully")
