"""
title: Cohere Manifold Pipe
author: justinh-rahb
author_url: https://github.com/justinh-rahb
funding_url: https://github.com/open-webui
version: 0.1.5
license: MIT
"""

import os
import json
import requests
from typing import List, Union, Generator, Iterator
from pydantic import BaseModel, Field

# Set DEBUG to True to enable detailed logging
DEBUG = False


class Pipe:
    class Valves(BaseModel):
        COHERE_API_BASE_URL: str = Field(default="https://api.cohere.com/v1")
        COHERE_API_KEY: str = Field(default="")

    def __init__(self):
        self.type = "manifold"
        self.id = "cohere"
        self.name = "cohere/"
        self.valves = self.Valves(**{"COHERE_API_KEY": os.getenv("COHERE_API_KEY", "")})

    def get_cohere_models(self):
        if self.valves.COHERE_API_KEY:
            try:
                headers = {
                    "Authorization": f"Bearer {self.valves.COHERE_API_KEY}",
                    "Content-Type": "application/json",
                    "Accept": "application/json",
                }
                if DEBUG:
                    print(
                        f"Fetching models from: {self.valves.COHERE_API_BASE_URL}/models"
                    )
                    print(f"Headers: {headers}")
                r = requests.get(
                    f"{self.valves.COHERE_API_BASE_URL}/models", headers=headers
                )
                r.raise_for_status()
                models = r.json()
                if DEBUG:
                    print(f"Models response: {models}")
                return [
                    {
                        "id": model["name"],
                        "name": model.get("name", model["name"]),
                    }
                    for model in models["models"]
                ]
            except Exception as e:
                if DEBUG:
                    print(f"Error fetching Cohere models: {e}")
                return [
                    {
                        "id": "cohere",
                        "name": f"Could not fetch models from Cohere: {str(e)}",
                    },
                ]
        else:
            if DEBUG:
                print("COHERE_API_KEY is not set")
            return []

    def pipes(self) -> List[dict]:
        return self.get_cohere_models()

    def pipe(self, body: dict) -> Union[str, Generator, Iterator]:
        try:
            model = body["model"]
            messages = body["messages"]
            stream = body.get("stream", False)

            if DEBUG:
                print("Incoming body:", json.dumps(body, indent=2))

            # Strip "cohere." prefix from model name
            if model.startswith("cohere."):
                model = model[7:]

            if stream:
                return self.stream_response(model, messages)
            else:
                return self.get_completion(model, messages)
        except Exception as e:
            if DEBUG:
                print(f"Error in pipe method: {e}")
            return f"Error: {e}"

    def stream_response(self, model: str, messages: List[dict]) -> Generator:
        headers = {
            "Authorization": f"Bearer {self.valves.COHERE_API_KEY}",
            "Content-Type": "application/json",
            "Accept": "application/json",
        }

        chat_history = self._format_chat_history(messages[:-1])
        user_message = messages[-1]["content"]
        system_message = self._get_system_message(messages)

        payload = {
            "model": model,
            "chat_history": chat_history,
            "message": user_message,
            "stream": True,
            **({"system": str(system_message)} if system_message else {}),
        }

        if DEBUG:
            print("Stream request:")
            print(f"URL: {self.valves.COHERE_API_BASE_URL}/chat")
            print(f"Headers: {headers}")
            print(f"Payload: {json.dumps(payload, indent=2)}")

        try:
            r = requests.post(
                url=f"{self.valves.COHERE_API_BASE_URL}/chat",
                json=payload,
                headers=headers,
                stream=True,
            )
            r.raise_for_status()

            for line in r.iter_lines():
                if line:
                    try:
                        event = json.loads(line.decode("utf-8"))
                        if DEBUG:
                            print(f"Received event: {event}")
                        if event["event_type"] == "text-generation":
                            yield event["text"]
                        elif event["event_type"] == "stream-end":
                            break
                    except json.JSONDecodeError:
                        if DEBUG:
                            print(f"Failed to decode JSON: {line}")
                        pass
        except requests.RequestException as e:
            if DEBUG:
                print(f"Request exception in stream_response: {e}")
                print(
                    f"Response content: {r.content if 'r' in locals() else 'No response'}"
                )
            yield f"Error: {str(e)}"

    def get_completion(self, model: str, messages: List[dict]) -> str:
        headers = {
            "Authorization": f"Bearer {self.valves.COHERE_API_KEY}",
            "Content-Type": "application/json",
            "Accept": "application/json",
        }

        chat_history = self._format_chat_history(messages[:-1])
        user_message = messages[-1]["content"]
        system_message = self._get_system_message(messages)

        payload = {
            "model": model,
            "chat_history": chat_history,
            "message": user_message,
            **({"system": str(system_message)} if system_message else {}),
        }

        if DEBUG:
            print("Completion request:")
            print(f"URL: {self.valves.COHERE_API_BASE_URL}/chat")
            print(f"Headers: {headers}")
            print(f"Payload: {json.dumps(payload, indent=2)}")

        try:
            r = requests.post(
                url=f"{self.valves.COHERE_API_BASE_URL}/chat",
                json=payload,
                headers=headers,
            )
            r.raise_for_status()
            data = r.json()
            if DEBUG:
                print(f"Completion response: {json.dumps(data, indent=2)}")
            return data.get("text", "No response from Cohere.")
        except requests.RequestException as e:
            if DEBUG:
                print(f"Request exception in get_completion: {e}")
                print(
                    f"Response content: {r.content if 'r' in locals() else 'No response'}"
                )
            return f"Error: {str(e)}"

    def _format_chat_history(self, messages: List[dict]) -> List[dict]:
        return [
            {
                "role": "USER" if message["role"] == "user" else "CHATBOT",
                "message": message["content"],
            }
            for message in messages
            if message["role"] != "system"
        ]

    def _get_system_message(self, messages: List[dict]) -> Union[str, None]:
        system_messages = [
            msg["content"] for msg in messages if msg["role"] == "system"
        ]
        return str(system_messages[0]) if system_messages else None