import os
import re
import json
import torch
import pandas as pd
from sentence_transformers import SentenceTransformer, util
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import os, json, datetime

def load_training_data(data_folder):
    data = []
    missing_code = 0
    missing_query = 0
    for file in os.listdir(data_folder):
        if file.endswith(".json"):
            path = os.path.join(data_folder, file)
            with open(path, "r", encoding="utf-8") as f:
                try:
                    entries = json.load(f)
                    for e in entries:
                        query_text = e.get("english") or e.get("English") or e.get("query") or e.get("Query")
                        if not query_text:
                            missing_query += 1
                            continue

                        code_text = e.get("pandas_code") or e.get("Pandas_Code") or e.get("code")
                        if not code_text:
                            missing_code += 1
                            continue

                        data.append({"english": query_text, "pandas_code": code_text})
                except Exception as e:
                    print(f" Skipped {file}: {e}")
    print(f" Loaded {len(data)} valid query–code pairs from {data_folder}")
    print(f" Skipped {missing_code} missing-code and {missing_query} missing-query entries.")
    return data

class EnhancedRetriever:
    def __init__(self, data):
        self.model = SentenceTransformer("all-MiniLM-L6-v2")
        valid_data = [item for item in data if (item.get("pandas_code") and (item.get("english") or item.get("query")))]
        if not valid_data:
            raise ValueError("No valid query–code pairs found in dataset!")
        self.queries = [item.get("english") or item.get("query") for item in valid_data]
        self.codes = [item["pandas_code"] for item in valid_data]
        print(f" Using {len(valid_data)} valid items for retrieval.")
        print(" Encoding queries for retrieval...")
        self.query_embeddings = self.model.encode(self.queries, convert_to_tensor=True)

    def retrieve_best_match(self, user_query, top_k=3):
        user_emb = self.model.encode(user_query, convert_to_tensor=True)
        similarity = util.pytorch_cos_sim(user_emb, self.query_embeddings)[0]
        top_results = torch.topk(similarity, k=top_k)
        results = []
        for i in range(top_k):
            results.append({
                "query": self.queries[top_results.indices[i]],
                "pandas_code": self.codes[top_results.indices[i]],
                "similarity": float(top_results.values[i])
            })
        return results



class Generator:
    def __init__(self, model_dir="model"):

        try:
            self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
            self.model = AutoModelForSeq2SeqLM.from_pretrained(model_dir)
           

            print(" Base model loaded successfully.")
        except Exception as e:
            print(f" Error loading base model: {e}")
            
            self.tokenizer = None
            self.model = None

    def generate(self, query):
                                print("\n================= GENERATE DEBUG =================")
                                print("Incoming Query :", query)

                                try:
                                    prompt = query
                                    print("Prompt Created:", prompt)

                                    print("Tokenizer Loaded:", self.tokenizer is not None)
                                    try:
                                        inputs = self.tokenizer(prompt, return_tensors="pt")
                                        
                                    except Exception as e:
                                        print(" TOKENIZER ERROR:", e)
                                        return f"[Tokenizer Error] {str(e)}"

                               
                                    print("Model Loaded:", self.model is not None)
                                    try:
                                        outputs = self.model.generate(**inputs, max_length=128)
                                        print("Model Generation Successful")
                                    except Exception as e:
                                        print(" MODEL ERROR:", e)
                                        return f"[Model Error] {str(e)}"

                                
                                    try:
                                        decoded = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
                                        print("Decoding Successful →", decoded)
                                    except Exception as e:
                                        print(" DECODING ERROR:", e)
                                        return f"[Decode Error] {str(e)}"

                                   
                                    return decoded

                                except Exception as e:
                                    print(" UNEXPECTED ERROR:", e)
                                    return f"[Unexpected Error] {str(e)}"

def normalize_name(name):
    if not isinstance(name, str):
        return name
    return re.sub(r'[^a-z0-9]', '', name.lower())

