from __future__ import annotations

import asyncio
import inspect
import logging
import time
import uuid
from dataclasses import dataclass
from typing import Any, Dict, Optional

from .errors import SecurityBlockError
from .security import get_injection_scanner

logger = logging.getLogger("agentid")


try:
    from langchain_core.callbacks import BaseCallbackHandler  # type: ignore
except Exception:  # pragma: no cover
    BaseCallbackHandler = object  # type: ignore


@dataclass
class _RunState:
    input_text: str
    started_at: float
    model: Optional[str] = None
    client_event_id: Optional[str] = None
    guard_event_id: Optional[str] = None


def _safe_text(val: Any) -> str:
    if isinstance(val, str):
        return val
    return ""


def _extract_prompt_from_prompts(prompts: Any) -> Optional[str]:
    if isinstance(prompts, list) and prompts:
        last = prompts[-1]
        return _safe_text(last)
    return None


def _extract_prompt_from_messages(messages: Any) -> Optional[str]:
    # LangChain chat hooks can deliver BaseMessage[][] or similar.
    flat: list[Any] = []
    if isinstance(messages, list):
        for item in messages:
            if isinstance(item, list):
                flat.extend(item)
            else:
                flat.append(item)

    last = None
    for msg in flat:
        role = getattr(msg, "type", None) or getattr(msg, "role", None)
        if role in ("human", "user"):
            last = msg
        if isinstance(msg, dict) and msg.get("role") == "user":
            last = msg

    if last is None:
        return None

    if isinstance(last, dict):
        return _safe_text(last.get("content"))

    return _safe_text(getattr(last, "content", None))


def _set_prompt_in_prompts(prompts: Any, sanitized: str) -> bool:
    if not isinstance(prompts, list) or not prompts:
        return False
    prompts[-1] = sanitized
    return True


def _set_prompt_in_messages(messages: Any, sanitized: str) -> bool:
    if not isinstance(messages, list):
        return False

    flat: list[Any] = []
    for item in messages:
        if isinstance(item, list):
            flat.extend(item)
        else:
            flat.append(item)

    for msg in reversed(flat):
        role = getattr(msg, "type", None) or getattr(msg, "role", None)
        if role not in ("human", "user") and not (
            isinstance(msg, dict) and msg.get("role") == "user"
        ):
            continue

        if isinstance(msg, dict):
            if "content" in msg:
                msg["content"] = sanitized
                return True
            if "text" in msg:
                msg["text"] = sanitized
                return True
            msg["content"] = sanitized
            return True

        if hasattr(msg, "content"):
            try:
                setattr(msg, "content", sanitized)
                return True
            except Exception:
                continue

    return False


def _extract_model_name(serialized: Any, kwargs: Any) -> Optional[str]:
    # Best-effort; varies by integration.
    if isinstance(kwargs, dict):
        model = kwargs.get("model") or kwargs.get("model_name")
        if isinstance(model, str) and model:
            return model
    if isinstance(serialized, dict):
        model = serialized.get("name") or serialized.get("id")
        if isinstance(model, str) and model:
            return model
    return None


def _as_bool(value: Any) -> Optional[bool]:
    if isinstance(value, bool):
        return value
    return None


def _extract_stream_flag(serialized: Any, kwargs: Any) -> bool:
    if isinstance(kwargs, dict):
        direct = _as_bool(kwargs.get("stream"))
        if direct is None:
            direct = _as_bool(kwargs.get("streaming"))
        if direct is not None:
            return direct

        invocation = kwargs.get("invocation_params")
        if isinstance(invocation, dict):
            inv_stream = _as_bool(invocation.get("stream"))
            if inv_stream is None:
                inv_stream = _as_bool(invocation.get("streaming"))
            if inv_stream is not None:
                return inv_stream

    if isinstance(serialized, dict):
        serialized_kwargs = serialized.get("kwargs")
        if isinstance(serialized_kwargs, dict):
            ser_stream = _as_bool(serialized_kwargs.get("stream"))
            if ser_stream is None:
                ser_stream = _as_bool(serialized_kwargs.get("streaming"))
            if ser_stream is not None:
                return ser_stream

    return False


