"""
title: Grok Manifold Pipe
author: Siraj Florida
author_url: https://github.com/sirajperson
funding_url: https://www.kinothen.com
version: 0.1.0
license: GNUv3
"""


import os
import json
import requests
from typing import List, Union, Generator, Iterator, Optional
from pydantic import BaseModel, Field
from utils.misc import pop_system_message


class Pipe:
    class Valves(BaseModel):
        """Configuration for the Grok API interactions."""

        GROK_API_KEY: str = Field(default="", description="API key for Grok services.")
        GROK_API_BASE_URL: str = Field(
            default="https://api.x.ai/v1",
            description="Base URL for Grok API endpoints.",
        )
        MAX_TOKENS: int = Field(
            default=4096, description="Maximum number of tokens to generate."
        )
        TEMPERATURE: float = Field(default=0.8, description="Sampling temperature.")
        TOP_P: float = Field(default=0.9, description="Nucleus sampling top_p value.")
        STREAM: bool = Field(default=False, description="Whether to stream responses.")

    def __init__(self):
        self.valves = self.Valves(
            GROK_API_KEY=os.getenv("GROK_API_KEY", ""),
            GROK_API_BASE_URL=os.getenv("GROK_API_BASE_URL", "https://api.x.ai/v1"),
        )
        self.type = "manifold"
        self.id = "grok"
        self.name = "grok/"

    def get_model_id(self, model_name: str) -> str:
        """Extract just the base model name from any format"""
        # Split on both / and . to handle any format
        parts = model_name.replace(".", "/").split("/")
        # Return only the actual model name (e.g. "grok-beta")
        return parts[-1]

    def get_grok_models(self):
        headers = {
            "Authorization": f"Bearer {self.valves.GROK_API_KEY}",
            "Content-Type": "application/json",
        }
        try:
            response = requests.get(
                f"{self.valves.GROK_API_BASE_URL}/models", headers=headers
            )
            response.raise_for_status()
            models_data = response.json()
            return [
                {"id": model["id"], "name": model["id"]}
                for model in models_data.get("data", [])
            ]
        except Exception as e:
            print(f"Error fetching models: {e}")
            return []

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

    def process_image(self, image_data):
        if image_data["image_url"]["url"].startswith("data:image"):
            mime_type, base64_data = image_data["image_url"]["url"].split(",", 1)
            media_type = mime_type.split(":")[1].split(";")[0]
            return {
                "type": "image",
                "source": {
                    "type": "base64",
                    "media_type": media_type,
                    "data": base64_data,
                },
            }
        else:
            return {
                "type": "image",
                "source": {"type": "url", "url": image_data["image_url"]["url"]},
            }

    def pipe(self, body: dict) -> Union[str, Generator, Iterator]:
        system_message, messages = pop_system_message(body.get("messages", []))

        processed_messages = []
        for message in messages:
            if isinstance(message.get("content"), list):
                for item in message["content"]:
                    if item["type"] == "text":
                        processed_messages.append(
                            {"role": message["role"], "content": item["text"]}
                        )
                    elif item["type"] == "image_url":
                        processed_image = self.process_image(item)
                        processed_messages.append(processed_image)
            else:
                processed_messages.append(
                    {"role": message["role"], "content": message.get("content", "")}
                )

        # Include system message if present
        if system_message:
            processed_messages.insert(
                0, {"role": "system", "content": str(system_message)}
            )

        # Extract just the base model name
        model_id = self.get_model_id(body["model"])

        # Structure payload according to API spec
        payload = {
            "model": model_id,
            "messages": processed_messages,
            "stream": body.get("stream", self.valves.STREAM),
            "temperature": body.get("temperature", self.valves.TEMPERATURE),
            "max_tokens": body.get("max_tokens", self.valves.MAX_TOKENS),
            "top_p": body.get("top_p", self.valves.TOP_P),
            "frequency_penalty": body.get("frequency_penalty", 0),
            "presence_penalty": body.get("presence_penalty", 0),
            "stop": body.get("stop", []),
            "user": body.get("user", ""),
            "n": body.get("n", 1),
            "logprobs": body.get("logprobs", False),
            "top_logprobs": body.get("top_logprobs", 0),
        }

        headers = {
            "Authorization": f"Bearer {self.valves.GROK_API_KEY}",
            "Content-Type": "application/json",
        }

        url = f"{self.valves.GROK_API_BASE_URL}/chat/completions"

        try:
            if payload["stream"]:
                return self.stream_response(url, headers, payload)
            else:
                return self.non_stream_response(url, headers, payload)
        except Exception as e:
            print(f"Error in pipe method: {e}")
            return f"Error: {e}"

    def stream_response(self, url, headers, payload):
        with requests.post(url, headers=headers, json=payload, stream=True) as response:
            if response.status_code != 200:
                raise Exception(f"HTTP Error {response.status_code}: {response.text}")

            for line in response.iter_lines():
                if line:
                    line = line.decode("utf-8")
                    if line.startswith("data: "):
                        try:
                            data = json.loads(line[6:])
                            if "choices" in data and len(data["choices"]) > 0:
                                delta = data["choices"][0].get("delta", {})
                                content = delta.get("content", "")
                                if content:
                                    yield content
                        except json.JSONDecodeError:
                            print(f"Failed to parse JSON: {line}")
                        except KeyError as e:
                            print(f"Unexpected data structure: {e}")
                            print(f"Full data: {data}")

    def non_stream_response(self, url, headers, payload):
        response = requests.post(url, headers=headers, json=payload)
        if response.status_code != 200:
            raise Exception(f"HTTP Error {response.status_code}: {response.text}")

        res = response.json()
        return res["choices"][0]["message"]["content"] if res.get("choices") else ""