def extract_column_names(text):
    words = re.findall(r'\b[a-zA-Z_][a-zA-Z0-9_]*\b', text)
    stopwords = {
        'show','display','find','get','the','and','or','where','what','how','many','much','list','give','me','all','with','for','bottom','top','average','mean','sum','median','count','minimum','maximum','highest','lowest'
    }
    cols = [w for w in words if w.lower() not in stopwords and len(w) > 2]
    return [normalize_name(c) for c in cols]

def extract_values(text):
    quoted = re.findall(r"'([^']*)'", text)
    numbers = re.findall(r'\b\d+\b', text)
    return quoted + numbers

def enhanced_adaptation(user_query, code, original_retrieved_query):
    query_columns = extract_column_names(user_query)
    original_columns = extract_column_names(original_retrieved_query)
    query_values = extract_values(user_query)
    original_values = extract_values(original_retrieved_query)
    new_code = code
    for orig_col, new_col in zip(original_columns, query_columns):
        if orig_col and new_col and normalize_name(orig_col) != normalize_name(new_col):
            for pattern in [rf"'{orig_col}'", rf'\"{orig_col}\"', rf"\\b{orig_col}\\b"]:
                new_code = re.sub(pattern, new_col, new_code, flags=re.IGNORECASE)
    for orig_val, new_val in zip(original_values, query_values):
        if orig_val and new_val and orig_val != new_val:
            new_code = re.sub(rf"'{re.escape(orig_val)}'", f"'{new_val}'", new_code)
            new_code = re.sub(rf'\"{re.escape(orig_val)}\"', f'\"{new_val}\"', new_code)
            new_code = re.sub(rf"\\b{re.escape(orig_val)}\\b", new_val, new_code)
    return new_code

def adapt_operations_based_on_query(query, code):
    q = query.lower()
    c = code
    if any(word in q for word in ["average", "mean", "avg"]):
        c = re.sub(r"\.(sum|min|max|count)\(\)", ".mean()", c)
    elif any(word in q for word in ["total", "sum", "add", "together"]):
        c = re.sub(r"\.(mean|min|max|count)\(\)", ".sum()", c)
    return c

def select_best_template(retrieved_results, user_query):
    best_score = -1
    best_result = retrieved_results[0]
    for result in retrieved_results:
        score = result["similarity"]
        if score > best_score:
            best_score = score
            best_result = result
    return best_result

def post_process_code(code, user_query):
    code = re.sub(r'\.\.', '.', code)
    return code

def detect_vague_query(user_query):
    q = user_query.lower()
    intent_keywords = [
        "amount", "sum", "total", "average", "mean", "minimum", "maximum",
        "count", "number", "list", "show", "display", "details", "records","how many","what is"
    ]
    data_keywords = [
        "claim", "bill", "policy", "approved", "rejected", "tariff", "package", "department", "provider"
    ]
    has_intent = any(word in q for word in intent_keywords)
    has_data = any(word in q for word in data_keywords)
    if has_data and not has_intent:
        return True
    return False

class RobustHybridText2Code:
    def __init__(self, data_folder="data", model_dir="model/my_finetuned_model"):
        self.data = load_training_data(data_folder)
        self.retriever = EnhancedRetriever(self.data)
        self.generator = MultiModelGenerator(
    {
        "modelA": "model/my_finetuned_model",
        "modelB": "model/second_model_folder"
    }
)

        
        
        
        # self.generator = Generator(model_dir)
        all_cols = set()
        for item in self.data:
            code = item.get("pandas_code", "")
            cols = re.findall(r"df\[['\"]([^'\"]+)['\"]\]", code)
            all_cols.update(cols)
        self.col_map = {normalize_name(c): c for c in all_cols}

    def get_code(self, user_query):
        retrieved_results = self.retriever.retrieve_best_match(user_query, top_k=3)
        best = select_best_template(retrieved_results, user_query)
        print(f" [Best Match Similarity: {best['similarity']:.2f}]")
        print(f" Original Query: {best['query']}")
        if best["similarity"] > 0.95:
            print(" High similarity — using code directly.")
            code = best["pandas_code"]
        elif best["similarity"] >= 0.90:
            print(" Moderate similarity — adapting code.")
            code = enhanced_adaptation(user_query, best["pandas_code"], best["query"])
        else:
            print(" Low similarity — generating new code.")
            #code = self.generator.generate(user_query)
            
            
            
            generated, used_model = self.generator.generate(user_query)
            print(f"[INFO] Model used for this query → {used_model}")

            code = generated

        return post_process_code(code, user_query)
     

