import json
import chardet
import pandas as pd
import streamlit as st
import pymysql
import ast
import re
from utils import word_sentence_similarity, get_list_meaning_word, iast_process
from llama_index.core.tools.tool_spec.base import BaseToolSpec
from database import execute_query, get_details_mantra_json
import pandas as pd
import json
import ast
import logging


# Constants
SCRIPTURE_DESCRIPTIONS_CSV_PATH = "Data/scripture_descriptions.csv"
VEDAMANTRA_CSV_PATH = "Data/veda_content_modified_v5.csv"
PADA_CSV_PATH = "Data/term_data_processed_v2.csv"

class ScriptureDescriptionToolSpec(BaseToolSpec):
    '''
    Purpose: Obtains the description or summary about vedas, mandalas, kandas, shuktas, archakah, adhyaya, and other scriptural elements.
    Returns: A dictionary containing the description or basic information about the specified scriptural element.
    Structure of the levels in each Vedas:
    RigVeda->Mandala->Shukta
    ShuklaYajurVeda->Adhyaya
    SamaVeda->Archika->Shukta
    Atharvaveda->Kandah->Shukta
    Sample query:
    1. Describe the first kandah, second shukta from Atharvaveda? 
       scripture_name: Atharvaveda, level_0: 1, level_1:2
    2. Summarize ShuklaYajurVeda?
    3. What is the difference between ShuklaYajurVeda and KrishnaYajurVeda?
    '''
    spec_functions = ["get_description"]

    def __init__(self):
        super().__init__()
        with open(SCRIPTURE_DESCRIPTIONS_CSV_PATH, 'rb') as f:
            result = chardet.detect(f.read())
        encoding = result['encoding']
        self.df = pd.read_csv(SCRIPTURE_DESCRIPTIONS_CSV_PATH, encoding=encoding)

    def _query_description(self, conditions):
        try:
            result = self.df[conditions]
            if not result.empty:
                return result.iloc[0].to_dict()
            else:
                raise IndexError("Scripture description not found.")
        except IndexError as e:
            raise ValueError(f"Failed to get scripture description: {e}")

    def get_description(self, scripture_name: str, level_1: int=None, level_2: int=None, level_3: int=None):
        try:
            conditions = (self.df['scripture_name'].str.lower() == scripture_name.lower())
            if level_3 is not None:
                conditions &= (self.df['level_1'] == str(level_1)) & (self.df['level_2'] == str(level_2)) & (self.df['level_3'] == str(level_3))
            elif level_2 is not None:
                conditions &= (self.df['level_1'] == str(level_1)) & (self.df['level_2'] == str(level_2))
            elif level_1 is not None:
                conditions &= (self.df['level_1'] == str(level_1))
            return self._query_description(conditions)
        except ValueError as e:
            return {"error": str(e)}

