"""
LLM Service for generating AI summaries using LangChain
"""

import logging
import os
import re
import sys
import uuid as uuid_mod
from pathlib import Path
from typing import Any, Dict, List, Optional

from deep_translator import GoogleTranslator
from jinja2 import Environment, FileSystemLoader
from langchain_core.callbacks import UsageMetadataCallbackHandler
from langchain_core.messages import HumanMessage, SystemMessage

# Add utils to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent.parent))
from utils.langsmith_util import setup_langsmith_tracing  # noqa: E402
from utils.llm_factory import get_llm  # noqa: E402

# Setup LangSmith tracing
setup_langsmith_tracing()

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# LangSmith trace URL helpers
# ---------------------------------------------------------------------------


def is_langsmith_tracing_enabled() -> bool:
    """Check if LangSmith tracing is enabled."""
    return os.getenv("LANGCHAIN_TRACING_V2", "").lower() == "true" and bool(
        os.getenv("LANGCHAIN_API_KEY")
    )


def _resolve_langsmith_url_prefix() -> Optional[str]:
    """Resolve the LangSmith project URL prefix (cached after first call)."""
    cache_attr = "_cached_prefix"
    if hasattr(_resolve_langsmith_url_prefix, cache_attr):
        return getattr(_resolve_langsmith_url_prefix, cache_attr)

    result: Optional[str] = None
    try:
        from langsmith import Client

        client = Client()
        tenant_id = client._get_tenant_id()
        project_name = os.getenv("LANGCHAIN_PROJECT", "default")
        project = client.read_project(project_name=project_name)
        host = os.getenv("LANGCHAIN_ENDPOINT", "https://api.smith.langchain.com")
        app_host = host.replace("api.", "").rstrip("/")
        if "api" in app_host:
            app_host = "https://smith.langchain.com"
        result = f"{app_host}/o/{tenant_id}/projects/p/{project.id}"
    except Exception as exc:
        logger.debug("Could not resolve LangSmith URL prefix: %s", exc)

    setattr(_resolve_langsmith_url_prefix, cache_attr, result)
    return result


def summarize_usage_metadata(
    handler: UsageMetadataCallbackHandler,
    model_key: Optional[str],
) -> Dict[str, Any]:
    """Collapse a ``UsageMetadataCallbackHandler`` into activity-log fields.

    The handler accumulates per-model usage as a nested dict; we sum
    across whichever models the LLM reported (typically just one) so a
    single PATCH call carries the totals plus the configured ``model_key``
    (the user-facing key from ``supported_llms``).
    """
    payload: Dict[str, Any] = {}
    if model_key:
        payload["llm_model"] = model_key
    try:
        per_model = handler.usage_metadata or {}
    except Exception:  # pragma: no cover - handler API surface
        per_model = {}
    if not per_model:
        return payload
    prompt = sum(int(v.get("input_tokens", 0) or 0) for v in per_model.values())
    completion = sum(int(v.get("output_tokens", 0) or 0) for v in per_model.values())
    if prompt:
        payload["prompt_tokens"] = prompt
    if completion:
        payload["completion_tokens"] = completion
    return payload


def get_langsmith_trace_url(run_id: uuid_mod.UUID) -> Optional[str]:
    """Construct LangSmith trace URL for a given run ID.

    Returns None if LangSmith tracing is not configured or the URL
    prefix could not be resolved.
    """
    if not is_langsmith_tracing_enabled():
        return None
    try:
        prefix = _resolve_langsmith_url_prefix()
        if prefix:
            return f"{prefix}/r/{run_id}?poll=true"
    except Exception:
        pass
    return None


# Initialize Jinja2 environment for prompt templates
PROMPTS_DIR = Path(__file__).resolve().parents[3] / "prompts"
jinja_env = Environment(loader=FileSystemLoader(str(PROMPTS_DIR)), autoescape=True)

# Load templates at module level
_system_template = jinja_env.get_template("ai_summary_system.j2")
_user_template = jinja_env.get_template("ai_summary_user.j2")


def render_prompt(
    query: str,
    results: List[Dict[str, Any]],
    max_results: int = 20,
    system_prompt_override: str | None = None,
) -> str:
    """
    Render the full prompt for debugging/transparency purposes.

    Args:
        query: The search query string
        results: List of search results
        max_results: Maximum number of results to include
        system_prompt_override: Optional custom system prompt (from group settings)

    Returns:
        The fully rendered prompt text with system and user messages
    """
    # Limit to top N results
    top_results = results[:max_results]

    # Render prompts from templates (use override if provided)
    system_prompt = system_prompt_override or _system_template.render()
    user_prompt = _user_template.render(query=query, results=top_results)

    # Combine into full prompt
    full_prompt = f"SYSTEM MESSAGE:\n{system_prompt}\n\nUSER MESSAGE:\n{user_prompt}"

    return full_prompt