# class MultiDatasetHybridText2Code:
#     def __init__(self):
#         print(" Initializing Multi-Dataset Hybrid System with CSV Execution and Clarification Memory...\n")

#         self.datasets = {
#             "bill": {"folder": "data/bill_data", "csv_path": "data/bill_data/bill_dataset.csv"},
#             "status": {"folder": "data/status_data", "csv_path": "data/status_data/dataset.csv"},
#             "history": {"folder": "data/history_data", "csv_path": "data/history_data/status_history_dataset.csv"}  
#         }

#         self.models = {}
#         for name, meta in self.datasets.items():
#             print(f" Loading dataset: {name}")
#             self.models[name] = RobustHybridText2Code(
#                 data_folder=meta["folder"],
#                 model_dir="model/my_finetuned_model" )
#             if os.path.exists(meta["csv_path"]):
#                 df = pd.read_csv(meta["csv_path"])
#                 self.datasets[name]["df"] = df
#                 print(f" Loaded {len(df)} records for {name} dataset.")
#             else:
#                 print(f" CSV not found for {name}: {meta['csv_path']}")

#         print("\n Models and CSVs loaded successfully!\n")

#         self.routing_model = SentenceTransformer("all-MiniLM-L6-v2")
#         self.dataset_embeddings = {}
#         for name in self.models:
#             all_queries = [e.get("english") or e.get("query") for e in self.models[name].data]
#             mean_emb = self.routing_model.encode(all_queries, convert_to_tensor=True).mean(dim=0)
#             self.dataset_embeddings[name] = mean_emb


#     def detect_dataset(self, user_query):
#         q = user_query.lower().strip()

    
#         if any(word in q for word in ["bill", "billing", "bill details", "bill status","invoice"]):
#             print(" Rule-based routing: 'bill' detected → Using BILL dataset.")
#             return "bill"
#         elif any(word in q for word in ["claim history", "history", "previous claims", "old claims"]):
#             print(" Rule-based routing: 'claim history' detected → Using HISTORY dataset.")
#             if "history" not in self.datasets:
#                 print(" History dataset not found — falling back to default routing.")
#             else:
#                 return "history"
#         else:
#             return "status"

        
#         user_emb = self.routing_model.encode(user_query, convert_to_tensor=True)
#         sims = {name: float(util.pytorch_cos_sim(user_emb, emb)) for name, emb in self.dataset_embeddings.items()}
#         best_match = max(sims, key=sims.get)
#         print(f" Dataset routing via embeddings: {sims}")
#         print(f"Fallback selected: {best_match}")
#         return best_match

#     def execute_code(self, code, df):
#         try:
#             local_env = {"df": df, "pd": pd}
#             result = eval(code, {"__builtins__": {}}, local_env)
#             return result
#         except Exception as e:
#             print(f" Error executing code: {e}")
#             return None
        
#     def get_code(self, user_query,api_mode=False):
#         if detect_vague_query(user_query):
#             print(" Please specify what you want — number of claims, total amount, or claim details?")
#             return None
#         dataset_name = self.detect_dataset(user_query)
#         print(f" Using dataset: {dataset_name}")
#         bot = self.models[dataset_name]
#         code = bot.get_code(user_query)
#         df = self.datasets[dataset_name].get("df")
#         if df is not None and code:
#             print(f"\n Suggested Pandas Code:\n{code}")
#             result = self.execute_code(code, df)
#             if result is not None:
#                 print(f"\n Result:\n{result}")

              
#                 log_interaction(user_query, code, result, dataset_name)

                
#                 # if not api_mode:
#                 #         save_user_example(dataset_name, user_query, code,result)
#         return code,result





