# Copyright (c) 2025 Microsoft Corporation.
# Licensed under the MIT License

"""OpenRouter LLM providers for NLWeb."""

import asyncio
import json
import logging
from typing import Any

from nlweb_core.llm import GenerativeLLMProvider
from nlweb_core.scoring import ScoringContext, ScoringLLMProvider
from openai import AsyncOpenAI

logger = logging.getLogger(__name__)

OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1"


def normalize_schema_for_structured_output(schema: dict[str, Any]) -> dict[str, Any]:
    """Normalize schema for OpenRouter structured outputs."""
    normalized = schema.copy()
    normalized["additionalProperties"] = False
    if "properties" in normalized and "required" not in normalized:
        normalized["required"] = list(normalized["properties"].keys())
    return normalized


class OpenRouterProvider(GenerativeLLMProvider):
    """Implementation of GenerativeLLMProvider for OpenRouter."""

    def __init__(
        self,
        api_key: str,
        model: str,
        site_url: str | None = None,
        site_name: str | None = None,
        **kwargs,
    ):
        """Initialize the OpenRouter provider.

        Args:
            api_key: OpenRouter API key
            model: Model identifier (e.g., "openai/gpt-4o", "anthropic/claude-3.5-sonnet")
            site_url: Optional site URL for OpenRouter ranking
            site_name: Optional site name for OpenRouter ranking
            **kwargs: Additional configuration (ignored)
        """
        super().__init__(**kwargs)
        if not api_key:
            raise ValueError("api_key is required for OpenRouter")

        self.api_key = api_key
        self.model = model
        self.site_url = site_url
        self.site_name = site_name
        self._client: AsyncOpenAI | None = None

    async def _ensure_client(self) -> None:
        """Create client if not already initialized."""
        if self._client is not None:
            return

        extra_headers = {}
        if self.site_url:
            extra_headers["HTTP-Referer"] = self.site_url
        if self.site_name:
            extra_headers["X-Title"] = self.site_name

        self._client = AsyncOpenAI(
            base_url=OPENROUTER_BASE_URL,
            api_key=self.api_key,
            default_headers=extra_headers if extra_headers else None,
            timeout=30.0,
        )

    async def get_completion(
        self,
        prompt: str,
        schema: dict[str, Any],
        temperature: float = 0.7,
        max_tokens: int = 2048,
        timeout: float = 8.0,
        **kwargs,
    ) -> dict[str, Any]:
        """Get completion from OpenRouter.

        Args:
            prompt: The prompt to send to the model
            schema: JSON schema for the expected response
            temperature: Model temperature
            max_tokens: Maximum tokens in the generated response
            timeout: Request timeout in seconds
            **kwargs: Additional provider-specific arguments (ignored)

        Returns:
            Parsed JSON response

        Raises:
            ValueError: If the response cannot be parsed as valid JSON
            TimeoutError: If the request times out
        """
        await self._ensure_client()
        assert self._client is not None

        normalized_schema = normalize_schema_for_structured_output(schema)
        system_prompt = "You are a helpful assistant. Respond with JSON matching the required schema."

        try:
            response = await asyncio.wait_for(
                self._client.chat.completions.create(
                    messages=[
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": prompt},
                    ],
                    max_tokens=max_tokens,
                    temperature=temperature,
                    top_p=0.1,
                    stream=False,
                    model=self.model,
                    response_format={
                        "type": "json_schema",
                        "json_schema": {
                            "name": "response_schema",
                            "strict": True,
                            "schema": normalized_schema,
                        },
                    },
                ),
                timeout=timeout,
            )

            if not response or not hasattr(response, "choices") or not response.choices:
                return {}

            if not hasattr(response.choices[0], "message") or not hasattr(
                response.choices[0].message, "content"
            ):
                return {}

            content = response.choices[0].message.content
            if content is None:
                return {}

            return json.loads(content)

        except TimeoutError:
            raise
        except Exception:
            raise

    async def close(self) -> None:
        """Close the OpenRouter client and release resources."""
        if self._client is not None:
            await self._client.close()
            self._client = None