def _extract_output_text(response: Any) -> str:
    # LLMResult.generations: list[list[Generation]]
    gens = getattr(response, "generations", None)
    if isinstance(gens, list) and gens:
        first_list = gens[0]
        if isinstance(first_list, list) and first_list:
            gen0 = first_list[0]
            text = getattr(gen0, "text", None)
            if isinstance(text, str):
                return text
            message = getattr(gen0, "message", None)
            content = getattr(message, "content", None)
            if isinstance(content, str):
                return content
    return ""


def _extract_token_usage(response: Any) -> Optional[Dict[str, Any]]:
    llm_output = getattr(response, "llm_output", None) or getattr(response, "llmOutput", None)
    if isinstance(llm_output, dict):
        usage = llm_output.get("token_usage") or llm_output.get("tokenUsage") or llm_output.get("usage")
        if isinstance(usage, dict):
            return usage

    gens = getattr(response, "generations", None)
    if isinstance(gens, list) and gens:
        first_list = gens[0]
        if isinstance(first_list, list) and first_list:
            gen0 = first_list[0]
            message = getattr(gen0, "message", None)

            usage_meta = getattr(message, "usage_metadata", None)
            if isinstance(usage_meta, dict):
                return usage_meta

            response_meta = getattr(message, "response_metadata", None)
            if isinstance(response_meta, dict):
                token_usage = response_meta.get("token_usage") or response_meta.get("tokenUsage")
                if isinstance(token_usage, dict):
                    return token_usage

            generation_info = getattr(gen0, "generation_info", None)
            if isinstance(generation_info, dict):
                token_usage = generation_info.get("token_usage") or generation_info.get("tokenUsage")
                if isinstance(token_usage, dict):
                    return token_usage
    return None


def _coerce_uuid_str(value: Any) -> Optional[str]:
    if not isinstance(value, str):
        return None
    candidate = value.strip()
    if not candidate:
        return None
    try:
        return str(uuid.UUID(candidate))
    except Exception:
        return None


