import os
import json
import logging
import anthropic
from openai import OpenAI
from abc import ABC, abstractmethod
from tools.github_tool import GitHubTool

# Initialize GitHubTool and get functions
github_tool = GitHubTool()
functions = github_tool.get_functions()

class AIProvider(ABC):
    @abstractmethod
    def get_chat_response(self, messages):
        pass

    @abstractmethod
    def format_messages(self, messages):
        pass

    @abstractmethod
    def format_tool_calls(self, response):
        pass

    @abstractmethod
    def format_tool_result(self, tool_call, tool_response):
        pass

class AnthropicProvider(AIProvider):
    def __init__(self):
        self.client = anthropic.Anthropic(
            api_key=os.environ.get("ANTHROPIC_API_KEY"),
            default_headers={"anthropic-beta": "max-tokens-3-5-sonnet-2024-07-15"}
        )
        self.model = "claude-3-5-sonnet-20240620"

    def get_chat_response(self, messages):
        try:
            response = self.client.messages.create(
                model=self.model,
                system=messages[0]['content'],
                messages=self.format_messages(messages[1:]),
                max_tokens=8192,
                tools=self.format_tools()
            )
            return response
        except Exception as e:
            logging.error(f"An error occurred: {str(e)}")
            return None

    def format_messages(self, messages):
        return messages

    def format_tool_calls(self, response):
        tool_calls = []
        for message in response.content:
            if message.type == "tool_use":
                tool_calls.append(message)
        return tool_calls

    def format_tools(self):
        return [
            {
                "name": function['name'],
                "description": function['description'],
                "input_schema": function['parameters'] if function['parameters'] not in [None, {}] else {"type": "object", "properties": {"param1": {"type": "string", "description": "Unnecessary"}}, "required": []}
            }
            for function in functions
        ]

    def format_assistant_reply(self, response):
        for message in response.content:
            if message.type == "text":
                return message.text
        return ""

    def get_reply_text(self, response):
        return self.format_assistant_reply(response)

    def get_model(self):
        return self.model

    def format_tool_result(self, tool_call, tool_response):
        return {
            "role": "function",
            "name": tool_call.name,
            "content": json.dumps(tool_response)
        }

class OpenAIProvider(AIProvider):
    def __init__(self, use_smart_model=True):
        self.client = OpenAI()
        self.use_smart_model = use_smart_model
        self.model = self.get_model()

    def get_model(self):
        return "gpt-4o" if self.use_smart_model else "gpt-4o-mini"

    def get_chat_response(self, messages):
        response = self.client.chat.completions.create(
            model=self.model,
            messages=self.format_messages(messages),
            functions=functions,
            function_call="auto",
            max_tokens=self.get_max_tokens()
        )
        return response

    def format_messages(self, messages):
        return messages

    def format_tool_calls(self, response):
        tool_calls = []
        assistant_message = response.choices[0].message
        if hasattr(assistant_message, 'function_call') and assistant_message.function_call is not None:
            tool_calls.append(assistant_message.function_call)
        return tool_calls

    def get_max_tokens(self):
        return 4096 if self.model == "gpt-4o" else 16384

    def format_tool_result(self, tool_call, tool_response):
        return {
            "role": "function",
            "name": tool_call.name,
            "content": json.dumps(tool_response)
        }

def create_ai_provider(provider_name="anthropic", use_smart_model=True):
    if provider_name.lower() == "anthropic":
        return AnthropicProvider()
    elif provider_name.lower() == "openai":
        return OpenAIProvider(use_smart_model)
    else:
        raise ValueError(f"Unknown provider: {provider_name}")