"""
============================================================
AgriExplain
Explainable AI Agricultural Decision Support System
============================================================
"""

# IMPORTANT: 'spaces' MUST BE IMPORTED FIRST BEFORE TORCH/CUDA
import spaces

import os
import re
from threading import Thread

import gradio as gr
import torch
from peft import PeftModel
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TextIteratorStreamer,
)

# ============================================================
# CONFIGURATION
# ============================================================

APP_NAME = "AgriExplain"
BASE_MODEL = "unsloth/llama-3-8b-Instruct-bnb-4bit"
LORA_MODEL = "BenKhemical/llama_agri_lora"

MAX_NEW_TOKENS = 300
TEMPERATURE = 0.5
REPETITION_PENALTY = 1.2
TOP_P = 0.9

HF_TOKEN = os.getenv("HF_TOKEN", None)

# Global model and tokenizer containers
tokenizer = None
model = None

# ============================================================
# LAZY MODEL INITIALIZATION FOR ZEROGPU
# ============================================================

@spaces.GPU
def load_model_if_needed():
    global tokenizer, model
    if model is not None and tokenizer is not None:
        return

    print("\n" + "=" * 60)
    print("Loading AgriExplain on ZeroGPU...")
    print("=" * 60)

    print("\nLoading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, token=HF_TOKEN)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    print("\nLoading base model...")
    base_model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        device_map="auto",
        token=HF_TOKEN,
    )

    print("\nLoading AgriExplain LoRA adapter...")
    model = PeftModel.from_pretrained(
        base_model,
        LORA_MODEL,
        token=HF_TOKEN,
    )
    model.eval()

    print("\n" + "=" * 60)
    print("AgriExplain model is ready.")
    print("=" * 60)

# ============================================================
# AGRICULTURAL SCOPE
# ============================================================

AGRICULTURE_KEYWORDS = [
    "agriculture", "agricultural", "farm", "farmer", "farmers", "farming",
    "crop", "crops", "field", "maize", "corn", "rice", "cassava", "yam",
    "plantain", "banana", "cocoa", "tomato", "tomatoes", "pepper", "onion",
    "potato", "beans", "bean", "soybean", "groundnut", "groundnuts",
    "vegetable", "vegetables", "pest", "pests", "insect", "insects", "aphid",
    "aphids", "thrips", "whitefly", "whiteflies", "caterpillar", "caterpillars",
    "weevil", "weevils", "beetle", "beetles", "borer", "borers", "worm",
    "worms", "disease", "diseases", "fungus", "fungal", "bacteria", "bacterial",
    "virus", "viral", "blight", "rust", "mildew", "rot", "leaf spot",
    "leaf spots", "soil", "soil fertility", "soil nutrient", "soil nutrients",
    "nutrient", "nutrients", "nitrogen", "phosphorus", "potassium",
    "fertilizer", "fertiliser", "compost", "manure", "soil ph", "soil pH",
    "planting", "plant", "plants", "seed", "seeds", "germination", "spacing",
    "planting season", "planting date", "harvest", "harvesting", "irrigation",
    "watering", "water", "weed", "weeds", "weeding", "mulching", "pruning",
    "yield", "leaf", "leaves", "stem", "roots", "root", "yellow leaves",
    "yellowing", "wilting", "wilt", "spots", "holes", "discoloration",
    "brown leaves", "dry leaves", "curling leaves"
]

OUT_OF_SCOPE_MESSAGE = (
    "🌱 **AgriExplain** is an agricultural decision-support system.\n\n"
    "Please ask a question related to agriculture, such as "
    "crop management, pests, crop diseases, soil management, "
    "fertilizer, planting, irrigation, or harvesting."
)

# ============================================================
# HELPER FUNCTIONS
# ============================================================

def is_in_scope(message: str) -> bool:
    if not message:
        return False
    msg = message.lower().strip()
    return any(keyword in msg for keyword in AGRICULTURE_KEYWORDS)

def remove_repeated_sentences(text: str) -> str:
    if not text:
        return ""
    sentences = re.split(r"(?<=[.!?])\s+", text)
    unique_sentences = []
    seen = set()
    for sentence in sentences:
        sentence = sentence.strip()
        if not sentence:
            continue
        normalized = re.sub(r"\s+", " ", sentence.lower())
        if normalized in seen:
            continue
        seen.add(normalized)
        unique_sentences.append(sentence)
    return " ".join(unique_sentences)

