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

logging.basicConfig(level=logging.DEBUG)

class Pipeline:
    class Valves(BaseModel):
        OR_MODEL: str
        OR_URL: str
        OR_KEY: str

    def __init__(self):
        self.name = "Basic Openrouter Chat"
        self.valves = self.Valves(
            **{
                "pipelines": ["*"],
                "OR_MODEL": os.getenv("OR_MODEL", "anthropic/claude-3.5-haiku:beta"),
                "OR_URL": os.getenv("OR_URL", "https://openrouter.ai/api/v1/chat/completions"),
                "OR_KEY": os.getenv("OR_KEY", "OPENROUTER_API_KEY")
            }
        )

    async def on_startup(self):
        pass

    async def on_shutdown(self):
        pass

    def pipe(self, user_message: str, model_id: str, messages: List[dict], body: dict) -> Union[str, Generator, Iterator]:
        try:
            response = requests.post(
                url = self.valves.OR_URL,
                headers = {
                    "Authorization": f"Bearer {self.valves.OR_KEY}"
                },

                data = json.dumps({
                    "model": self.valves.OR_MODEL,
                    "messages": [
                        {
                            "role": "user",
                            "content": user_message
                        }
                    ]
                })
            )
        
            if response.status_code == 200:
                response_data = response.json()

                return response_data["choices"][0]["message"]["content"]
            else:
                logging.error(f"Error response {response.status_code}: {response.text}")
                return f"Error: {response.status_code}"
        except requests.HTTPError as e:
            logging.error(f"Clientresponse error: {e}")
            return "HTTP backend error"
        except Exception as e:
            logging.error(f"Unexpected error: {e}")
            return f"Unexpected error: {e}"