import os
import re
import gradio as gr
from functools import lru_cache
from camel_tools.ner import NERecognizer
from camel_tools.disambig.bert import BERTUnfactoredDisambiguator
from camel_tools.tokenizers.word import simple_word_tokenize
from camel_tools.utils.normalize import normalize_unicode
from camel_tools.utils.dediac import dediac_ar

# ── Data Download (Skipping shell sub-processes if already exists) ────────
data_path = os.path.expanduser("~/.camel_tools")
os.environ["CAMELTOOLS_DATA"] = data_path

# Fix paths to match how camel_tools unzips them under the hood
ner_path = os.path.join(data_path, "data", "ner", "arabert")
disambig_path = os.path.join(data_path, "data", "disambig_bert_unfactored", "msa")

if not os.path.exists(ner_path):
    print("Downloading CAMEL tools NER data...")
    os.system("camel_data -i ner-arabert")

if not os.path.exists(disambig_path):
    print("Downloading CAMEL tools BERT Disambiguation data...")
    os.system("camel_data -i disambig-bert-unfactored-msa")

# ── Configuration & Regex ──────────────────────────────────────────────────
VERB_POS = {"verb", "verb_pseudo"}
DANGLING_PREFIXES = ('وأضافت', 'وأضاف', 'وتابع', 'وتابعت', 'وأكدت', 'وأشار', 'وأشارت', 'وبين', 'وأوضح', 'وأوضحت')

# Pre-compile sentence splitting Regex
SENTENCE_SPLIT_REGEX = re.compile(r'[.!?؟。]+')

# Layer 2 Split setup - compiled into a single Regex operation for speed
CONNECTORS = sorted([
    'بينما', 'في حين', 'حيث إن', 'حيث أن',
    'كما أعلنت', 'كما أعلن', 'كما أكدت', 'كما أكد',
    'كما أشارت', 'كما أشار', 'كما',
    'وأضافت', 'وأضاف', 'وأكدت', 'وأكد',
    'وأشارت', 'وأشار', 'وأعلنت', 'وأعلن',
    'وأوضح', 'وأوضحت', 'وتابع', 'وتابعت',
    'وصرح', 'وصرحت', 'وأفاد', 'وأفادت',
    'وقال', 'وقالت', 'وأضاف أن', 'وأكد أن',
], key=len, reverse=True)

# Build regex pattern for capturing spaces + connector + space
SPLIT_PATTERN = re.compile(f"( {' | '.join(CONNECTORS)} )")