class MantraToolSpec(BaseToolSpec):
    '''
    Use the function `get_vedamantra_details` to retrieve detailed information about Vedic mantras, including vedamantra, padapatha, devata, chandah, 
    and rishi, from all Vedas (RigVeda, AtharvaVeda, SamaVeda, KrishnaYajurVeda, and ShuklaYajurVeda). 
    Use the function `get_vedamantra_summary` to access the information such as anvaya of the mantra, mantraVishaya of the mantra, 
    adhibautic (or adhyatmic or adhidyvic) meaning (or bhavarth) of the mantra, purpose of the mantra, usage of the mantra, and tippani of the mantra.
    Sample Questions:
    1. Obtain the vedamantra of the mantra whose id is 1.1.1.1?
    2. Retrieve the devata of the vedamantra from Rigveda, first mandala, first shukta, and first mantra.
    3. Provide the meaning of the vedamantra from Rigveda, first mandala, first shukta, and first mantra written by Tulsi Ram.
    4. Explain the adhibautic meaning of the first mantra from RigVeda, first mandala, and first shukta.
    5. Identify the mantraVishaya of the vedamantra from RigVeda, first mandala, first shukta, and first mantra.
    6. What is the adibhautic meaning of the mantra 1.1.1.9?
    7. What is the adhyatmic meaning of the mantra 1.1.1.7?
    8. What is the adhidyic meaning of the 6th mantra from RigVeda, first mandala, and first shukta?
    '''
    spec_functions = ["get_vedamantra_details", "get_vedamantra_summary"]

    def __init__(self):
        super().__init__()
        self.df_vedamantra = pd.read_csv(VEDAMANTRA_CSV_PATH, encoding='utf-8')

    def _get_mantra_details_by_scripture(self, scripture_name=None, KandahNumber=None,MandalaNumber=None, ArchikahNumber=None, 
                                                           ShuktaNumber=None, PrapatakNumber=None, MantraNumber=None,
                                                           AnuvakNumber=None, AdhyayaNumber=None):
        try:
            condition = True
            if scripture_name:
                condition = (self.df_vedamantra['scripture_name'].str.lower() == scripture_name.lower())
            if KandahNumber:
                condition &= (self.df_vedamantra['KandahNumber'] == KandahNumber)
            if MandalaNumber:
                condition &= (self.df_vedamantra['MandalaNumber'] == MandalaNumber)
            if ArchikahNumber:
                condition &= (self.df_vedamantra['ArchikahNumber'] == ArchikahNumber)
            if ShuktaNumber:
                condition &= (self.df_vedamantra['ShuktaNumber'] == ShuktaNumber)
            if PrapatakNumber:
                condition &= (self.df_vedamantra['PrapatakNumber'] == PrapatakNumber)
            if MantraNumber:
                condition &= (self.df_vedamantra['MantraNumber'] == MantraNumber)
            if AnuvakNumber:
                condition &= (self.df_vedamantra['AnuvakNumber'] == AnuvakNumber)
            if AdhyayaNumber:
                condition &= (self.df_vedamantra['AdhyayaNumber'] == AdhyayaNumber)
            
            filtered_df = self.df_vedamantra[condition]
            if not filtered_df.empty:
                return filtered_df
            else:
                return None
        except Exception as e:
            logging.error(f"Error in _get_pada_details_by_scripture: {e}")


    def get_vedamantra_details(self, mantraid=None, scripture_name=None, KandahNumber=None,MandalaNumber=None, ArchikahNumber=None, 
                                                           ShuktaNumber=None, PrapatakNumber=None, MantraNumber=None,
                                                           AnuvakNumber=None, AdhyayaNumber=None):
        try:
            if mantraid:
                MantraID = mantraid
            else:
                filter_df = self._get_mantra_details_by_scripture(scripture_name=scripture_name, KandahNumber=KandahNumber,MandalaNumber=MandalaNumber, ArchikahNumber=ArchikahNumber, 
                                                           ShuktaNumber=ShuktaNumber, PrapatakNumber=PrapatakNumber, MantraNumber=MantraNumber,
                                                           AnuvakNumber=AnuvakNumber, AdhyayaNumber=AdhyayaNumber)
                if filter_df is not None:
                    MantraID = filter_df.iloc[0]['mantra_number']
            query = f"SELECT mantra_json FROM veda_content WHERE mantra_number = '{MantraID}'"
            details = get_details_mantra_json(query)
            mantra_details = details['mantraHeader']['language'][1]
            return mantra_details
        except Exception as e:
            return json.dumps({"error": str(e)})

    def get_vedamantra_summary(self, mantraid=None, scripture_name=None, KandahNumber=None,MandalaNumber=None, ArchikahNumber=None, 
                                                           ShuktaNumber=None, PrapatakNumber=None, MantraNumber=None,
                                                           AnuvakNumber=None, AdhyayaNumber=None):
        '''
        Sample Query:
        1. Obtain the anvaya of the mantra whose id (mantraid) is 1.1.1.1?
        2. Retrieve tha adibhautic meaning of the mantra from RigVeda, first mandala, first shukta, and first mantra.
        3. Provide the adhyatmic meaning of the mantra 1.1.1.9?
        4. What is the tippani of the mantra 1.1.1.7?
        5. What is the adhyatmic meaning of the mantra 1.1.1.7?
        6. What is the mantravishaya of the 6th mantra from RigVeda, first mandala, and first shukta?
        '''
        try:
            if mantraid:
                MantraID = mantraid
            else:
                filtered_df = self._get_mantra_details_by_scripture(scripture_name=scripture_name, KandahNumber=KandahNumber,MandalaNumber=MandalaNumber, ArchikahNumber=ArchikahNumber, 
                                                           ShuktaNumber=ShuktaNumber, PrapatakNumber=PrapatakNumber, MantraNumber=MantraNumber,
                                                           AnuvakNumber=AnuvakNumber, AdhyayaNumber=AdhyayaNumber)   
                if filtered_df is not None:
                    MantraID = filtered_df.iloc[0]['mantra_number']
            query = f"SELECT mantra_json FROM veda_content WHERE mantra_number = '{MantraID}'"
            json_dict = get_details_mantra_json(query)
            mantra_summary = json_dict['mantraSummary']['language']
            summary_dict = {"Roman-IAST summary of vedamantra": json_dict['mantraSummary']['language'][1]}
            for item in mantra_summary:
                if item['languageName'] == 'English':
                    mahatma = item['mahatma']['mahatmaName']
                    summary_dict[f"English summary of vedamantra by {mahatma}"] = item
            return summary_dict
        except Exception as e:
            return {"error": str(e)}


