from __future__ import annotations

import asyncio
import random
from collections.abc import Awaitable, Callable
from typing import TYPE_CHECKING

from loguru import logger
from pydantic_ai import Agent, DeferredToolRequests, DeferredToolResults, Tool
from pydantic_ai.exceptions import ModelHTTPError, UnexpectedModelBehavior
from pydantic_ai.models.groq import GroqModelSettings

from .. import constants as cs
from .. import exceptions as ex
from .. import logs as ls
from ..config import ModelConfig, settings
from ..prompts import (
    CYPHER_SYSTEM_PROMPT,
    LOCAL_CYPHER_SYSTEM_PROMPT,
    MESSAGE_SUMMARIZER_PROMPT,
    build_rag_orchestrator_prompt,
)
from ..providers.base import get_provider_from_config

if TYPE_CHECKING:
    from pydantic_ai.models import Model
    from pydantic_ai.settings import ModelSettings


def _create_provider_model(config: ModelConfig) -> Model:
    provider = get_provider_from_config(config)
    return provider.create_model(config.model_id)


def _is_rate_limit_error(error: Exception) -> bool:
    if isinstance(error, ModelHTTPError):
        return error.status_code == 429
    if isinstance(error, UnexpectedModelBehavior):
        message = str(error).lower()
        return "429" in message or "rate limit" in message or "too many requests" in message
    return False


def _build_model_settings(config: ModelConfig) -> ModelSettings | None:
    if config.provider != cs.Provider.GROQ:
        return None

    groq_settings: GroqModelSettings = {
        "parallel_tool_calls": bool(getattr(settings, "GROQ_PARALLEL_TOOL_CALLS", False)),
        "temperature": float(getattr(settings, "GROQ_TEMPERATURE", 0.0)),
    }

    max_tokens = getattr(settings, "GROQ_MAX_TOKENS", None)
    if max_tokens:
        groq_settings["max_tokens"] = int(max_tokens)

    return groq_settings


async def run_agent_with_rate_limit_retry[T](
    run_factory: Callable[[], Awaitable[T]],
    *,
    operation_name: str,
) -> T:
    max_retries = int(getattr(settings, "MODEL_RATE_LIMIT_RETRIES", 5))
    base_delay = float(getattr(settings, "MODEL_RATE_LIMIT_BASE_DELAY", 1.0))
    max_delay = float(getattr(settings, "MODEL_RATE_LIMIT_MAX_DELAY", 20.0))

    for attempt in range(max_retries + 1):
        try:
            return await run_factory()
        except Exception as e:
            if not _is_rate_limit_error(e) or attempt >= max_retries:
                raise

            delay = min(max_delay, base_delay * (2**attempt)) + random.uniform(0, 0.25)
            logger.warning(
                "Rate limited during {} (attempt {}/{}). Retrying in {:.2f}s.",
                operation_name,
                attempt + 1,
                max_retries + 1,
                delay,
            )
            await asyncio.sleep(delay)

    # Unreachable: loop either returns or raises, but mypy needs a final raise.
    raise RuntimeError("Rate-limit retry loop exited unexpectedly")


def _clean_cypher_response(response_text: str) -> str:
    query = response_text.strip().replace(cs.CYPHER_BACKTICK, "")
    if query.startswith(cs.CYPHER_PREFIX):
        query = query[len(cs.CYPHER_PREFIX) :].strip()
    if not query.endswith(cs.CYPHER_SEMICOLON):
        query += cs.CYPHER_SEMICOLON
    return query


class CypherGenerator:
    def __init__(self) -> None:
        try:
            config = settings.active_cypher_config
            llm = _create_provider_model(config)

            system_prompt = (
                LOCAL_CYPHER_SYSTEM_PROMPT
                if config.provider == cs.Provider.OLLAMA
                else CYPHER_SYSTEM_PROMPT
            )

            self.agent = Agent(
                model=llm,
                system_prompt=system_prompt,
                output_type=str,
                retries=settings.AGENT_RETRIES,
                model_settings=_build_model_settings(config),
            )
        except Exception as e:
            raise ex.LLMGenerationError(ex.LLM_INIT_CYPHER.format(error=e)) from e

    async def generate(self, natural_language_query: str) -> str:
        logger.info(ls.CYPHER_GENERATING.format(query=natural_language_query))
        try:
            result = await run_agent_with_rate_limit_retry(
                lambda: self.agent.run(natural_language_query),
                operation_name="cypher generation",
            )
            if (
                not isinstance(result.output, str)
                or cs.CYPHER_MATCH_KEYWORD not in result.output.upper()
            ):
                raise ex.LLMGenerationError(
                    ex.LLM_INVALID_QUERY.format(output=result.output)
                )

            query = _clean_cypher_response(result.output)
            logger.info(ls.CYPHER_GENERATED.format(query=query))
            return query
        except Exception as e:
            logger.error(ls.CYPHER_ERROR.format(error=e))
            raise ex.LLMGenerationError(ex.LLM_GENERATION_FAILED.format(error=e)) from e


