import os
import boto3
import openai
from paperqa import Docs
from urllib.parse import urlparse
from typing import Dict, List, Union
from langchain.base_language import BaseLanguageModel
from langchain.chat_models import ChatOpenAI, ChatAnthropic

from pydantic import BaseModel, Field

class KnowledgeBase(Docs, BaseModel):
    llms: Dict[str, BaseLanguageModel] = Field(default_factory=dict)
    documents: List[str] = Field(default_factory=list)
    temperature: float = Field(0.7)

    def __init__(self, documents: List[str], *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.llms = {}
        self.documents = documents
        self.temperature = kwargs.get('temperature', 0.7)

        # Always Initialize both LLMs
        self.llms['gpt4'] = ChatOpenAI(model='gpt-4', temperature=self.temperature)
        self.llms['claude'] = ChatAnthropic(model="claude-2", temperature=self.temperature)
        self.llms['both'] = Docs(llm=self.llms['gpt4'], summary_llm=self.llms['claude'])

        # Process documents for each LLM
        for llm in self.llms.values():
            self.process_documents(llm)

    def process_documents(self, llm: BaseLanguageModel):
        # Add documents to Paper QA
        for path in self.documents:
            parsed_url = urlparse(path)
            scheme = parsed_url.scheme
            ext = os.path.splitext(path)[-1].lower()

            if scheme in ['http', 'https']:
                self.add_url(path)
            elif scheme == 's3':
                s3 = boto3.client('s3')
                bucket = parsed_url.netloc
                key = parsed_url.path.lstrip('/')

                if key.endswith('/'):  # It's a directory
                    try:
                        objects = s3.list_objects(Bucket=bucket, Prefix=key)
                        for obj in objects.get('Contents', []):
                            file_key = obj['Key']
                            local_path = f"/tmp/{os.path.basename(file_key)}"
                            s3.download_file(bucket, file_key, local_path)
                            self.add(local_path)
                    except NoCredentialsError:
                        print("No AWS credentials found.")
                else:  # It's a single file
                    local_path = f"/tmp/{os.path.basename(key)}"
                    s3.download_file(bucket, key, local_path)

                    if ext == '.pdf':
                        self.add(local_path)
                    else:
                        with open(local_path, 'r') as f:
                            text_content = f.read()
                            self.add_file(f, text_content)
            else:
                if ext == '.pdf':
                    self.add(path)
                else:
                    with open(path, 'r') as f:
                        text_content = f.read()
                        self.add_file(f, text_content)