"""
title: Mancer Manifold Pipe
author: ncks024
version: 0.1.0
"""

from pydantic import BaseModel, Field
from typing import Union, Generator, Iterator
import requests
import logging

logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)


class Pipe:
    class Valves(BaseModel):
        NAME_PREFIX: str = Field(
            default="MANCER/",
            description="The prefix applied before the model names.",
        )
        MANCER_API_BASE_URL: str = Field(
            default="https://mancer.tech/oai/v1",
            description="The base URL for Mancer API endpoints.",
        )
        MANCER_API_KEY: str = Field(
            default="",
            description="Required API key to retrieve the model list.",
        )

    class UserValves(BaseModel):
        MANCER_API_KEY: str = Field(
            default="",
            description="User-specific API key for accessing Mancer services.",
        )

    def __init__(self):
        self.type = "manifold"
        self.valves = self.Valves()
        logger.debug(f"Initialized Valves: {self.valves}")

    def pipes(self):
        if self.valves.MANCER_API_KEY:
            try:
                headers = {
                    "Authorization": f"Bearer {self.valves.MANCER_API_KEY}",
                    "Content-Type": "application/json",
                }
                logger.debug(
                    f"Fetching models with API Key: {self.valves.MANCER_API_KEY}"
                )

                response = requests.get(
                    f"{self.valves.MANCER_API_BASE_URL}/models", headers=headers
                )
                response.raise_for_status()

                models = response.json()
                logger.debug(f"Fetched models: {models}")
                return [
                    {
                        "id": model["id"],
                        "name": f'{self.valves.NAME_PREFIX}{model.get("name", model["id"])}',
                    }
                    for model in models["data"]
                ]
            except requests.exceptions.RequestException as e:
                logger.error(f"Error while fetching models: {e}")
                return [
                    {
                        "id": "error",
                        "name": "Could not fetch models from Mancer. Please check the API Key.",
                    }
                ]
        else:
            logger.error("Global API Key not provided.")
            return [
                {
                    "id": "error",
                    "name": "Global API Key not provided.",
                }
            ]

    def pipe(self, body: dict, __user__: dict) -> Union[str, Generator, Iterator]:
        logger.info(f"Executing pipe function: {__name__}")
        logger.debug(f"User info: {__user__}")

        user_valves = __user__.get("valves")

        if not user_valves:
            error_message = "Error: User Valves not configured."
            logger.error(error_message)
            return error_message

        if isinstance(user_valves, dict):
            api_key = user_valves.get("MANCER_API_KEY", self.valves.MANCER_API_KEY)
        else:
            api_key = (
                user_valves.MANCER_API_KEY or self.valves.MANCER_API_KEY
            )  # Access attributes directly if it's a Pydantic model

        logger.debug(f"Using API Key: {api_key}")

        if not api_key:
            error_message = (
                "Error: MANCER_API_KEY not provided by the user or globally."
            )
            logger.error(error_message)
            return error_message

        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        }

        model_id = body.get("model", "").split(".")[-1]
        payload = {**body, "model": model_id}
        logger.debug(f"Payload for Mancer API: {payload}")

        try:
            response = requests.post(
                url=f"{self.valves.MANCER_API_BASE_URL}/chat/completions",
                json=payload,
                headers=headers,
                stream=True,
            )

            response.raise_for_status()

            logger.info("Mancer API request successful.")
            if body.get("stream"):
                return response.iter_lines()
            else:
                response_json = response.json()
                logger.debug(f"Mancer API response: {response_json}")
                return response_json
        except requests.exceptions.RequestException as e:
            error_message = f"Error: Could not complete request to Mancer API - {e}"
            logger.error(error_message)
            return error_message