class MessageSummarizer:
    """Compresses conversation history after every 4 messages to manage token usage."""

    SUMMARIZE_THRESHOLD = 4  # Trigger summarization after this many messages
    KEEP_RECENT = 2  # Keep last N messages fresh (don't summarize them)

    def __init__(self) -> None:
        try:
            config = settings.active_orchestrator_config
            self.llm = _create_provider_model(config)
            self.model_settings = _build_model_settings(config)
        except Exception as e:
            logger.warning("Failed to initialize message summarizer: {}", e)
            self.llm = None

    async def maybe_summarize(self, message_history: list) -> list:
        """Check if messages need summarization and compress if threshold reached."""
        if not message_history or not self.llm:
            return message_history

        # Only summarize if we have threshold messages
        if len(message_history) < self.SUMMARIZE_THRESHOLD:
            return message_history

        # Don't summarize further if already compressed (check for summary marker)
        if any("Previous conversation summary:" in str(msg) for msg in message_history):
            return message_history

        try:
            # Split messages: old ones to summarize, recent ones to keep
            recent_count = self.KEEP_RECENT
            messages_to_summarize = message_history[:-recent_count]
            messages_to_keep = message_history[-recent_count:]

            # Format old messages for the summarizer
            conversation_text = "\n\n".join([str(msg) for msg in messages_to_summarize])

            # Call summarizer
            summary = await self._summarize(conversation_text)

            # Create summary message and preserve recent messages
            summary_msg = f"Previous conversation summary:\n{summary}"

            result = [summary_msg] + messages_to_keep
            logger.info(
                "Compressed {} messages into summary + {} recent messages",
                len(messages_to_summarize),
                len(messages_to_keep),
            )
            return result

        except Exception as e:
            logger.warning("Message summarization failed: {}. Keeping original history.", e)
            return message_history

    async def _summarize(self, conversation_text: str) -> str:
        """Use LLM to create a concise summary of conversation."""
        try:
            summary_agent = Agent(
                model=self.llm,
                system_prompt=MESSAGE_SUMMARIZER_PROMPT,
                output_type=str,
                retries=settings.AGENT_RETRIES,
                model_settings=self.model_settings,
            )

            result = await run_agent_with_rate_limit_retry(
                lambda: summary_agent.run(conversation_text),
                operation_name="message summarization",
            )
            return str(result.output).strip()
        except Exception as e:
            logger.error("Failed to generate message summary: {}", e)
            raise


def create_rag_orchestrator(tools: list[Tool]) -> Agent:
    try:
        config = settings.active_orchestrator_config
        llm = _create_provider_model(config)

        return Agent(
            model=llm,
            system_prompt=build_rag_orchestrator_prompt(tools),
            tools=tools,
            retries=settings.AGENT_RETRIES,
            output_retries=settings.ORCHESTRATOR_OUTPUT_RETRIES,
            output_type=[str, DeferredToolRequests],
            model_settings=_build_model_settings(config),
        )
    except Exception as e:
        raise ex.LLMGenerationError(ex.LLM_INIT_ORCHESTRATOR.format(error=e)) from e


# Global message summarizer instance (created on demand)
_message_summarizer: MessageSummarizer | None = None


def get_message_summarizer() -> MessageSummarizer:
    """Get or create the global message summarizer."""
    global _message_summarizer
    if _message_summarizer is None:
        _message_summarizer = MessageSummarizer()
    return _message_summarizer


async def run_rag_agent_with_summarization(
    agent: Agent,
    prompt: str | None,
    message_history: list | None = None,
    deferred_tool_results: DeferredToolResults | None = None,
    model: Model | None = None,
) -> object:
    """Run RAG agent with automatic message history compression.
    
    Summarizes old messages after every 4 messages to prevent token overflow
    on models with small limits (e.g., Codestral-2501 with 8k limit).
    """
    # Apply message summarization if history is present
    if message_history:
        summarizer = get_message_summarizer()
        message_history = await summarizer.maybe_summarize(message_history)

    # Run agent with potentially compressed history
    return await agent.run(
        prompt,
        message_history=message_history,
        deferred_tool_results=deferred_tool_results,
        model=model,
    )