def clean_agri_output(raw_text: str, user_message: str = "") -> str:
    """
    Cleans model output. 
    - For Pest/Disease queries: Displays Problem, Possible Cause, Recommended Action, and Explanation.
    - For other queries: Displays Recommended Action and Explanation.
    """
    if not raw_text:
        return ""

    text = raw_text.strip()

    # 1. Strip raw metadata key=value pairs and structural tags
    text = re.sub(
        r'(?:Context\s*:?\s*)?(?:Region|Owner|Farm|Season|Soil|Severity|Source)\s*=\s*.*?(?=\.|\n|\Z)',
        '', text, flags=re.IGNORECASE
    )
    text = re.sub(r'\bContext\s*:?\s*', '', text, flags=re.IGNORECASE)
    text = re.sub(r'https?://\S+|\[.*?\]\(.*?\)', '', text)
    text = re.sub(r"<\|.*?\|>", "", text)

    # 2. Check if the query is related to Pests or Diseases
    pest_disease_keywords = [
        "pest", "pests", "insect", "insects", "aphid", "thrips", "whitefly", 
        "caterpillar", "weevil", "beetle", "borer", "worm", "disease", 
        "diseases", "fungus", "fungal", "bacteria", "bacterial", "virus", 
        "viral", "blight", "rust", "mildew", "rot", "leaf spot", "spots"
    ]
    
    combined_check_text = (user_message + " " + text).lower()
    is_pest_or_disease = any(k in combined_check_text for k in pest_disease_keywords)

    # 3. Extract sections using Regex
    problem_match = re.search(
        r"Problem\s*:\s*(.*?)(?=\n\s*(?:Possible Cause|Recommended Action|Recommendation|Explanation)\s*:|\Z)",
        text, flags=re.IGNORECASE | re.DOTALL
    )
    cause_match = re.search(
        r"Possible Cause\s*:\s*(.*?)(?=\n\s*(?:Problem|Recommended Action|Recommendation|Explanation)\s*:|\Z)",
        text, flags=re.IGNORECASE | re.DOTALL
    )
    recommendation_match = re.search(
        r"(?:Recommended Action|Recommendation)\s*:\s*(.*?)(?=\n\s*(?:Problem|Possible Cause|Explanation)\s*:|\Z)",
        text, flags=re.IGNORECASE | re.DOTALL
    )
    explanation_match = re.search(
        r"Explanation\s*:\s*(.*?)(?=\n\s*(?:Problem|Possible Cause|Recommended Action|Recommendation)\s*:|\Z)",
        text, flags=re.IGNORECASE | re.DOTALL
    )

    problem = remove_repeated_sentences(problem_match.group(1).strip()) if problem_match else ""
    cause = remove_repeated_sentences(cause_match.group(1).strip()) if cause_match else ""
    recommendation = remove_repeated_sentences(recommendation_match.group(1).strip()) if recommendation_match else ""
    explanation = remove_repeated_sentences(explanation_match.group(1).strip()) if explanation_match else ""

    # Fallback cleanup for unstructured generation
    if not recommendation:
        cleaned = re.sub(r"^(#+\s*)?(Response|Instruction|Recommendation)\s*:\s*", "", text, flags=re.IGNORECASE)
        recommendation = remove_repeated_sentences(re.split(r"\n#+\s*|\nContext:|\nExplanation:|\bRegion=", cleaned)[0].strip())

    recommendation = re.sub(r"\s+", " ", recommendation).strip()
    explanation = re.sub(r"\s+", " ", explanation).strip()
    problem = re.sub(r"\s+", " ", problem).strip()
    cause = re.sub(r"\s+", " ", cause).strip()

    output = []

    # 4. Construct response according to query type
    if is_pest_or_disease:
        if problem and problem != ".":
            output.append(f"### ⚠️ Problem\n\n{problem}")
        if cause and cause != ".":
            output.append(f"### 🔍 Possible Cause\n\n{cause}")
        if recommendation and recommendation != ".":
            output.append(f"### 🌱 Recommended Action\n\n{recommendation}")
        if explanation and explanation != ".":
            output.append(f"### 💡 Explanation\n\n{explanation}")
    else:
        if recommendation and recommendation != ".":
            output.append(f"### 🌱 Recommended Action\n\n{recommendation}")
        if explanation and explanation != ".":
            output.append(f"### 💡 Explanation\n\n{explanation}")

    if not output:
        return "I was unable to generate a clear agricultural recommendation. Please try rephrasing your question."

    final_output = "\n\n".join(output)
    return re.sub(r"\n{3,}", "\n\n", final_output).strip()

# ============================================================
# RESPONSE GENERATION (MESSAGES FORMAT)
# ============================================================

