from langchain.schema import HumanMessage
from output_parser import attachment_parser, bigfive_parser, personality_parser
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import FAISS
from llm_loader import load_model
from config import openai_api_key
from langchain.chains import RetrievalQA
import os
import json

embedding_model = OpenAIEmbeddings(openai_api_key=openai_api_key)

knowledge_files = {
    "attachments": "knowledge/bartholomew_attachments_definitions.txt",
    "bigfive": "knowledge/bigfive_definitions.txt",
    "personalities": "knowledge/personalities_definitions.txt"
}

documents = []
for key, file_path in knowledge_files.items():
    with open(file_path, 'r', encoding='utf-8') as file:
        content = file.read().strip()
        documents.append(content)

faiss_index = FAISS.from_texts(documents, embedding_model)

llm = load_model(openai_api_key)

qa_chain = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=faiss_index.as_retriever())

def load_text(file_path: str) -> str:
    with open(file_path, 'r', encoding='utf-8') as file:
        return file.read().strip()

def truncate_text(text: str, max_tokens: int = 10000) -> str:
    words = text.split()
    if len(words) > max_tokens:
        return ' '.join(words[:max_tokens])
    return text

def process_input(input_text: str, llm):
    general_task = load_text("tasks/general_task.txt")
    attachments_task = load_text("tasks/Attachments_task.txt")
    bigfive_task = load_text("tasks/BigFive_task.txt")
    personalities_task = load_text("tasks/Personalities_task.txt")

    truncated_input = truncate_text(input_text)

    relevant_docs = qa_chain.invoke({"query": truncated_input})
    
    if isinstance(relevant_docs, dict) and 'result' in relevant_docs:
        retrieved_knowledge = relevant_docs['result']
    else:
        retrieved_knowledge = str(relevant_docs)

    prompt = f"""{general_task}

Attachment Styles Task:
{attachments_task}

Big Five Traits Task:
{bigfive_task}

Personality Disorders Task:
{personalities_task}

Retrieved Knowledge: {retrieved_knowledge}

Input: {truncated_input}

Please provide a comprehensive analysis for each speaker, including:
1. Attachment styles (use the format from the Attachment Styles Task)
2. Big Five traits (use the format from the Big Five Traits Task)
3. Personality disorders (use the format from the Personality Disorders Task)

Respond with a JSON object containing an array of speaker analyses under the key 'speaker_analyses'. Each speaker analysis should include all three aspects mentioned above.

Analysis:"""

    messages = [HumanMessage(content=prompt)]
    response = llm.invoke(messages)
    
    print("Raw LLM Model Output:")
    print(response.content)

    try:
        content = response.content
        if content.startswith("```json"):
            content = content.split("```json", 1)[1]
        if content.endswith("```"):
            content = content.rsplit("```", 1)[0]
        
        parsed_json = json.loads(content.strip())
        
        results = {}
        speaker_analyses = parsed_json.get('speaker_analyses', [])
        for i, speaker_analysis in enumerate(speaker_analyses, 1):
            speaker_id = f"Speaker {i}"
            results[speaker_id] = {
                'attachments': attachment_parser.parse_object(speaker_analysis.get('attachment_styles', {})),
                'bigfive': bigfive_parser.parse_object(speaker_analysis.get('big_five_traits', {})),
                'personalities': personality_parser.parse_object(speaker_analysis.get('personality_disorders', {}))
            }
        
        if not results:
            print("Warning: No speaker analyses found in the parsed JSON.")
            return {"Speaker 1": {
                'attachments': attachment_parser.parse_object({}),
                'bigfive': bigfive_parser.parse_object({}),
                'personalities': personality_parser.parse_object({})
            }}
        
        return results
    except Exception as e:
        print(f"Error processing input: {e}")
        return {"Speaker 1": {
            'attachments': attachment_parser.parse_object({}),
            'bigfive': bigfive_parser.parse_object({}),
            'personalities': personality_parser.parse_object({})
        }}