async def stream_ai_summary(
    query: str,
    results: List[Dict[str, Any]],
    max_results: int = 20,
    model_key: str | None = None,
    temperature: float | None = None,
    max_tokens: int | None = None,
    system_prompt_override: str | None = None,
):
    """
    Stream an AI summary of search results token by token.

    Args:
        query: The search query string
        results: List of search results (dictionaries with title, organization, text, etc.)
        max_results: Maximum number of results to include in the prompt (default: 20)
        system_prompt_override: Optional custom system prompt (from group settings)

    Yields:
        Individual tokens as they are generated by the LLM.
        The *last* yielded item is a dict with metadata including
        ``langsmith_trace_url`` when LangSmith tracing is active.
    """
    try:
        # Limit to top N results
        top_results = results[:max_results]

        # Render prompts from templates (use override if provided)
        system_prompt = system_prompt_override or _system_template.render()
        user_prompt = _user_template.render(query=query, results=top_results)

        logger.info(
            f"Streaming AI summary for query: '{query}' using {len(top_results)} results"
        )

        # Log the prompts
        logger.info("=" * 80)
        logger.info("AI Summary Request (Streaming)")
        logger.info("=" * 80)
        logger.info("SYSTEM PROMPT:")
        logger.info(system_prompt)
        logger.info("-" * 80)
        logger.info("USER PROMPT:")
        logger.info(user_prompt)
        logger.info("=" * 80)

        # Get LLM instance
        llm = get_llm(model=model_key, temperature=temperature, max_tokens=max_tokens)

        # Use LangChain's astream for async streaming
        messages = [
            SystemMessage(content=system_prompt),
            HumanMessage(content=user_prompt),
        ]

        # Generate a deterministic run_id for LangSmith tracing
        run_id: Optional[uuid_mod.UUID] = None
        usage_handler = UsageMetadataCallbackHandler()
        config: Dict[str, Any] = {"callbacks": [usage_handler]}
        if is_langsmith_tracing_enabled():
            run_id = uuid_mod.uuid4()
            config["run_id"] = run_id

        # Stream tokens from the LLM
        accumulated = ""
        async for chunk in llm.astream(
            messages, config=config  # type: ignore[arg-type]
        ):
            if chunk.content:
                token = str(chunk.content)
                accumulated += token
                yield token

        logger.info(f"✓ Streamed AI summary ({len(accumulated)} chars)")

        # Yield metadata dict as final item (picked up by the route layer)
        metadata: Dict[str, Any] = summarize_usage_metadata(usage_handler, model_key)
        if run_id:
            trace_url = get_langsmith_trace_url(run_id)
            if trace_url:
                metadata["langsmith_trace_url"] = trace_url
        if metadata:
            yield metadata

    except Exception as e:
        logger.error(f"Error streaming AI summary: {e}", exc_info=True)
        raise


async def generate_ai_summary(
    query: str,
    results: List[Dict[str, Any]],
    max_results: int = 20,
    model_key: str | None = None,
    temperature: float | None = None,
    max_tokens: int | None = None,
    system_prompt_override: str | None = None,
) -> str:
    """
    Generate an AI summary of search results using LangChain + HuggingFace Inference API.

    Args:
        query: The search query string
        results: List of search results (dictionaries with title, organization, text, etc.)
        max_results: Maximum number of results to include in the prompt (default: 20)
        system_prompt_override: Optional custom system prompt (from group settings)

    Returns:
        Generated summary text
    """
    summary, _ = await generate_ai_summary_with_usage(
        query=query,
        results=results,
        max_results=max_results,
        model_key=model_key,
        temperature=temperature,
        max_tokens=max_tokens,
        system_prompt_override=system_prompt_override,
    )
    return summary