class AgentIDCallbackHandler(BaseCallbackHandler):
    """
    LangChain callback handler:
    - runs shared local enforcement before request leaves the client
    - blocks execution immediately when AgentID guard returns allowed=False
    - logs results to AgentID ingest (fire-and-forget via AgentID client)
    """

    def __init__(self, agent: Any, system_id: str, api_key: Optional[str] = None):
        self.agent = agent
        self.system_id = system_id
        self.api_key = api_key.strip() if isinstance(api_key, str) else None
        self._runs: Dict[str, _RunState] = {}

    def _scan_injection_sync(self, input_text: str) -> None:
        cfg = getattr(self.agent, "_config", None)
        if cfg is None or not bool(getattr(cfg, "check_injection", False)):
            return

        effective_api_key = self.api_key or str(getattr(cfg, "api_key", "") or "").strip()
        base_url = str(getattr(cfg, "base_url", "") or "").strip()
        if not effective_api_key or not base_url:
            return

        get_injection_scanner().scan(
            prompt=input_text,
            api_key=effective_api_key,
            base_url=base_url,
            ai_scan_enabled=bool(getattr(cfg, "ai_scan_enabled", True)),
            store_pii=bool(getattr(cfg, "store_pii", False)),
            pii_manager=getattr(self.agent, "_pii", None),
            source="python_langchain",
        )

    def _langchain_capabilities(self) -> Dict[str, Any]:
        cfg = getattr(self.agent, "_config", None)
        pii_masking_enabled = bool(getattr(cfg, "pii_masking", False)) if cfg else False
        return {
            "capabilities": {
                "has_feedback_handler": True,
                "pii_masking_enabled": pii_masking_enabled,
                "framework": "langchain",
            }
        }

    async def _scan_injection_async(self, input_text: str) -> None:
        loop = asyncio.get_running_loop()
        await loop.run_in_executor(None, self._scan_injection_sync, input_text)

    def _guard_sync(
        self,
        input_text: str,
        model: Optional[str] = None,
        client_event_id: Optional[str] = None,
    ) -> Dict[str, Any]:
        return self.agent.guard(
            input=input_text,
            system_id=self.system_id,
            model=model,
            client_event_id=client_event_id,
            api_key=self.api_key,
            client_capabilities=self._langchain_capabilities(),
        )

    async def _guard_async(
        self,
        input_text: str,
        model: Optional[str] = None,
        client_event_id: Optional[str] = None,
    ) -> Dict[str, Any]:
        guard_fn = getattr(self.agent, "guard", None)
        if guard_fn is None:
            return {"allowed": True, "fail_open": True}

        if inspect.iscoroutinefunction(guard_fn):
            return await guard_fn(
                input=input_text,
                system_id=self.system_id,
                model=model,
                client_event_id=client_event_id,
                api_key=self.api_key,
                client_capabilities=self._langchain_capabilities(),
            )

        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(
            None, self._guard_sync, input_text, model, client_event_id
        )

    def _enforce_local_sync(self, input_text: str, stream: bool) -> str:
        enforce_fn = getattr(self.agent, "_apply_local_security", None)
        if not callable(enforce_fn):
            return input_text

        sanitized, _flags = enforce_fn(
            system_id=self.system_id,
            input_text=input_text,
            stream=stream,
            api_key=self.api_key,
        )
        return sanitized

    async def _enforce_local_async(self, input_text: str, stream: bool) -> str:
        async_enforce_fn = getattr(self.agent, "_apply_local_security_async", None)
        if callable(async_enforce_fn) and inspect.iscoroutinefunction(async_enforce_fn):
            sanitized, _flags = await async_enforce_fn(
                system_id=self.system_id,
                input_text=input_text,
                stream=stream,
                api_key=self.api_key,
            )
            return sanitized

        sync_enforce_fn = getattr(self.agent, "_apply_local_security", None)
        if not callable(sync_enforce_fn):
            return input_text

        loop = asyncio.get_running_loop()

        def _run_sync() -> str:
            sanitized, _flags = sync_enforce_fn(
                system_id=self.system_id,
                input_text=input_text,
                stream=stream,
                api_key=self.api_key,
            )
            return sanitized

        return await loop.run_in_executor(None, _run_sync)

    def _block_if_needed(self, verdict: Dict[str, Any]) -> None:
        if verdict.get("allowed") is False:
            raise SecurityBlockError(verdict.get("reason"))

    @staticmethod
    def _resolve_transformed_input(verdict: Dict[str, Any], fallback_input: str) -> str:
        transformed = verdict.get("transformed_input")
        if isinstance(transformed, str) and transformed:
            return transformed
        return fallback_input

    # ---- Sync hooks ----
    def on_llm_start(self, serialized: Dict[str, Any], prompts: Any, **kwargs: Any) -> Any:
        run_id = str(kwargs.get("run_id") or kwargs.get("runId") or "")
        input_text = _extract_prompt_from_prompts(prompts) or ""
        if not input_text:
            raise ValueError("AgentID: No prompt found. Security guard requires string input.")

        self._scan_injection_sync(input_text)

        stream = _extract_stream_flag(serialized, kwargs)
        sanitized_input = self._enforce_local_sync(input_text, stream)
        if sanitized_input != input_text and not _set_prompt_in_prompts(prompts, sanitized_input):
            raise ValueError("AgentID: Strict PII mode requires mutable LangChain prompt payload.")

        model = _extract_model_name(serialized, kwargs.get("invocation_params") or kwargs)
        request_client_event_id = str(uuid.uuid4())
        verdict = self._guard_sync(
            sanitized_input, model=model, client_event_id=request_client_event_id
        )
        self._block_if_needed(verdict)
        canonical_client_event_id = (
            _coerce_uuid_str(verdict.get("client_event_id")) or request_client_event_id
        )
        guard_event_id = verdict.get("guard_event_id")
        guard_event_id_str = guard_event_id if isinstance(guard_event_id, str) and guard_event_id else None
        transformed_input = self._resolve_transformed_input(verdict, sanitized_input)
        if transformed_input != sanitized_input and not _set_prompt_in_prompts(prompts, transformed_input):
            raise ValueError(
                "AgentID: Guard transformed input could not be applied to LangChain prompt payload."
            )
        self._runs[run_id] = _RunState(
            input_text=transformed_input,
            started_at=time.perf_counter(),
            model=_extract_model_name(serialized, kwargs.get("invocation_params") or kwargs),
            client_event_id=canonical_client_event_id,
            guard_event_id=guard_event_id_str,
        )

    def on_chat_model_start(self, serialized: Dict[str, Any], messages: Any, **kwargs: Any) -> Any:
        run_id = str(kwargs.get("run_id") or kwargs.get("runId") or "")
        input_text = _extract_prompt_from_messages(messages) or ""
        if not input_text:
            raise ValueError(
                "AgentID: No user message found. Security guard requires string input."
            )

        self._scan_injection_sync(input_text)

        stream = _extract_stream_flag(serialized, kwargs)
        sanitized_input = self._enforce_local_sync(input_text, stream)
        if sanitized_input != input_text and not _set_prompt_in_messages(messages, sanitized_input):
            raise ValueError("AgentID: Strict PII mode requires mutable LangChain message payload.")

        model = _extract_model_name(serialized, kwargs.get("invocation_params") or kwargs)
        request_client_event_id = str(uuid.uuid4())
        verdict = self._guard_sync(
            sanitized_input, model=model, client_event_id=request_client_event_id
        )
        self._block_if_needed(verdict)
        canonical_client_event_id = (
            _coerce_uuid_str(verdict.get("client_event_id")) or request_client_event_id
        )
        guard_event_id = verdict.get("guard_event_id")
        guard_event_id_str = guard_event_id if isinstance(guard_event_id, str) and guard_event_id else None
        transformed_input = self._resolve_transformed_input(verdict, sanitized_input)
        if transformed_input != sanitized_input and not _set_prompt_in_messages(messages, transformed_input):
            raise ValueError(
                "AgentID: Guard transformed input could not be applied to LangChain message payload."
            )
        self._runs[run_id] = _RunState(
            input_text=transformed_input,
            started_at=time.perf_counter(),
            model=_extract_model_name(serialized, kwargs.get("invocation_params") or kwargs),
            client_event_id=canonical_client_event_id,
            guard_event_id=guard_event_id_str,
        )

    def on_llm_end(self, response: Any, **kwargs: Any) -> Any:
        run_id = str(kwargs.get("run_id") or kwargs.get("runId") or "")
        state = self._runs.pop(run_id, None)
        if not state:
            return

        latency_ms = int((time.perf_counter() - state.started_at) * 1000)
        output_text = _extract_output_text(response)
        usage = _extract_token_usage(response)
        metadata: Dict[str, Any] = {}
        if state.client_event_id:
            metadata["client_event_id"] = state.client_event_id
        if state.guard_event_id:
            metadata["guard_event_id"] = state.guard_event_id

        self.agent.log(
            system_id=self.system_id,
            input=state.input_text,
            output=output_text,
            event_id=state.client_event_id,
            model=state.model or "unknown",
            usage=usage if isinstance(usage, dict) else None,
            metadata=metadata or None,
            latency=latency_ms,
            api_key=self.api_key,
            client_capabilities=self._langchain_capabilities(),
        )

    def on_llm_error(self, error: BaseException, **kwargs: Any) -> Any:
        run_id = str(kwargs.get("run_id") or kwargs.get("runId") or "")
        state = self._runs.pop(run_id, None)
        input_text = state.input_text if state else ""
        metadata: Dict[str, Any] = {
            "error_type": error.__class__.__name__,
            "error_message": str(error),
        }
        if state and state.client_event_id:
            metadata["client_event_id"] = state.client_event_id
        if state and state.guard_event_id:
            metadata["guard_event_id"] = state.guard_event_id

        self.agent.log(
            system_id=self.system_id,
            input=input_text,
            output="",
            event_id=state.client_event_id if state else None,
            model=state.model if state and state.model else "unknown",
            event_type="error",
            severity="error",
            metadata=metadata,
            api_key=self.api_key,
            client_capabilities=self._langchain_capabilities(),
        )

    # ---- Async hooks ----
    async def aon_llm_start(self, serialized: Dict[str, Any], prompts: Any, **kwargs: Any) -> Any:
        run_id = str(kwargs.get("run_id") or kwargs.get("runId") or "")
        input_text = _extract_prompt_from_prompts(prompts) or ""
        if not input_text:
            raise ValueError("AgentID: No prompt found. Security guard requires string input.")

        await self._scan_injection_async(input_text)

        stream = _extract_stream_flag(serialized, kwargs)
        sanitized_input = await self._enforce_local_async(input_text, stream)
        if sanitized_input != input_text and not _set_prompt_in_prompts(prompts, sanitized_input):
            raise ValueError("AgentID: Strict PII mode requires mutable LangChain prompt payload.")

        model = _extract_model_name(serialized, kwargs.get("invocation_params") or kwargs)
        request_client_event_id = str(uuid.uuid4())
        verdict = await self._guard_async(
            sanitized_input, model=model, client_event_id=request_client_event_id
        )
        self._block_if_needed(verdict)
        canonical_client_event_id = (
            _coerce_uuid_str(verdict.get("client_event_id")) or request_client_event_id
        )
        guard_event_id = verdict.get("guard_event_id")
        guard_event_id_str = guard_event_id if isinstance(guard_event_id, str) and guard_event_id else None
        transformed_input = self._resolve_transformed_input(verdict, sanitized_input)
        if transformed_input != sanitized_input and not _set_prompt_in_prompts(prompts, transformed_input):
            raise ValueError(
                "AgentID: Guard transformed input could not be applied to LangChain prompt payload."
            )
        self._runs[run_id] = _RunState(
            input_text=transformed_input,
            started_at=time.perf_counter(),
            model=_extract_model_name(serialized, kwargs.get("invocation_params") or kwargs),
            client_event_id=canonical_client_event_id,
            guard_event_id=guard_event_id_str,
        )

    async def aon_chat_model_start(
        self, serialized: Dict[str, Any], messages: Any, **kwargs: Any
    ) -> Any:
        run_id = str(kwargs.get("run_id") or kwargs.get("runId") or "")
        input_text = _extract_prompt_from_messages(messages) or ""
        if not input_text:
            raise ValueError(
                "AgentID: No user message found. Security guard requires string input."
            )

        await self._scan_injection_async(input_text)

        stream = _extract_stream_flag(serialized, kwargs)
        sanitized_input = await self._enforce_local_async(input_text, stream)
        if sanitized_input != input_text and not _set_prompt_in_messages(messages, sanitized_input):
            raise ValueError("AgentID: Strict PII mode requires mutable LangChain message payload.")

        model = _extract_model_name(serialized, kwargs.get("invocation_params") or kwargs)
        request_client_event_id = str(uuid.uuid4())
        verdict = await self._guard_async(
            sanitized_input, model=model, client_event_id=request_client_event_id
        )
        self._block_if_needed(verdict)
        canonical_client_event_id = (
            _coerce_uuid_str(verdict.get("client_event_id")) or request_client_event_id
        )
        guard_event_id = verdict.get("guard_event_id")
        guard_event_id_str = guard_event_id if isinstance(guard_event_id, str) and guard_event_id else None
        transformed_input = self._resolve_transformed_input(verdict, sanitized_input)
        if transformed_input != sanitized_input and not _set_prompt_in_messages(messages, transformed_input):
            raise ValueError(
                "AgentID: Guard transformed input could not be applied to LangChain message payload."
            )
        self._runs[run_id] = _RunState(
            input_text=transformed_input,
            started_at=time.perf_counter(),
            model=_extract_model_name(serialized, kwargs.get("invocation_params") or kwargs),
            client_event_id=canonical_client_event_id,
            guard_event_id=guard_event_id_str,
        )

    async def aon_llm_end(self, response: Any, **kwargs: Any) -> Any:
        # logging is fire-and-forget already; keep sync call
        self.on_llm_end(response, **kwargs)

    async def aon_llm_error(self, error: BaseException, **kwargs: Any) -> Any:
        self.on_llm_error(error, **kwargs)