class PadaToolSpec(BaseToolSpec):
    '''
    Purpose: To obtains a complete or meaningful meaning of a word or pada based on context information.
    1. The function 'get_meaning_pada' used to get all the possible meanings of the pada based on the given information.
    2. The function 'get_adibauatic_adidaivic_adhyatmic_meaning_of_pada' used to get the adibhautic, adidaivic and sdyatmic meaning 
        of a word based on the information about the mantra (or vedamantra).
    Use the context to generate a meaningful meaning of the pada in the vedamantra.
    Sample query:
    1. What is the meaning of the word apratidhṛṣṭa-śavasam?
    2. What is the adibauatic meaning of the word agnim in the context of the vedamantra from Rigveda, first mandala, first shukta, and first mantra?
    3. Whats the adidaivic meaning of the word apratidhṛṣṭa-śavasam in the mantra 1.1.1.1?
    4. What is the adhyatmic meaning of the word apratidhṛṣṭa-śavasam in the context of the vedamantra from Rigveda, first mandala, first shukta, and first mantra?
    '''
    spec_functions = ["get_pada_meaning","get_adibauatic_adidaivic_adhyatmic_meaning_of_pada"]
    
    def __init__(self):
        super().__init__()
        self.df_terms = pd.read_csv(PADA_CSV_PATH, dtype={'AnuvakNumber': 'Int64', 'PrapatakNumber': 'Int64', 'KandahNumber': 'Int64', 'ShuktaNumber': 'Int64', 'ArchikahNumber': 'Int64', 'AdhyayaNumber': 'Int64', 'MandalaNumber': 'Int64', 'ParyayaNumber': 'Int64'}, encoding='utf-8')
        self.df_vedic_content = pd.read_csv(VEDAMANTRA_CSV_PATH,encoding = 'utf-8')

    def _get_pada_details_by_scripture(self, pada, scripture_name=None, KandahNumber=None,MandalaNumber=None, ArchikahNumber=None, 
                                       ShuktaNumber=None, PrapatakNumber=None, MantraNumber=None,AnuvakNumber=None, AdhyayaNumber=None):
        try:
            #pada = iast_process(pada)
            condition = (self.df_terms['Pada'] == pada)

            if scripture_name:
                condition &= (self.df_terms['scripture_name'].str.lower() == scripture_name.lower())
            if KandahNumber:
                condition &= (self.df_terms['KandahNumber'] == KandahNumber)
            if MandalaNumber:
                condition &= (self.df_terms['MandalaNumber'] == MandalaNumber)
            if ArchikahNumber:
                condition &= (self.df_terms['ArchikahNumber'] == ArchikahNumber)
            if ShuktaNumber:
                condition &= (self.df_terms['ShuktaNumber'] == ShuktaNumber)
            if PrapatakNumber:
                condition &= (self.df_terms['PrapatakNumber'] == PrapatakNumber)
            if MantraNumber:
                condition &= (self.df_terms['MantraNumber'] == MantraNumber)
            if AnuvakNumber:
                condition &= (self.df_terms['AnuvakNumber'] == AnuvakNumber)
            if AdhyayaNumber:
                condition &= (self.df_terms['AdhyayaNumber'] == AdhyayaNumber)
            
            filtered_df = self.df_terms[condition]
            
            if not filtered_df.empty:
                return filtered_df
            else:
                return None
        
        except KeyError as ke:
            logging.error(f"KeyError in _get_pada_details_by_scripture: {ke}")
        except Exception as e:
            logging.error(f"Error in _get_pada_details_by_scripture: {e}")
        
        return None
        
    def _get_vedamantra_meaning(self, mantraID, MahatmaName=None):
        try:
            query = f"SELECT mantra_json FROM veda_content WHERE mantra_number = '{mantraID}'"
            jsonDict = get_details_mantra_json(query)
            mantraSummary = jsonDict['mantraSummary']['language']
            if MahatmaName is not None:
                filtered_summary = [data_dict for data_dict in mantraSummary if data_dict.get('mahatma', {}).get('mahatmaName') == MahatmaName]
                if filtered_summary:
                    mantraSummary = filtered_summary
            best_meaning = None
            best_count = 0
            for data_dict in mantraSummary:
                if data_dict.get('languageName') == "English":
                    meanings = data_dict['mahatma']['bhavartha']
                    count = sum(bool(meanings.get(cat, None)) for cat in ['adibhautic', 'adidaivic', 'adhyatmic'])
                    if count >= best_count:
                        best_meaning = {cat: meanings.get(cat, None) for cat in ['adibhautic', 'adidaivic', 'adhyatmic']}
                        best_count = count
            return best_meaning if best_meaning else json.dumps({"error": "Required meaning associated with vedamantra is not available."})
        except Exception as e:
            logging.error(f"Error in _get_vedamantra_meaning: {e}")
            return json.dumps({"error": f"An error occurred: {e}"})
    
    def _get_pada_morphology(self, term_details, meanings):
        try:
            morphology_list = ast.literal_eval(term_details['Morphology'])
            term_morph_list = []
            for morphs in morphology_list:
                term_info = {}
                for field in ['stem', 'root']:
                    morph_word = morphs.get(field)
                    if morph_word:
                        meaning = word_sentence_similarity(meanings, morph_word)
                        term_info[f'{field}_word'] = morph_word
                        term_info[f'{field}_meaning'] = meaning[0][0] if meaning else None
                        term_info[f'{field}_score'] = meaning[0][1] if meaning else None
                term_info['grammar'] = morphs['grammar']
                term_morph_list.append(term_info)
            return term_morph_list
        except Exception as e:
            logging.error(f"Error in _get_pada_morphology: {e}")
            return []

    def get_pada_meaning(self, pada):
        #pada=iast_process(pada)
        try:
            pada_details = self.df_terms[self.df_terms['Pada'] == pada]
            meanings_list = []
            for morphs in ast.literal_eval(pada_details['Morphology'].values[0]):
                for field in ['stem', 'root']:
                    word = morphs.get(field)
                    if word:
                        meanings_list.append(get_list_meaning_word(word))
            return meanings_list
        except Exception as e:
            logging.error(f"Error in get_pada_meaning: {e}")
            return json.dumps({"error": f"Required meaning associated with pada is not available. {e}"})    
    

    def get_adibauatic_adidaivic_adhyatmic_meaning_of_pada(self, pada, mantraid=None, scripture_name=None, 
                                                           KandahNumber=None,MandalaNumber=None, ArchikahNumber=None, 
                                                           ShuktaNumber=None, PrapatakNumber=None, MantraNumber=None,
                                                           AnuvakNumber=None, AdhyayaNumber=None,MahatmaName=None):
        '''
        Sample query:
        1. What is the meaning of pada 'agnim' from RigVeda, first mandala, first shukta and first mantra?
        2. What is the adhyatmic meaning of the pada agnim in the context of the mantra whose id is '1.1.1.1?'
        '''
        try:
            #pada = iast_process(pada)
            if mantraid:
                details = self.df_terms[(self.df_terms['mantra_id'] == mantraid) & (self.df_terms['Pada'] == pada)]
            else:
                details = self._get_pada_details_by_scripture(pada, scripture_name=scripture_name, KandahNumber=KandahNumber,MandalaNumber=MandalaNumber, ArchikahNumber=ArchikahNumber, 
                                                           ShuktaNumber=ShuktaNumber, PrapatakNumber=PrapatakNumber, MantraNumber=MantraNumber,
                                                           AnuvakNumber=AnuvakNumber, AdhyayaNumber=AdhyayaNumber)
        
            if not details.empty:
                pada_details = details.iloc[0]
                print(pada_details)
                mantraID = pada_details['mantra_id']
                meanings = self._get_vedamantra_meaning(mantraID,MahatmaName=MahatmaName)
                if 'error' in meanings:
                    return meanings
                ab_term_morph_list = self._get_pada_morphology(pada_details, meanings['adibhautic'])
                ad_term_morph_list = self._get_pada_morphology(pada_details, meanings['adidaivic'])
                at_term_morph_list = self._get_pada_morphology(pada_details, meanings['adhyatmic'])
                return json.dumps({
                    f'adibhautic_info_{pada}': ab_term_morph_list,
                    'vedamantra_adibhautic_meaning': meanings['adibhautic'],
                    f'adidavic_info_{pada}': ad_term_morph_list,
                    'vedamantra_adidavic_meaning': meanings['adidaivic'],
                    f'adhyatmic_info_{pada}': at_term_morph_list,
                    'vedamantra_adhyatmic_meaning': meanings['adhyatmic']
                })
            else:
                return json.dumps({"error": f"No details found for pada '{pada}'"})
        except Exception as e:
            logging.error(f"Error in get_adibauatic_adidaivic_adhyatmic_meaning_of_pada: {e}")
            return json.dumps({"error": f"Failed to get meaning of the word {pada}. {e}"})