async def generate_ai_summary_with_usage(
    query: str,
    results: List[Dict[str, Any]],
    max_results: int = 20,
    model_key: str | None = None,
    temperature: float | None = None,
    max_tokens: int | None = None,
    system_prompt_override: str | None = None,
) -> tuple[str, Dict[str, Any]]:
    """Generate an AI summary and return ``(summary, usage_metadata)``.

    ``usage_metadata`` carries the same shape emitted by the streaming
    endpoint: ``llm_model``, ``prompt_tokens``, ``completion_tokens``
    when the provider reports usage.
    """
    try:
        # Limit to top N results
        top_results = results[:max_results]

        # Render prompts from templates (use override if provided)
        system_prompt = system_prompt_override or _system_template.render()
        user_prompt = _user_template.render(query=query, results=top_results)

        logger.info(
            f"Generating AI summary for query: '{query}' using {len(top_results)} results"
        )

        # Log the prompts
        logger.info("=" * 80)
        logger.info("AI Summary Request")
        logger.info("=" * 80)
        logger.info("SYSTEM PROMPT:")
        logger.info(system_prompt)
        logger.info("-" * 80)
        logger.info("USER PROMPT:")
        logger.info(user_prompt)
        logger.info("=" * 80)

        # Get LLM instance
        llm = get_llm(model=model_key, temperature=temperature, max_tokens=max_tokens)

        # Use LangChain's ainvoke for async invocation with chat format
        messages = [
            SystemMessage(content=system_prompt),
            HumanMessage(content=user_prompt),
        ]
        usage_handler = UsageMetadataCallbackHandler()
        response = await llm.ainvoke(
            messages, config={"callbacks": [usage_handler]}  # type: ignore[arg-type]
        )

        # Return raw content, matching the streaming endpoint behaviour
        summary = str(response.content).strip()

        # Log the response
        logger.info("AI SUMMARY RESPONSE:")
        logger.info(summary)
        logger.info("=" * 80)

        logger.info(f"✓ Generated AI summary ({len(summary)} chars)")

        usage = summarize_usage_metadata(usage_handler, model_key)
        return summary, usage

    except Exception as e:
        logger.error(f"Error generating AI summary: {e}", exc_info=True)
        raise


async def translate_text(
    text: str, target_language: str, source_language: str | None = None
) -> str:
    """
    Translate text using deep-translator (Google Translate) instead of LLM.
    Uses regex to protect reference numbers like [64] from being mangled.

    Args:
        source_language: ISO code of the original text language.
            When provided, passed to GoogleTranslator instead of "auto".
    """
    if not text:
        return ""

    # Map full language names or codes to deep-translator ISO codes.
    # Note: "zh" is not accepted by GoogleTranslator; use "zh-CN".
    lang_map = {
        "english": "en",
        "french": "fr",
        "spanish": "es",
        "arabic": "ar",
        "chinese": "zh-CN",
        "portuguese": "pt",
        "russian": "ru",
        "swahili": "sw",
        "hindi": "hi",
        "bengali": "bn",
        "german": "de",
        "greek": "el",
        "italian": "it",
        "lithuanian": "lt",
        "vietnamese": "vi",
        "dutch": "nl",
        "polish": "pl",
        "turkish": "tr",
        "japanese": "ja",
        "korean": "ko",
        "en": "en",
        "fr": "fr",
        "es": "es",
        "ar": "ar",
        "zh": "zh-CN",
        "pt": "pt",
        "ru": "ru",
        "sw": "sw",
        "hi": "hi",
        "bn": "bn",
        "de": "de",
        "el": "el",
        "it": "it",
        "lt": "lt",
        "vi": "vi",
        "nl": "nl",
        "pl": "pl",
        "tr": "tr",
        "ja": "ja",
        "ko": "ko",
    }

    target_lang_code = lang_map.get(target_language.lower(), "en")
    source_lang_code = (
        lang_map.get(source_language.lower(), "auto") if source_language else "auto"
    )

    try:
        # 1. Protect references: [64] -> __REF_64__
        protected_text = text
        ref_pattern = r"\[(\d+)\]"
        refs = re.findall(ref_pattern, text)
        for ref_num in set(refs):
            protected_text = protected_text.replace(
                f"[{ref_num}]", f"__REF_{ref_num}__"
            )

        # 2. Protect newlines to prevent flattening
        # Replace \n\n with __PARA__ and \n with __BR__
        protected_text = protected_text.replace("\n\n", " __PARA__ ")
        protected_text = protected_text.replace("\n", " __BR__ ")

        # 3. Perform translation
        # deep-translator is synchronous, suitable for direct call here.
        translator = GoogleTranslator(source=source_lang_code, target=target_lang_code)
        translated_text = translator.translate(protected_text)

        # 4. Restore references: __REF_64__ -> [64]
        if translated_text:
            # Restore references
            restore_ref_pattern = r"__\s*REF\s*_\s*(\d+)\s*__"

            def replace_match(match):
                return f"[{match.group(1)}]"

            final_text = re.sub(restore_ref_pattern, replace_match, translated_text)

            # Restore newlines (handling potential extra spaces added by translator)
            # __PARA__ -> \n\n
            final_text = re.sub(r"\s*__\s*PARA\s*__\s*", "\n\n", final_text)
            # __BR__ -> \n
            final_text = re.sub(r"\s*__\s*BR\s*__\s*", "\n", final_text)

            return final_text.strip()

        return text

    except Exception as e:
        logger.error(f"Translation failed: {e}")
        return text
