import torch
import torch.serialization        # the module that controls loading behaviour
from supar.config import Config   # import the exact class that's being blocked

from transformers import BertTokenizerFast

# supar was built against transformers <4.30 which didn't have split_special_tokens.
# This patches the attribute in so supar's internal tokenizer calls don't crash.
if not hasattr(BertTokenizerFast, 'split_special_tokens'):
    BertTokenizerFast.split_special_tokens = False
import re
import os
import numpy as np
import requests
import subprocess
from pathlib import Path
import json
import sys
from haystack import Pipeline, component, Document
from haystack.components.builders import PromptBuilder
from haystack.components.embedders import SentenceTransformersTextEmbedder

# CAMeL Tools imports
from camel_tools.morphology.analyzer import Analyzer
from camel_tools.tokenizers.word import simple_word_tokenize

from haystack.components.embedders import SentenceTransformersTextEmbedder, SentenceTransformersDocumentEmbedder
from haystack.components.builders.prompt_builder import PromptBuilder
from haystack.components.generators import HuggingFaceAPIGenerator
from haystack import Document
from haystack_integrations.components.generators.google_genai import (
    GoogleGenAIChatGenerator,
)
from pathlib import Path
from transformers import GPT2Tokenizer
from haystack_integrations.document_stores.qdrant import QdrantDocumentStore
from haystack_integrations.components.retrievers.qdrant import QdrantEmbeddingRetriever
from haystack.utils import Secret
from haystack.document_stores.types import DuplicatePolicy
from bs4 import BeautifulSoup
from getpass import getpass
from haystack import Pipeline
from haystack.components.writers import DocumentWriter
from haystack import component
from haystack_integrations.components.embedders.fastembed import FastembedDocumentEmbedder, FastembedTextEmbedder
from haystack.components.builders import ChatPromptBuilder
from dotenv import load_dotenv
from camel_tools.morphology.database import MorphologyDB
from camel_tools.morphology.analyzer import Analyzer
from camel_tools.tokenizers.word import simple_word_tokenize
from google import genai
from pymongo import MongoClient
import json
import os
from dotenv import load_dotenv

load_dotenv()

connection_uri = os.getenv("MONGODB_URL")
client = MongoClient(connection_uri, tls=True) # temp fix


try: 
    db = client.get_database("Sebayhi_convos")
    collection = db.get_collection("collection_0")

except Exception as e:
    print("An error occured when trying to connect to the cluster: ", e)

docs = list(collection.find({
    "role": {"$in": ["ai-assistant", "user"]}
}))

msg_history = json.dumps(docs, indent=2, default=str)

qdrant_doc_store = QdrantDocumentStore(
    url="https://a7ec1c00-7522-415a-9152-d6fd4ac3479e.eu-central-1-0.aws.cloud.qdrant.io",
    index="Document",
    embedding_dim=768, # based on the embedding model
    recreate_index=False, # enable only to recreate the index and not connect to the existing one
    api_key = Secret.from_token(os.getenv("Qdrant_key")) 
        
)

doc_embedder = SentenceTransformersDocumentEmbedder(model="akhooli/Arabic-SBERT-100K")

gemini_chat = GoogleGenAIChatGenerator(
    model="gemini-3.1-flash-lite", 
    generation_kwargs={
        "max_output_tokens": 800,
        #"temperature": 0.9
    },
    api_key=Secret.from_token(os.getenv("GEMINI_API_KEY")),
)

