from fastapi import Form, APIRouter, Depends, HTTPException, status, Request
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from fastapi.responses import JSONResponse
from dotenv import load_dotenv
import os
import json
from typing import List, Optional
from openai import OpenAI
from sqlalchemy import text
from dependencies.HybridSearch import HybridSearch
from dependencies.helper import *
import db_config.database as dbase
from db_config.models import User, Folder, Documents, ChatHistory
import db_config.auth as auth


router = APIRouter()
security = HTTPBearer()
load_dotenv()

collection_name = f"{os.getenv('COLLECTION_NAME')}"
DB_FOLDER = str(os.getenv('DB_PATH'))+"/"
UPLOAD_FOLDER = str(os.getenv('STORAGE_PATH'))+"/"
openai_client=OpenAI()


@router.post("/text")
async def search_text(
    query: str = Form(...), 
    doc_id: Optional[str] = Form(None), 
    folder_id: int = Form(...),
    hybrid_weight: float = Form(...),
    top_k: Optional[int] = Form(15),
    limits: str = Form(None),
    # request: Request, 
    credentials: HTTPAuthorizationCredentials = Depends(security)
):
    # 
    # form = await request.form()
    # form_data = dict(form)
    
    # question = form_data.get('query')
    # filename = form_data.get('filename')
    # doc_id = form_data.get('doc_id')
    # folder_id = form_data.get('folder_id')
    # hybrid_weight = form_data.get('hybrid_weight',0.5)
    # answer_prompt_size = form_data.get('answer_prompt_size')
    # prompt_total_size = form_data.get('prompt_total_size')

    db = dbase.SessionLocal()
    token = credentials.credentials
    user = auth.get_current_user(db, token)
    try:
        folder = db.query(Folder).filter(Folder.user_id == user.user_id).filter(Folder.folder_id == folder_id).first()
        if doc_id:
            document = db.query(Documents).filter(Documents.doc_id == doc_id).first()
            if not document:
                return errorResponse(
                    "Document not found.",
                    404
                )  #    
            else:
                filename = document.sanitized_name
        if not folder:
            return errorResponse(
                "Folder does not exist.",
                404
            )  #
        else:
            # answer, img_path, html, citation  = answer_query(query, filename, user, folder.sanitized_name,doc_id)
            # Initialize HybridSearch instance
            hs = HybridSearch(collection_name, user.user_code, folder.sanitized_name)
            
            # Retrieve the UUID of the selected PDF
            document_uuid = hs.get_document_uuid(filename)
            
            # Perform a hybrid search to get relevant chunks
            search_results = hs.search(query, document_uuid, alpha=hybrid_weight,top_k=top_k)
            img_path=None
            html = None
            chunk_data = []
            if(search_results):
                citation = (int(json.loads(search_results[0]['metadata'])['page_number']) + 1, search_results[0]['content'])
                # Combine the content from the search results
                for result in search_results[:1]:
                    metadata = json.loads(result['metadata'])
                    if metadata['type']=='image' and img_path==None:
                        img_path = metadata['path']
                    if metadata['type']=='table' and html==None:
                        html = metadata['html']
                        html = re.sub(r"\\t(?![hd])", " ", html)
                passage = "\n".join(result["content"] for result in search_results)

                for result in search_results:
                    # If the reranker score is less than 0, skip the result
                    # if result['reranker_score'] < 0:
                    #     continue
                    result['content'] = result['content'].replace("\n", " ")
                    result['content'] = result['content'].replace("\r", " ")
                    result['content'] = result['content'].replace("\t", " ")
                    result['content'] = result['content'].replace("  ", " ")
                    
                    chunk_data.append({
                        "page_num": int(json.loads(result['metadata'])['page_number']) + 1,
                        "content": result['content']
                    })
            else:
                citation = ""
                passage = ""
            # page_num, content = citation

            # return successReponse(
            #     "",
            #     {"page_num":page_num,"content":passage,"chunk_data":chunk_data}
            # )
            return successReponse(
                "",
                {"chunk_data":chunk_data}
            )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Something went wrong: {str(e)}")
    

