import re
import json
from openai import OpenAI
from langchain.prompts import ChatPromptTemplate
from langchain.document_loaders import PyPDFLoader
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_openai import ChatOpenAI

class AiService:
    def __init__(self, model="gpt-4.1-mini"):
        self.OPENAI_MODEL = model
        self.embeddings = OpenAIEmbeddings(
            model="text-embedding-3-small"
        )
        self.openai = OpenAI()

    def generate_response(self, question_history=[], docs = [], extra_context: str = "", json_structure=None, instructions: str = None):
        # simplified $SELECTION_PLACEHOLDER$
        original_question = question_history[-1]["content"]

        # attach supporting docs
        if docs:
            print(f"Found {len(docs)} supporting documents")
            fragments = [str(d["payload"]) for d in docs]
            question_history[-1]["content"] += "\n\nSupporting documents:\n" + "\n\n".join(fragments)

        # attach extra context
        if extra_context:
            question_history[-1]["content"] += "\n\nExtra context:\n" + extra_context

        try:
            # common parameters
            params = {
                "instructions": instructions,
                "input": question_history,
                "model": self.OPENAI_MODEL,
                "temperature": 0.0,
                "top_p": 0.25
            }

            if json_structure:
                params["text_format"] = json_structure
                result = self.openai.responses.parse(**params)
                content = result.output_parsed
            else:
                result = self.openai.responses.create(**params)
                content = re.sub(r"\$", r"\\$", result.output_text)

            return content

        except Exception as e:
            print(f"Error during API request: {e}")
            return ""

        finally:
            # restore original question text
            question_history[-1]["content"] = original_question
    
    def generate_single_response(self, query, system_message, json_output=False, json_structure=None):
        question_history = [{"role": "system", "content": system_message}, {"role": "user", "content": query}]
        return self.generate_response(question_history, json_structure=json_structure)


    def json_response_cleaner(self, response):
        response = response.replace("\n", "").replace(
            "\t", "").replace("\\", "")
        return json.loads(response)

    def init_messages(self, system_message):
        return [{"role": "system", "content": system_message}]

    def add_user_message(self, messages, user_message):
        messages.append({"role": "user", "content": user_message})
        return messages

    def add_assistant_message(self, messages, assistant_message):
        messages.append({"role": "assistant", "content": assistant_message})
        return messages

    def search_pdf(self, pdf_path, query, schema, output_format):
        """
        Search through a PDF document and extract structured information.
        
        Args:
            pdf_path (str): Path to the PDF file
            query (str): Search query
            schema (dict): Schema defining the required output structure
        
        Returns:
            dict: Extracted structured information
        """
        # Load and split the PDF
        loader = PyPDFLoader(pdf_path)
        pages = loader.load()
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=100
        )
        splits = text_splitter.split_documents(pages)

        combined_content = "\n".join([chunk.page_content for chunk in splits])

        prompt = ChatPromptTemplate.from_template(
            """Extract the requested information from the following text, just write the answer to the question do not include that how the document discusses it or some other form of this. DIRECTLY ANSWER THE QUESTION.:
            
            {text}
            
            {schema}
            
            Query: {query}"""
        )

        messages = prompt.format_messages(
            text=combined_content,
            schema=schema,
            query=query
        )

        llm = ChatOpenAI(temperature=0, model=self.OPENAI_MODEL).with_structured_output(output_format)
        response = llm.invoke(messages)
        return response