class MultiModelGenerator:

    def __init__(self, model_dirs):
       
        self.model_dirs = {
            "modelA": "model/modelA_dir",
            "modelB": "model/modelB_dir"
        }
        
        self.models = {}
        self.tokenizers = {}

        for name, path in model_dirs.items():
            try:
                tok = AutoTokenizer.from_pretrained(path)
                try:
                    model = AutoModelForSeq2SeqLM.from_pretrained(path)
                except:
                    from transformers import T5ForConditionalGeneration
                    model = T5ForConditionalGeneration.from_pretrained(path)

                self.models[name] = model
                self.tokenizers[name] = tok
                print(f"[MultiModel] Loaded model: {name} from {path}")

            except Exception as e:
                print(f"[MultiModel] Failed loading {name}: {e}")

        self.router = SentenceTransformer("all-MiniLM-L6-v2")

    
    def pick_best_model(self, query):
        query_l = query.lower()

        #  KEYWORD-BASED ROUTING RULES
        trigger_keywords = [
            "settelement age", 
            "outstanding age", 
            "top",
            "claim age", 
            "rank"
        ]

        for kw in trigger_keywords:
            if kw in query_l:
                print(f"[MultiModel] Keyword '{kw}' detected → Using Model 2")
                return "modelB"   # your second model

        # Otherwise → ALWAYS use model 1
        print("[MultiModel] Default routing → Using Model 1")
        return "modelA"


    def generate(self, query):
        print("\n============= MULTI-MODEL GENERATE =============")

        chosen = self.pick_best_model(query)
        model = self.models[chosen]
        tokenizer = self.tokenizers[chosen]

        inputs = tokenizer(query, return_tensors="pt")
        out = model.generate(**inputs, max_length=128)
        decoded = tokenizer.decode(out[0], skip_special_tokens=True)

        print(f"[MultiModel] Final Answer from {chosen} → {decoded}")
        return decoded, chosen