class OpenRouterScoringProvider(ScoringLLMProvider):
    """Implementation of ScoringLLMProvider for OpenRouter."""

    def __init__(
        self,
        api_key: str,
        model: str,
        site_url: str | None = None,
        site_name: str | None = None,
        **kwargs,
    ):
        """Initialize the OpenRouter scoring provider.

        Args:
            api_key: OpenRouter API key
            model: Model identifier (e.g., "openai/gpt-4o-mini")
            site_url: Optional site URL for OpenRouter ranking
            site_name: Optional site name for OpenRouter ranking
            **kwargs: Additional configuration (ignored)
        """
        super().__init__(**kwargs)
        self.api_key = api_key
        self.model = model
        self.site_url = site_url
        self.site_name = site_name
        self._client: AsyncOpenAI | None = None

    async def _ensure_client(self) -> None:
        """Create client if not already initialized."""
        if self._client is not None:
            return

        extra_headers = {}
        if self.site_url:
            extra_headers["HTTP-Referer"] = self.site_url
        if self.site_name:
            extra_headers["X-Title"] = self.site_name

        self._client = AsyncOpenAI(
            base_url=OPENROUTER_BASE_URL,
            api_key=self.api_key,
            default_headers=extra_headers if extra_headers else None,
            timeout=30.0,
        )

    def _build_scoring_prompt(self, context: ScoringContext) -> str:
        """Build a scoring prompt from context.

        Args:
            context: Structured context for scoring

        Returns:
            Formatted prompt string
        """
        if context.item_description and context.item_type:
            item_type = context.item_type or "item"

            prompt = f"""\
Assign a score between 0 and 100 to the following {item_type} based on how relevant it is to the user's question. Use your knowledge from other sources, about the item, to make a judgement.

If the score is above 50, provide a short description of the item highlighting the relevance to the user's question, without mentioning the user's question.

Provide an explanation of the relevance of the item to the user's question, without mentioning the user's question or the score or explicitly mentioning the term relevance.

If the score is below 75, in the description, include the reason why it is still relevant."""

            if context.publication_date and context.age_days is not None:
                prompt += f"""

FRESHNESS CONTEXT:
- Publication date: {context.publication_date}
- Age: {context.age_days} days old

When considering relevance, factor in the item's freshness based on the query intent:
- For queries asking for "latest", "recent", "new", or "today's" content, give higher scores to more recent items
- For queries about specific events, news, or time-sensitive topics, prioritize fresher content
- For evergreen topics (recipes, how-to guides, general information), age is less important
- Very recent items (< 7 days) should get a bonus for time-sensitive queries"""

            prompt += f"""

The user's question is: {context.query}

The item's description is: {context.item_description}"""
        elif context.intent:
            prompt = f"""\
Assign a score between 0 and 100 based on how well the user's query matches the specified intent.

User's query: {context.query}

Intent to check: {context.intent}

Provide a score indicating the match strength (0 = no match, 100 = perfect match)."""
        elif context.required_info:
            prompt = f"""\
Assign a score between 0 and 100 based on whether the user's query contains the required information.

User's query: {context.query}

Required information: {context.required_info}

Provide a score indicating presence (0 = not present, 100 = fully present)."""
        else:
            prompt = f"""Assign a score between 0 and 100 for the following query.

Query: {context.query}

Provide a relevance score."""

        return prompt

    async def score(
        self,
        questions: list[str],
        context: ScoringContext,
        timeout: float = 30.0,
        **kwargs,
    ) -> float:
        """Score a single context.

        Args:
            questions: List of scoring questions (ignored for LLM-based scoring)
            context: Structured context for the scoring operation
            timeout: Request timeout in seconds
            **kwargs: Additional provider-specific arguments

        Returns:
            Score between 0-100

        Raises:
            ValueError: If the request fails or response is invalid
            TimeoutError: If the request times out
        """
        await self._ensure_client()
        assert self._client is not None

        prompt = self._build_scoring_prompt(context)

        scoring_schema = {
            "type": "object",
            "properties": {
                "score": {
                    "type": "integer",
                    "description": "Relevance score between 0 and 100",
                },
                "description": {
                    "type": "string",
                    "description": "Brief explanation of the score",
                },
            },
            "required": ["score", "description"],
            "additionalProperties": False,
        }
        system_prompt = "You are a scoring assistant. Provide a relevance score and brief explanation."

        try:
            response = await asyncio.wait_for(
                self._client.chat.completions.create(
                    messages=[
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": prompt},
                    ],
                    max_tokens=500,
                    temperature=0.3,
                    top_p=0.1,
                    stream=False,
                    model=self.model,
                    response_format={
                        "type": "json_schema",
                        "json_schema": {
                            "name": "scoring_response",
                            "strict": True,
                            "schema": scoring_schema,
                        },
                    },
                ),
                timeout=timeout,
            )

            if (
                not response
                or not response.choices
                or not response.choices[0].message.content
            ):
                raise ValueError("Empty response from OpenRouter")

            content = response.choices[0].message.content
            result = json.loads(content)
            score = result.get("score", 0)

            if isinstance(score, (int, float)):
                return float(max(0, min(100, score)))
            else:
                raise ValueError(f"Invalid score type: {type(score)}")

        except TimeoutError:
            raise TimeoutError(f"OpenRouter scoring request timed out after {timeout}s")
        except Exception as e:
            raise ValueError(f"OpenRouter scoring failed: {e}")

    async def score_batch(
        self,
        questions: list[str],
        contexts: list[ScoringContext],
        timeout: float = 30.0,
        **kwargs,
    ) -> list[float | BaseException]:
        """Score multiple contexts with the given questions in parallel.

        Args:
            questions: List of scoring questions to ask
            contexts: List of contexts to score
            timeout: Request timeout in seconds
            **kwargs: Additional provider-specific arguments

        Returns:
            List of scores (0-100) or Exception for each failed request
        """
        tasks = [
            self.score(questions, context, timeout=timeout, **kwargs)
            for context in contexts
        ]
        results = await asyncio.gather(*tasks, return_exceptions=True)
        return list(results)

    async def close(self) -> None:
        """Close the OpenRouter client and release resources."""
        if self._client is not None:
            await self._client.close()
            self._client = None
