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):
    '''
    To obtain the description of the vedic scriptures, mandalas, kandahas, shukta, Anuvaka etc.
    Sample Query:
    1. Describe RigVeda?
    2. What is the summary of the first mandala from RigVeda?
    3. What is the brief description of the 4th Shukta of the 2nd Kandah from AtharvaVeda? 
    '''
    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):
    '''
    You can retrieve detailed information about Vedic mantras, including details such as vedamantra, padapatha, devata, chandah, 
    and rishi, from all Vedas (RigVeda, AtharvaVeda, SamaVeda, KrishnaYajurVeda, and ShuklaYajurVeda) using the function 
    `get_vedamantra_details`. Additionally, you can access a summary of the mantra, including anvaya, mantraVishaya, and 
    adhibautic (or adhyatmic or adhidyvic) meaning (or bhavartha), purpose, usage, and tippani of the vedamantra with the function 
    `get_vedamantra_summary`.

    Here's a sample query format:
    
    1. Obtain the vedamantra of the mantra 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.
    '''
    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 _get_mantra_details_by_scripture(self, scripture_name=None, **kwargs):
        try:
            if scripture_name:
                condition = (self.df_vedamantra['scripture_name'].str.lower() == scripture_name.lower())
            
            for key, value in kwargs.items():
                condition &= (self.df_vedamantra[key] == value)
            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_mantra_id(self, scripture_name, **kwargs):
        filtered_df = self._get_mantra_details_by_scripture(scripture_name, **kwargs)
        mantraID = filtered_df['mantra_number']
        return mantraID.values[0]

    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)
                #print(mantra_id)
                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)
                print(mantra_id)
                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_adhyatmic_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 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?
    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):
        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())
            
            for key, value in kwargs.items():
                condition &= (self.df_terms[key] == value)
                
            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 {"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, 
                                                           KandahNumber=None,MandalaNumber=None, ArchikahNumber=None, 
                                                           ShuktaNumber=None, PrapatakNumber=None, MantraNumber=None,
                                                           AnuvakNumber=None, AdhyayaNumber=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}"}