@router.post("/answer")
async def search_answer(
    query: str = Form(...), 
    doc_id: Optional[str] = Form(None), 
    folder_id: int = Form(...), 
    hybrid_weight: float = Form(0.5),
    llm: Optional[str] = Form(None),
    prompt: Optional[str] = Form(os.getenv('OPENAI_MODEL')),
    temperature: Optional[int] = Form(1),
    top_p: Optional[int] = Form(1),
    top_k: Optional[int] = Form(15),
    limits: Optional[str] = Form(None),
    email: Optional[str] = Form(None),
    metadata: Optional[str] = Form(None),
    query_response: Optional[str] = Form(None),
    credentials: HTTPAuthorizationCredentials = Depends(security)
):
    db = dbase.SessionLocal()
    token = credentials.credentials
    user = auth.get_current_user(db, token)
    prompt = prompt.strip()
    
    if 0 <= temperature > 2:
        temperature = 1
    if 0 <= top_p > 2:
        top_p = 1
    
    try:
        folder = db.query(Folder).filter(Folder.user_id == user.user_id).filter(Folder.folder_id == folder_id).first()
        if doc_id:
            if metadata:
                metadata = json.loads(metadata)
                json_filter_str = str(metadata).replace("'", '"')
                document = db.query(Documents).filter(
                    Documents.doc_id == doc_id,
                    text(f"JSON_CONTAINS(doc_metadata, '{json_filter_str}')")
                ).first()
            else:
                document = db.query(Documents).filter(Documents.doc_id == doc_id).first()
            if not document:
                return errorResponse(
                    "Document not found.",
                    404
                )  #    
            else:
                filename = document.sanitized_name
        else:
            filename = ''
        if not folder:
            return errorResponse(
                "Folder does not exist.",
                404
            )  #
        else:
            if email:
                chat_user = db.query(User).filter(User.email == email).filter(User.org_id == user.org_id).first()
                
                if not chat_user:
                    return errorResponse(
                        "User does not exist.",
                        404
                    )  #
                else:  
                    answer, img_path, html, citation, passage  = answer_query(query, filename, user, folder.sanitized_name,doc_id,llm,prompt,temperature,top_p,top_k,hybrid_weight,chat_user.user_id,query_response)
                    page_num, content = citation
                    return successReponse(
                        "",
                        {"full_result": answer,"page_num":page_num,"content":passage}
                    )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Something went wrong: {str(e)}")

# Function to get an answer from the selected PDF
def answer_query(query, pdf_name, user, folder_name,doc_id,llm,prompt,temperature,top_p,top_k,hybrid_weight,chat_user_id,query_responses):
    # Initialize HybridSearch instance
    hs = HybridSearch(collection_name, user.user_code, folder_name)
    
    # Retrieve the UUID of the selected PDF
    document_uuid = hs.get_document_uuid(pdf_name)
    
    # Perform a hybrid search to get relevant chunks
    search_results = hs.search(query, document_uuid, alpha=hybrid_weight,top_k=top_k)
    # print(search_results)
    img_path=None
    html = None
    if(search_results):
        citation = (int(json.loads(search_results[0]['metadata'])['page_number']) + 1, search_results[0]['content'])
        # Combine the content from the search results
        for result in search_results[:1]:
            metadata = json.loads(result['metadata'])
            if metadata['type']=='image' and img_path==None:
                img_path = metadata['path']
            if metadata['type']=='table' and html==None:
                html = metadata['html']
                html = re.sub(r"\\t(?![hd])", " ", html)
        passage = "\n".join(result["content"] for result in search_results)
    else:
        citation = ""
        passage = ""
    
    # Construct the prompt for OpenAI
    if not prompt:
        prompt = f"""
        You are a bot-assistant that answers the above question. When answering this question, use the following rules:
        - always answer in English language;
        - use ONLY the information from the sources below;
        - answer briefly in just a few sentences, strictly in accordance with the sources, and do not make any assumptions;
        - reference the source if you use it in the answer, e.g. [#1] or [#2][#4];
        - if there is no information on the question in the sources: say that you can't find the answer and ask the user to try to reformulate the question.
        """

    query_prompt =  f"""The question is {query}?==={prompt}===
    The sources are {passage}
    And again, the question is {query}?
    """

    gpt_model = os.getenv('OPENAI_MODEL')
    # Get the answer from OpenAI
    message_data = [
        {"role": "user", "content": query_prompt}
    ]
    if query_responses and isinstance(query_responses, str):
        query_responses = json.loads(query_responses)
    
    if query_responses:
        message_data.extend(query_responses)
    
    try:
        response = openai_client.chat.completions.create(
            model=llm,
            messages=message_data
        )
        chat_response = response.choices[0].message.content
    except Exception as e:
        chat_response = f"Error occurred: {e}"
    
    db = dbase.SessionLocal()
    try:
        chat_history = ChatHistory(
                user_id=chat_user_id,
                doc_id=doc_id,
                document_uuid=document_uuid,
                query=query, 
                search_results=passage,
                citation=int(json.loads(search_results[0]['metadata'])['page_number']) + 1,
                query_prompt=query_prompt,
                query_response=chat_response
            )
        db.add(chat_history)
        db.commit()
        db.refresh(chat_history)
    except Exception as e:
        db.rollback()
        print(f"Error occurred: {e}")

    return chat_response, img_path, html, citation, passage