knowledge_prompt_template = """
You are **Sebayhi (سيبويه)**, an Arabic grammar tutor inspired by the scholarly tradition of the foundational grammarian سيبويه.

Your role is to help users understand **Arabic grammar, morphology, and related linguistic concepts** clearly, accurately, and conversationally.

## Scope & Mode Routing

* This mode handles **Arabic grammar and linguistic explanations**.
* If the user explicitly requests **إعراب, grammatical parsing, morphological analysis, word-by-word analysis, or analysis of a sentence/word**, do not perform the analysis here. Reply exactly:
  **الرجاء اختيار نمط الإعراب**
* If the user asks a general question about a grammatical or morphological concept, you may explain it normally. For example, you may explain what a صيغة مبالغة is, its common patterns, or how a derivative is formed. Only redirect when the user asks to analyze a specific word or sentence.
* If the request is unrelated to Arabic grammar, morphology, or closely related Arabic linguistic concepts, reply exactly:
  **هذا ليس مجالي.**

## Response Style

* Respond directly to what the student asked. Do not use greetings, introductions, filler, or unnecessary preambles.
* Be conversational and responsive rather than sounding like a textbook.
* Give the answer first, then provide the necessary explanation.
* Keep explanations concise by default. Expand when the question requires it or the student asks for more detail.
* Use clear Arabic appropriate for learners.
* Use examples when they genuinely improve understanding.
* Adapt the explanation to the student's question rather than giving a generic lecture.

## Linguistic Accuracy

* Prioritize **accuracy and clarity over confidence or verbosity**.
* Never invent a grammatical rule, linguistic explanation, root, pattern, or example.
* If a grammatical interpretation depends on context, missing diacritics, or has multiple legitimate analyses, explicitly state the ambiguity and explain the relevant possibilities.
* When discussing derivatives, intensive forms, or other morphological concepts, explain the root and the relationship between the morphological pattern and its meaning when relevant.
* Distinguish clearly between established grammatical rules and interpretations that depend on context or scholarly disagreement.

## Interaction Principle

Treat the conversation as an ongoing tutoring session. React to the student's actual question and, when useful, build naturally on concepts discussed earlier.

Do not add information merely to make the response longer. The goal is to help the student **understand Arabic accurately with the least unnecessary cognitive effort**.

Context:
{% for doc in documents %}
---
{{ doc.content }}
{% endfor %}


CONVERSATION HISTORY:
{{ history }}


السؤال: {{ question }}

الإجابة:
"""



TAG_DICT_PATH = "parser_tags.json"  # adjust path if needed

with open(TAG_DICT_PATH, "r", encoding="utf-8") as f:
    TAG_DICT = json.load(f)

print(f"Loaded tag dictionary: {len(TAG_DICT)} categories")



#    # insert new parser code here
PARSER_DIR = Path("camel_parser")
PARSER_SCRIPT = PARSER_DIR / "text_to_conll_cli.py"


def get_camel_parser_output(user_input: str) -> str:
    """
    Runs the CAMeL parser on an Arabic sentence and returns only the CoNLL-U output.
    """

    result = subprocess.run(
        [
            "python",
            str(PARSER_SCRIPT),
            "-f",
            "text",
            "-s",
            user_input,
        ],
        
        capture_output=True,
        text=True,
    )

    if result.returncode != 0:
        raise RuntimeError(
            f"CAMeL parser failed (exit code {result.returncode})\n\n"
            f"STDOUT:\n{result.stdout}\n\n"
            f"STDERR:\n{result.stderr}"
        )

    clean_data = []

    for line in result.stdout.splitlines():
        line = line.strip()
        if line and (line.startswith("#") or line[0].isdigit()):
            clean_data.append(line)

    return "\n".join(clean_data)

def lookup_tag(tag_name: str, tag: str) -> str:
    """Look up a feature code in the loaded TAG_DICT. Returns 'English (Arabic)' or the raw code."""
    cat = TAG_DICT.get(tag_name)
#print(cat)
    if not cat:
        return 'Enter valid tag name'
    val = cat.get("values", {}).get(tag)
    if not val:
        return 'Enter valid tag letter'
    return f"{val['en']} ({val['ar']} {val['description']} )"
    #return val['en']


def parse_token(columns: list) -> dict:
    """
    Takes ONE split token line and returns a flat,
    resolved dict with only meaningful tags.
    """
    DROP_KEYS = ['token_type']
    token_id, form, lemma, upos, xpos, tags_str, root, deprel, *others = columns
    raw_features = {}
    for pair in tags_str.split('|'):
        if '=' in pair:
            k, v = pair.split('=', 1)
            raw_features[k] = v
    meaningful = {
        k: v for k, v in raw_features.items()
        if v not in ['0', 'na', ''] and k not in DROP_KEYS
    }
    resolved_tags = {
        k: lookup_tag(k, v)
        for k, v in meaningful.items()
    }
    return {
        "id": token_id,
        "word": form,
        "upos": upos,
        "root": root,
        "tags analysis": resolved_tags,
    }

def get_parsed_lines(parser_output: str) -> list:
        parsed_lines = []
        for line in parser_output.strip().split('\n'):
            if line and line[0].isdigit():
                parsed_lines.append(line)
        return [parse_token(line.split('\t')) for line in parsed_lines]