@spaces.GPU
def generate_response(message, history):
    load_model_if_needed()

    if history is None:
        history = []

    if isinstance(message, dict):
        user_message = message.get("text", "")
    elif isinstance(message, str):
        user_message = message
    else:
        user_message = ""

    user_message = user_message.strip()

    if not user_message:
        history.append({"role": "user", "content": ""})
        history.append({"role": "assistant", "content": "Please enter a valid agricultural question."})
        yield history
        return

    history.append({"role": "user", "content": user_message})

    if not is_in_scope(user_message):
        history.append({"role": "assistant", "content": OUT_OF_SCOPE_MESSAGE})
        yield history
        return

    history.append({"role": "assistant", "content": ""})

    messages = [{"role": "user", "content": user_message}]
    inputs = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt",
    ).to(model.device)

    if hasattr(inputs, "input_ids"):
        input_ids = inputs.input_ids
        attention_mask = getattr(inputs, "attention_mask", torch.ones_like(input_ids))
    else:
        input_ids = inputs
        attention_mask = torch.ones_like(input_ids)

    streamer = TextIteratorStreamer(
        tokenizer, skip_prompt=True, skip_special_tokens=True
    )

    generation_kwargs = dict(
        input_ids=input_ids,
        attention_mask=attention_mask,
        max_new_tokens=MAX_NEW_TOKENS,
        temperature=TEMPERATURE,
        top_p=TOP_P,
        repetition_penalty=REPETITION_PENALTY,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id,
        eos_token_id=tokenizer.eos_token_id,
        streamer=streamer,
    )

    thread = Thread(target=model.generate, kwargs=generation_kwargs, daemon=True)
    thread.start()

    partial_text = ""
    for new_text in streamer:
        partial_text += new_text
        cleaned = clean_agri_output(partial_text, user_message=user_message)
        if cleaned:
            history[-1]["content"] = cleaned
            yield history

    thread.join()
    final_response = clean_agri_output(partial_text, user_message=user_message)
    if final_response:
        history[-1]["content"] = final_response
    else:
        history[-1]["content"] = "I was unable to generate a recommendation. Please try asking your question in a different way."
    yield history

# ============================================================
# GRADIO APPLICATION
# ============================================================

CUSTOM_CSS = """
.gradio-container {
    max-width: 1100px !important;
    margin: auto !important;
}
.agri-header {
    text-align: center;
    padding: 20px 10px 10px 10px;
}
.agri-title {
    font-size: 42px !important;
    font-weight: 700 !important;
}
.agri-subtitle {
    font-size: 17px !important;
    opacity: 0.8;
}
.agri-footer {
    text-align: center;
    margin-top: 20px;
    opacity: 0.7;
    font-size: 13px;
}
"""

EXAMPLES = [
    ["My maize leaves are turning yellow. What should I do?"],
    ["What can I do to control pests on my tomato farm?"],
    ["Why are there spots on my cassava leaves?"],
    ["What fertilizer is suitable for maize?"],
    ["When is the best time to plant maize?"],
    ["How can I improve soil fertility?"],
]

with gr.Blocks(title="AgriExplain") as app:

    gr.HTML(
        """
        <div class="agri-header">
            <div class="agri-title">🌱 AgriExplain</div>
            <div class="agri-subtitle">Explainable AI Agricultural Decision Support System</div>
        </div>
        """
    )

    gr.Markdown(
        """
        **AgriExplain** helps farmers and agricultural users obtain AI-powered recommendations
        about crop management, pests, diseases, soil, planting, fertilizer, irrigation, and other farming activities.
        """
    )

    chatbot = gr.Chatbot(
        label="AgriExplain Assistant",
        height=520,
    )

    message_box = gr.Textbox(
        label="Ask your agricultural question",
        placeholder="Example: My maize leaves are turning yellow. What should I do?",
        lines=3,
        max_lines=6,
    )

    with gr.Row():
        submit_button = gr.Button("🌱 Ask AgriExplain", variant="primary")
        clear_button = gr.Button("🗑️ Clear")

    gr.Examples(
        examples=EXAMPLES,
        inputs=message_box,
        label="Try an example",
    )

    with gr.Accordion("About AgriExplain", open=False):
        gr.Markdown(
            """
            ### What is AgriExplain?
            AgriExplain is an Explainable Artificial Intelligence agricultural decision-support
            system designed to provide farming recommendations together with explanations.

            ### Areas of support
            - 🌽 Crop management
            - 🐛 Pest management
            - 🦠 Crop disease management
            - 🌱 Soil management
            - 💧 Irrigation
            - 🌾 Planting
            - 🧪 Fertilizer management
            - 🌿 General farm management
            """
        )

    gr.HTML(
        """
        <div class="agri-footer">
            AgriExplain — Explainable AI Agricultural Decision Support System <br> Final Year Project
        </div>
        """
    )

    submit_button.click(
        fn=generate_response,
        inputs=[message_box, chatbot],
        outputs=chatbot,
    )

    message_box.submit(
        fn=generate_response,
        inputs=[message_box, chatbot],
        outputs=chatbot,
    )

    clear_button.click(
        fn=lambda: [],
        inputs=None,
        outputs=chatbot,
    )

# ============================================================
# RUN APPLICATION
# ============================================================

if __name__ == "__main__":
    print("\nStarting Gradio application...\n")

    app.launch(
        show_error=True,
        theme=gr.themes.Soft(),
        css=CUSTOM_CSS,
    )