class ClaimBoticsEngine:
    def __init__(self):
        print("\n=== Initializing Multi-Dataset + Multi-Model Hybrid Engine ===\n")

        # 1️⃣ DATASET SETUP
        self.datasets = {
            "bill": {"folder": "data/bill_data", "csv_path": "data/bill_data/bill_dataset.csv"},
            "status": {"folder": "data/status_data", "csv_path": "data/status_data/dataset.csv"},
            "history": {"folder": "data/history_data", "csv_path": "data/history_data/history_dataset.csv"}
        }

        self.models_by_dataset = {}
        for name, meta in self.datasets.items():
            print(f"Loading dataset: {name}")

            # Load JSON training data for retrieval
            data_pairs = load_training_data(meta["folder"])

            # Load CSV dataframe
            if os.path.exists(meta["csv_path"]):
                df = pd.read_csv(meta["csv_path"])
                self.datasets[name]["df"] = df
                print(f"Loaded {len(df)} rows for dataset: {name}")
            else:
                self.datasets[name]["df"] = None
                print(f"CSV not found: {meta['csv_path']}")

            # Load retriever = trained JSON Q&A
            retriever = EnhancedRetriever(data_pairs)

            # Load two generation LLM models
            generator = MultiModelGenerator({
                "modelA": "model/my_finetuned_model",
                "modelB": "model/my_second_model"
            })

            # Store all dataset tools
            self.models_by_dataset[name] = {
                "retriever": retriever,
                "generator": generator,
                "training_data": data_pairs,
            }

        print("\nEngine fully loaded!\n")

    # 2️⃣ DATASET ROUTING (Bill / Status / History)
    def detect_dataset(self, user_query):
        q = user_query.lower()

        if any(x in q for x in ["bill", "billing", "invoice"]):
            print("Routing → BILL dataset")
            return "bill"

        if any(x in q for x in ["history", "older claims", "previous"]):
            print("Routing → HISTORY dataset")
            return "history"

        print("Routing → STATUS dataset")
        return "status"

    # 3️⃣ CODE EXECUTION
    def execute_code(self, code, df):
        try:
            local_env = {"df": df, "pd": pd}
            result = eval(code, {"__builtins__": {}}, local_env)
            return result
        except Exception as e:
            print("Execution error:", e)
            return None

    # 4️⃣ MAIN PIPELINE
    def process_query(self, user_query):

        # --- dataset routing ---
        dataset = self.detect_dataset(user_query)
        df = self.datasets[dataset]["df"]
        retriever = self.models_by_dataset[dataset]["retriever"]
        generator = self.models_by_dataset[dataset]["generator"]

        # --- retrieve top matches from JSON training files ---
        retrieved = retriever.retrieve_best_match(user_query, top_k=3)
        best = select_best_template(retrieved, user_query)

        print(f"Best similarity = {best['similarity']:.4f}")

        # --- use template if high similarity ---
        if best["similarity"] > 0.95:
            print("High similarity → using template")
            code = best["pandas_code"]

        # --- modify template if moderate similarity ---
        elif best["similarity"] >= 0.90:
            print("Medium similarity → adapting template")
            code = enhanced_adaptation(user_query, best["pandas_code"], best["query"])

        # --- low similarity → LLM calls & multi-model selection ---
        else:
            print("Low similarity → generating via LLM (ModelA / ModelB)")
            code, used_model = generator.generate(user_query)
            print(f"LLM Model Used: {used_model}")

        # Final code cleanup
        code = post_process_code(code, user_query)

        # --- Execute Pandas Code ---
        result = None
        if df is not None:
            result = self.execute_code(code, df)

        # Log interaction
        log_interaction(user_query, code, result, dataset)

        return {
            "dataset": dataset,
            "code": code,
            "result": result
        }


def log_interaction(query, code, result, dataset):
        
        os.makedirs("logs", exist_ok=True)
        log_path = "logs/interaction_log.jsonl"
        entry = {
            "timestamp": datetime.datetime.now().isoformat(),
            "query": query,
            "dataset": dataset,
            "code": code,
            "result_preview": str(result)[:300] 
        }
        with open(log_path, "a", encoding="utf-8") as f:
            f.write(json.dumps(entry) + "\n")

# def save_user_example(dataset, query, code,result):
        
#         os.makedirs(f"logs", exist_ok=True)
#         path = f"logs/new_user_data.jsonl"
#         new_entry = {
#             "timestamp": datetime.datetime.now().isoformat(),
#             "query": query,
#             "dataset": dataset,
#             "code": code,
#             "result_preview": str(result)[:300] 
#         }
#         with open(path, "a", encoding="utf-8") as f:
#             f.write(json.dumps(new_entry) + "\n")
#         print(f" Saved new user example to {path}")




if __name__ == "__main__":
    print("\nInitializing Hybrid Text-to-Pandas Bot...\n")
    bot = ClaimBoticsEngine()

    print("Chatbot ready! Type your queries below. Type 'exit' to quit.\n")

    while True:
        try:
            user_query = input("You: ").strip()
            if user_query.lower() in ["exit", "quit"]:
                print("Goodbye!")
                break
            if not user_query:
                continue

            code, result = bot.get_code(user_query)
            
            print("\nBot:")
            if code:
                print(f"  Suggested Pandas Code:\n  {code}")
            else:
                print("  Could not generate code. Please clarify your query.")
            
            if result is not None:
                print(f"\n  Execution Result:\n  {result}")
            print("\n" + "-"*50 + "\n")
            
        except KeyboardInterrupt:
            print("\nInterrupted. Exiting chatbot...")
            break
        except Exception as e:
            print(f"Error: {e}")
            continue