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, get_details_mantra_json, iast_process
from llama_index.core.tools.tool_spec.base import BaseToolSpec
from database import execute_query
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_v3.csv"
PADA_CSV_PATH = "Data/term_data_processed_v2.csv"

class ScriptureDescriptionToolSpec(BaseToolSpec):
    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, level_0, level_1=None, level_2=None, level_3=None):
        try:
            conditions = (self.df['scripture_name'].str.lower() == level_0.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):
    '''
    To obtain the vedamantra details such as vedamantra, padapatha, devata, chandah, rishi etc of vedamantras (or mantras or hyms) from all vedas (RigVeda, AtharvaVeda, SamaVeda, KrishnaYajurVeda, and ShuklaYajurVeda) using the function 
    `get_vedamantra_details`. The mantra summary like anvaya, mantraVishaya, bhavartha/meaning (adhibautic, ahyatmic, adhidaivic), purpose, usage, tippani of vedamantra accessible using the function 'get_vedamantra_summary'
    Sample Query:
    1. What is the vedamantra of the mantra from Rigveda, first mandala, first shukta, and first mantra?
    2. What is the devata of the vedamantra from Rigveda, first mandala, first shukta, and first mantra?
    3. What is the meaning of the vedamantra from Rigveda, first mandala, first shukta, and first mantra written by Tulsi Ram?
    4. What is the (adhibautic) meaning of the vedamantra from RigVeda, first mandala, first shukta, and first mantra?
    5. What is the mantraVishaya of the vedamantra from RigVeda, first mandala, first shukta, and first mantra?
    '''
    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(self, query):
        try:
            details = get_details_mantra_json(query)
            return details['mantraHeader']['language'][1]
        except Exception as e:
            raise ValueError(f"Failed to get mantra details: {e}")

    def _query_db(self, conditions):
        try:
            result = self.df_vedamantra[conditions]['mantra_number'].values
            if len(result) == 0:
                raise ValueError("Mantra not found.")
            return result[0]
        except Exception as e:
            raise ValueError("Failed to query database.")

    def _get_query_conditions(self, scripture_name, **kwargs):
        conditions = (self.df_vedamantra['scripture_name'].str.lower() == scripture_name.lower())
        for key, value in kwargs.items():
            conditions &= (self.df_vedamantra[key] == value)
        return conditions

    def _get_mantra_id(self, scripture_name, **kwargs):
        conditions = self._get_query_conditions(scripture_name, **kwargs)
        return self._query_db(conditions)

    def get_vedamantra_details(self, mantraid=None, scripture_name=None, **kwargs):
        try:
            if mantraid:
                query = f"SELECT mantra_json FROM veda_content WHERE mantra_number = '{mantraid}'"
            else:
                mantra_id = self._get_mantra_id(scripture_name, **kwargs)
                query = f"SELECT mantra_json FROM veda_content WHERE mantra_number = '{mantra_id}'"
            return self._get_mantra_details(query)
        except Exception as e:
            return {"error": str(e)}

    def get_vedamantra_summary(self, mantraid=None, scripture_name=None, **kwargs):
        try:
            if mantraid:
                query = f"SELECT mantra_json FROM veda_content WHERE mantra_number = '{mantraid}'"
            else:
                mantra_id = self._get_mantra_id(scripture_name, **kwargs)
                query = f"SELECT mantra_json FROM veda_content WHERE mantra_number = '{mantra_id}'"
            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_adyatmic_meaning_of_pada' used to get the adibhautic, adidaivic and sdyatmic meaning of a word based on context information.\
    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 apratidhṛṣṭa-śavasam?
    3. Whats the adidaivic meaning of the word apratidhṛṣṭa-śavasam?
    4. What is the adyatmic meaning of the word apratidhṛṣṭa-śavasam?
    '''
    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, **kwargs):
        pada = iast_process(pada)
        try:
            condition = (self.df_terms['Pada'] == pada)
            if scripture_name:
                condition &= (self.df_terms['scripture_name'].str.lower() == scripture_name.lower())
            for key, value in kwargs.items():
                if value is not None:
                    condition &= (self.df_terms[key] == value)
            filtered_df = self.df_terms[condition]
            return filtered_df if not filtered_df.empty else None
        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}'"
            print(query)
            jsonDict = get_details_mantra_json(query)
            print(jsonDict)
            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 {"error": "Required meaning associated with vedamantra is not available."}
        except Exception as e:
            logging.error(f"Error in _get_vedamantra_meaning: {e}")
            return {"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 {"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, **kwargs):
        pada = iast_process(pada)
        try:
            if mantraid:
                details = self.df_terms[(self.df_terms['mantra_id'] == mantraid) & (self.df_terms['Pada'] == pada)]
                
            else:
                if scripture_name is not None:
                    details = self._get_pada_details_by_scripture(pada, scripture_name, **kwargs)
        
            if not details.empty:
                pada_details = details.iloc[0]
                #print(pada_details)
                mantraID = pada_details['mantra_id']
                meanings = self._get_vedamantra_meaning(mantraID,MahatmaName=kwargs.get('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 {
                    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 {"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 {"error": f"Failed to get meaning of the word {pada}. {e}"}