def i3rab_final_output(final_tags_meaning: list) -> str:
    """
    Takes the resolved token list (output of get_parsed_lines)
    and asks Gemini to produce a full إعراب — in Arabic AND English —
    for each word in the sentence.
    """
    if not final_tags_meaning:
        return "خطأ: لم يتم العثور على مخرجات صالحة من المحلل."
    try:
        api_key = os.getenv("GEMINI_API_KEY")
        client = genai.Client(api_key=api_key)
    except Exception as e:
        return f"خطأ في المصادقة: {str(e)}"
    
    
    #final_tags_meaning = get_parsed_lines(parser_output)


    prompt = f"""You are Sebayhi, an Arabic grammar teacher specializing in i'rab (الإعراب). Your role is to guide the student to understand the i'rab, not just give them the answer.

Parsed sentence data:
{final_tags_meaning}

Teaching approach:
- Analyze all words in the sentence first, walking through each one with clear reasoning.
- After completing the full analysis, ask the student one thoughtful question about a grammatically interesting point in the sentence — something that tests understanding, not just memorization.
- If the student answers correctly, affirm it briefly and invite them to ask about another word or sentence.
- If the student answers incorrectly or seems confused, rephrase the explanation using a simpler example or analogy, then give them another chance.


Analysis style — follow this exactly, no labels or headers:
كبُر: فعل ماضٍ مبني على الفتح، والتاء: تاء التأنيث مبنية على السكون.
كلمة: تمييز منصوب بالفتحة.
من أفواههم: من حرف جر، وأفواه: مجرورة بـ(من) وعلامة جرها الكسرة، وهي مضاف، وهم: مضاف إليه.

Rules:
- One word or clitic group per line, colon after the word, then the analysis in plain Arabic grammatical prose.
- Clitics (ال، prepositions, attached pronouns) are analyzed together with their host word on the same line.
- Verbs: state مبني/معرب and what triggered it.
- Nouns: state case, its marker, and what assigned that case.
- If a word has a notable grammatical subtlety (hidden subject, elided verb, تعجب construction), mention it and explain why.
- No raw tags or parser values (asp=p, cas=n, NOUN, etc.) anywhere in the output.
- No greetings, intros, or closing remarks.

If the student asks a general knowledge question about Arabic grammar rather than asking about a specific sentence, direct them to the Knowledge section in the navigator sidebar.


"""


    # if sentence_text:
    #     prompt = f"Full sentence: {sentence_text}\n\n" + prompt

    try:
        response = client.models.generate_content(
            model='gemini-3.1-flash-lite',
            contents=prompt
        )
        lines = [line.strip() for line in response.text.splitlines() if line.strip()]
        return "\n".join(lines)
    
    except Exception as e:
        return f"خطأ أثناء توليد الإعراب: {str(e)}"


@component
class DebugPromptComponent:
    """
    Passthrough component for debugging.
    Prints the prompt it receives, then passes it to the next component.

    Input:  prompt (str)
    Output: prompt (str)  — identical, unchanged
    """

    @component.output_types(prompt=str)
    def run(self, prompt: str):
        print("\n" + "═" * 70)
        print("DEBUG — Prompt sent to LLM:")
        print("═" * 70)
        print(prompt)
        print("═" * 70 + "\n")
        #print(msg_history)
        return {"prompt": prompt}


'''

Conversation handling:
- If the user sends a greeting, respond warmly and briefly, then ask them their level before anything else: (طالب ابتدائي، طالب ثانوي، طالب جامعي، أو متعلم متقدم).
- Do not answer any grammar question until the user has stated their level.
- Once the user states their level, acknowledge it briefly and then answer their question adapted to that level.

Guidelines:

- When explaining a grammatical rule, provide examples suited to the user's level.

- Adapt vocabulary and depth to the student's level — simpler language for younger students, full terminology for advanced learners.
'''



txt_embedder = SentenceTransformersTextEmbedder(model="akhooli/Arabic-SBERT-100K")

knowledge_pipeline = Pipeline()

knowledge_pipeline.add_component("text_embedder",   txt_embedder)
knowledge_pipeline.add_component("retriever",        QdrantEmbeddingRetriever(document_store=qdrant_doc_store))
knowledge_pipeline.add_component("prompt_builder",   PromptBuilder(template=knowledge_prompt_template))
knowledge_pipeline.add_component("debug",            DebugPromptComponent())
knowledge_pipeline.add_component("llm",              gemini_chat)

knowledge_pipeline.connect("text_embedder.embedding", "retriever.query_embedding")
knowledge_pipeline.connect("retriever",               "prompt_builder.documents")
knowledge_pipeline.connect("prompt_builder.prompt",   "debug.prompt")
knowledge_pipeline.connect("debug.prompt",            "llm.messages")