class ClaimExtractorLogic:
    def __init__(self):
        print("Initializing models (Looking for GPU)...")
        # NERecognizer handles CPU/GPU placement relatively gracefully under the hood if device is not forced.
        self.ner = NERecognizer.pretrained()
        self.disambig = BERTUnfactoredDisambiguator.pretrained(model_name="msa")
        print("Models loaded.")

    @lru_cache(maxsize=128)  # <-- Caching repeated identical texts here
    def run(self, text: str) -> list[str]:
        if not text or not text.strip():
            return []

        text = normalize_unicode(text)
        text = ' '.join(text.split())
        base_sentences = SENTENCE_SPLIT_REGEX.split(text)
        base_sentences = [s.strip() for s in base_sentences if s.strip()]

        # 1. First pass: PREPARE DATA BATCH
        flat_tokens_batch = []
        chunk_meta = [] # Keeps track of original layout structure

        for sent_idx, base_sentence in enumerate(base_sentences):
            chunks = self._layer2_split(base_sentence)
            for chunk in chunks:
                tokens = simple_word_tokenize(chunk)
                if tokens:
                    flat_tokens_batch.append(tokens)
                    chunk_meta.append({
                        "sent_idx": sent_idx,
                        "chunk": chunk,
                        "tokens": tokens
                    })

        if not flat_tokens_batch:
            return []

        # 2. BATCH INFERENCE: Process everything simultaneously through the heavy ML models
        try:
            all_ner_labels = self.ner.predict_sentences(flat_tokens_batch)
            all_analyses = self.disambig.disambiguate_sentences(flat_tokens_batch)
        except AttributeError:
            # Fallback just in case you're using an older version of camel_tools
            # where batch execution names slightly differ
            all_ner_labels = [self.ner.predict_sentence(t) for t in flat_tokens_batch]
            all_analyses = [self.disambig.disambiguate(t) for t in flat_tokens_batch]

        # 3. Second pass: APPLY BUSINESS LOGIC (Sequential rules mapped to ML data)
        claims = []
        seen = set()
        current_sent_idx = -1
        running_context = set()

        for idx in range(len(flat_tokens_batch)):
            meta = chunk_meta[idx]
            
            # Reset context memory if we transitioned to a new base_sentence
            if meta["sent_idx"] != current_sent_idx:
                running_context.clear()
                current_sent_idx = meta["sent_idx"]

            tokens = meta["tokens"]
            ner_labels = all_ner_labels[idx]
            analyses = all_analyses[idx]
            chunk = meta["chunk"]

            # Process single chunk using our batch predictions
            chunk_claims, entities = self._parse_predictions(chunk, tokens, ner_labels, analyses)
            
            if entities:
                running_context.update([e["text"] for e in entities])

            for claim in chunk_claims:
                if claim.startswith(DANGLING_PREFIXES) and running_context:
                    context_str = "، ".join(list(running_context)[:3])
                    claim = f"[سياق: {context_str}] {claim}"
                
                # Filter uniques rapidly
                if claim not in seen:
                    seen.add(claim)
                    claims.append(claim)

        return claims

    def _layer2_split(self, sentence: str) -> list[str]:
        if len(sentence.split()) <= 20:
            return [sentence]

        MARKER = '<<<SPLIT>>>'
        
        # Uses C-compiled Regex instead of loop iterations for replacements
        marked = SPLIT_PATTERN.sub(lambda m: f" {MARKER}{m.group(1).strip()} ", sentence)

        chunks = []
        for part in marked.split(MARKER):
            part = part.strip().rstrip('،,')
            if part and len(part.split()) >= 4:
                chunks.append(part)

        return chunks if chunks else [sentence]

    def _parse_predictions(self, chunk, tokens, ner_labels, analyses):
        entities = self._merge_bio_entities(tokens, ner_labels)
        
        verbs_found = False
        nouns_found = False

        for i, analysis in enumerate(analyses):
            if analysis.analyses:
                best = analysis.analyses[0].analysis
                pos = best.get("pos", "")
                if pos in VERB_POS:
                    verbs_found = True
                    break  # Early exit if we only care "if" they exist
                elif pos in {"noun", "noun_prop"}:
                    nouns_found = True
        
        if not verbs_found or (not entities and not nouns_found):
            return [], entities

        claim = dediac_ar(chunk.strip())
        return [claim], entities

    def _merge_bio_entities(self, tokens: list[str], labels: list[str]) -> list[dict]:
        entities = []
        curr_tokens, curr_type, curr_start = [], None, None

        for i, (token, label) in enumerate(zip(tokens, labels)):
            if label.startswith("B-"):
                if curr_tokens:
                    entities.append({"text": " ".join(curr_tokens), "type": curr_type, "start": curr_start})
                curr_tokens, curr_type, curr_start = [token], label[2:], i
            elif label.startswith("I-") and curr_tokens:
                curr_tokens.append(token)
            else:
                if curr_tokens:
                    entities.append({"text": " ".join(curr_tokens), "type": curr_type, "start": curr_start})
                    curr_tokens, curr_type, curr_start = [], None, None

        if curr_tokens:
            entities.append({"text": " ".join(curr_tokens), "type": curr_type, "start": curr_start})
        return entities


# ── Gradio Interface ────────────────────────────────────────────────────────
extractor = ClaimExtractorLogic()

def predict(text):
    try:
        return extractor.run(text)
    except Exception as e:
        import traceback
        return [f"Error: {str(e)}", traceback.format_exc()]

# Passing api_name helps caching inside spaces
demo = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(label="Arabic News Text", lines=10),
    outputs=gr.JSON(label="Extracted Claims"),
    title="Arabic Claim Extractor",
    description="Extracts atomic verifiable claims from Arabic news articles using Camel-Tools NER and POS tagging.",
)

if __name__ == "__main__":
    demo.launch()