# ============================================================
#   LEXOSYNTH - BACKEND : FLASK API
#   File   : backend.py
#
#   Handles all engine logic and exposes a single /analyze
#   POST endpoint consumed by the Gradio frontend.
#
#   Engines supported :
#     1. supersonic      — CSV phrase lookup (no model load)
#     2. go_deep         — LLaMA 1B + LoRA adapter
#     3. hybimix         — LLaMA + Supersonic + confidence merge
#     4. go_deep_max     — Sarvam 2B + LoRA adapter
#     5. hybimix_max     — Sarvam + Supersonic + confidence merge
# ============================================================

import os
import re
import json
import threading
import torch
import pandas as pd
from flask import Flask, request, jsonify

app = Flask(__name__)

# ============================================================
#   Replace SECTION 1 (PATHS SETUP) with this
# ============================================================

import os
from huggingface_hub import snapshot_download, get_token

HF_TOKEN = os.environ.get("HF_TOKEN")   # Loaded from Space secrets

# Optional: fallback to get_token()
if not HF_TOKEN:
    HF_TOKEN = get_token()

# Supersonic CSV path — included directly in Space repo
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
CSV_PATH = os.path.join(BASE_DIR, "supersonic_helper_database", "search_tortured_correct.csv")

# Model paths — downloaded from HF Hub on first use (lazy load)
# These are set inside the lazy loader functions below
LLAMA_MODEL_ID  = "meta-llama/Llama-3.2-1B-Instruct"
LLAMA_ADAPTER_ID = os.environ.get("LLAMA_ADAPTER_REPO")   # from Space secrets
SARVAM_MODEL_ID  = "sarvamai/sarvam-1"
SARVAM_ADAPTER_ID= os.environ.get("SARVAM_ADAPTER_REPO")  # from Space secrets

# ============================================================
# SECTION 2 : LAZY MODEL STORAGE + THREAD LOCKS
# ============================================================

# Models loaded on first use only — avoids loading unused models
_supersonic_db    = None
_llama_model      = None
_llama_tokenizer  = None
_sarvam_model     = None
_sarvam_tokenizer = None

_db_lock          = threading.Lock()
_llama_lock       = threading.Lock()
_sarvam_lock      = threading.Lock()

# ============================================================
# SECTION 3 : SYSTEM PROMPT
# ============================================================

SYSTEM_PROMPT = (
    "You are a scientific language correction assistant specializing in "
    "Artificial Intelligence and Machine Learning (AIML). Your task is to "
    "detect 'tortured phrases' in the given sentence — these are unnatural, "
    "paraphrased substitutions of standard AIML terminology, often introduced "
    "by paraphrasing tools or paper mills to evade plagiarism detection. "
    "Replace each tortured phrase with the correct and widely accepted AIML "
    "term. Return only the corrected sentence with no explanations, labels, "
    "or additional commentary. If no tortured phrases are found, return the "
    "original sentence as-is."
)

# ============================================================
# SECTION 4 : LAZY LOADERS
# ============================================================

def get_supersonic_db():
    """
    Load search_tortured_correct.csv on first call.
    Cached in _supersonic_db for all subsequent calls.
    Thread-safe via _db_lock.
    """
    global _supersonic_db
    with _db_lock:
        if _supersonic_db is None:
            if not os.path.exists(CSV_PATH):
                raise FileNotFoundError(
                    f"Supersonic database not found at : {CSV_PATH}\n"
                    "Please run step8_create_csv_database.py first."
                )
            _supersonic_db = pd.read_csv(CSV_PATH)
    return _supersonic_db


def get_llama_model():
    global _llama_model, _llama_tokenizer
    with _llama_lock:
        if _llama_model is None:
            from transformers import AutoModelForCausalLM, AutoTokenizer
            from peft import PeftModel

            # Download base model from HuggingFace Hub
            llama_path   = snapshot_download(LLAMA_MODEL_ID, token=HF_TOKEN)
            adapter_path = snapshot_download(LLAMA_ADAPTER_ID, token=HF_TOKEN)

            tokenizer = AutoTokenizer.from_pretrained(llama_path)
            if tokenizer.pad_token is None:
                tokenizer.pad_token = tokenizer.eos_token

            base_model = AutoModelForCausalLM.from_pretrained(
                llama_path, torch_dtype=torch.float32
            )
            model = PeftModel.from_pretrained(base_model, adapter_path)
            model.eval()

            _llama_model     = model
            _llama_tokenizer = tokenizer
    return _llama_model, _llama_tokenizer


def get_sarvam_model():
    global _sarvam_model, _sarvam_tokenizer
    with _sarvam_lock:
        if _sarvam_model is None:
            from transformers import AutoModelForCausalLM, AutoTokenizer
            from peft import PeftModel

            sarvam_path  = snapshot_download(
                SARVAM_MODEL_ID, token=HF_TOKEN
            )
            adapter_path = snapshot_download(SARVAM_ADAPTER_ID, token=HF_TOKEN)

            tokenizer = AutoTokenizer.from_pretrained(
                sarvam_path, trust_remote_code=True
            )
            if tokenizer.pad_token is None:
                tokenizer.pad_token = tokenizer.eos_token

            base_model = AutoModelForCausalLM.from_pretrained(
                sarvam_path,
                torch_dtype=torch.float32,
                trust_remote_code=True
            )
            model = PeftModel.from_pretrained(base_model, adapter_path)
            model.eval()

            _sarvam_model     = model
            _sarvam_tokenizer = tokenizer
    return _sarvam_model, _sarvam_tokenizer


# ============================================================
# SECTION 5 : SENTENCE SPLITTER
# ============================================================

def split_into_sentences(text):
    """
    Split input text into individual sentences on . ! ? punctuation.
    Filters empty strings and strips whitespace.
    """
    raw = re.split(r'(?<=[.!?])\s+', text.strip())
    return [s.strip() for s in raw if s.strip()]

# ============================================================
# SECTION 6 : SUPERSONIC ENGINE
# ============================================================

def supersonic_check_already_correct(sentence, df):
    """
    Check if sentence already contains known correct AIML phrases
    using the correct and correct_lower columns.
    """
    found    = []
    s_lower  = sentence.lower()

    for _, row in df.iterrows():
        correct       = str(row["correct"]).strip()
        correct_lower = str(row["correct_lower"]).strip()

        if correct in sentence:
            found.append({"correct": correct, "match_type": "exact"})
        elif correct_lower in s_lower:
            found.append({"correct": correct, "match_type": "case_insensitive"})

    # Deduplicate
    seen, unique = set(), []
    for item in found:
        if item["correct"].lower() not in seen:
            seen.add(item["correct"].lower())
            unique.append(item)
    return unique


def supersonic_correct_sentence(sentence, df):
    """
    Correct a single sentence using CSV exact substring match.
    Also detects already-correct AIML phrases.
    """
    corrected    = sentence
    replacements = []

    for _, row in df.iterrows():
        tortured       = str(row["tortured"]).strip()
        correct        = str(row["correct"]).strip()
        tort_lower     = str(row["tortured_lower"]).strip()

        if tortured in corrected:
            corrected = corrected.replace(tortured, correct)
            replacements.append({
                "tortured": tortured, "correct": correct,
                "match_type": "exact"
            })
        elif tort_lower in corrected.lower():
            pattern   = re.compile(re.escape(tortured), re.IGNORECASE)
            corrected = pattern.sub(correct, corrected)
            replacements.append({
                "tortured": tortured, "correct": correct,
                "match_type": "case_insensitive"
            })

    correct_found = supersonic_check_already_correct(corrected, df)

    if replacements:
        status = "corrected"
    elif correct_found:
        status = "already_correct"
    else:
        status = "unchanged"

    return corrected, replacements, correct_found, status


def run_supersonic(sentences):
    """Run Supersonic engine on a list of sentences."""
    df      = get_supersonic_db()
    results = []

    for sentence in sentences:
        corrected, replacements, correct_found, status = supersonic_correct_sentence(
            sentence, df
        )
        results.append({
            "original"       : sentence,
            "corrected"      : corrected,
            "changed"        : len(replacements) > 0,
            "status"         : status,
            "replacements"   : replacements,
            "correct_found"  : correct_found
        })
    return results

# ============================================================
# SECTION 7 : AI ENGINE — GENERATE RESPONSE
# ============================================================

def ai_generate(sentence, model, tokenizer, max_new_tokens=150):
    """
    Generate a corrected sentence using an AI model + adapter.
    Uses two-step approach: apply_chat_template → tokenizer → generate.
    Strips EOS token if it appears as text in the output.
    """
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user",   "content": sentence}
    ]

    formatted = tokenizer.apply_chat_template(
        messages, add_generation_prompt=True, tokenize=False
    )
    input_ids = tokenizer(formatted, return_tensors="pt")["input_ids"]

    with torch.no_grad():
        output_ids = model.generate(
            input_ids,
            max_new_tokens = max_new_tokens,
            max_length     = None,
            do_sample      = False,
            pad_token_id   = tokenizer.eos_token_id
        )

    new_tokens = output_ids[0][input_ids.shape[-1]:]
    corrected  = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()

    # Strip EOS token if it appears as literal text
    if tokenizer.eos_token and corrected.endswith(tokenizer.eos_token):
        corrected = corrected[:-len(tokenizer.eos_token)].strip()

    return corrected if corrected else sentence


def run_go_deep(sentences):
    """Run Go Deep (LLaMA) engine."""
    model, tokenizer = get_llama_model()
    results = []

    for sentence in sentences:
        corrected = ai_generate(sentence, model, tokenizer)
        results.append({
            "original"  : sentence,
            "corrected" : corrected,
            "changed"   : corrected.strip() != sentence.strip()
        })
    return results


def run_go_deep_max(sentences):
    """Run Go Deep Max (Sarvam) engine."""
    model, tokenizer = get_sarvam_model()
    results = []

    for sentence in sentences:
        corrected = ai_generate(sentence, model, tokenizer)
        results.append({
            "original"  : sentence,
            "corrected" : corrected,
            "changed"   : corrected.strip() != sentence.strip()
        })
    return results

# ============================================================
# SECTION 8 : HYBIMIX CONFIDENCE LOGIC
# ============================================================

def hybimix_confidence(sup_result, ai_result):
    """
    Merge Supersonic and AI engine results using confidence rules.

    Case 1 : Both outputs identical         → 100% confidence
    Case 2 : Only one engine made a change  → 50% confidence
    Case 3 : Both changed differently       → <50% (40%) — show both
    """
    original    = sup_result["original"]
    sup_out     = sup_result["corrected"].strip()
    ai_out      = ai_result["corrected"].strip()
    sup_changed = sup_result["changed"]
    ai_changed  = ai_result["changed"]

    if sup_out == ai_out:
        return {
            "original"         : original,
            "final_corrected"  : sup_out,
            "confidence"       : 100,
            "confidence_label" : "100% — Both engines agree",
            "supersonic_out"   : sup_out,
            "ai_out"           : ai_out,
            "case"             : 1
        }

    if sup_changed and not ai_changed:
        return {
            "original"         : original,
            "final_corrected"  : sup_out,
            "confidence"       : 50,
            "confidence_label" : "50% — Supersonic only",
            "supersonic_out"   : sup_out,
            "ai_out"           : ai_out,
            "case"             : 2
        }

    if ai_changed and not sup_changed:
        return {
            "original"         : original,
            "final_corrected"  : ai_out,
            "confidence"       : 50,
            "confidence_label" : "50% — AI engine only",
            "supersonic_out"   : sup_out,
            "ai_out"           : ai_out,
            "case"             : 2
        }

    return {
        "original"         : original,
        "final_corrected"  : None,
        "confidence"       : 40,
        "confidence_label" : "<50% — Engines disagree",
        "supersonic_out"   : sup_out,
        "ai_out"           : ai_out,
        "case"             : 3
    }


def run_hybimix(sentences):
    """Run Hybimix (LLaMA + Supersonic) engine."""
    sup_results = run_supersonic(sentences)
    ai_results  = run_go_deep(sentences)
    return [hybimix_confidence(s, a) for s, a in zip(sup_results, ai_results)]


def run_hybimix_max(sentences):
    """Run Hybimix Max (Sarvam + Supersonic) engine."""
    sup_results = run_supersonic(sentences)
    ai_results  = run_go_deep_max(sentences)
    return [hybimix_confidence(s, a) for s, a in zip(sup_results, ai_results)]

# ============================================================
# SECTION 9 : OUTPUT FORMATTER
# ============================================================

def format_supersonic_output(results):
    lines = []
    for i, r in enumerate(results, 1):
        lines.append(f"─── Sentence {i} " + "─" * 40)
        lines.append(f"Original  : {r['original']}")
        if r["changed"]:
            lines.append(f"Corrected : {r['corrected']}")
            lines.append(f"Status    : ✅ CORRECTED")
            lines.append("Changes   :")
            for rep in r["replacements"]:
                lines.append(
                    f"  • '{rep['tortured']}'  →  '{rep['correct']}'"
                    f"  [{rep['match_type']}]"
                )
        elif r["status"] == "already_correct":
            lines.append(f"Status    : ✓ Already using correct AIML terms")
            lines.append("Correct phrases found :")
            for cp in r["correct_found"]:
                lines.append(f"  • '{cp['correct']}'")
        else:
            lines.append(f"Status    : — No tortured or correct phrases detected")
        lines.append("")
    return "\n".join(lines).strip()


def format_ai_output(results, engine_label):
    lines = []
    for i, r in enumerate(results, 1):
        lines.append(f"─── Sentence {i} " + "─" * 40)
        lines.append(f"Original  : {r['original']}")
        lines.append(f"Corrected : {r['corrected']}")
        lines.append(f"Changed   : {'✅ Yes' if r['changed'] else '— No change'}")
        lines.append("")
    return "\n".join(lines).strip()


def format_hybimix_output(results, engine_label):
    lines = []
    for i, r in enumerate(results, 1):
        lines.append(f"─── Sentence {i} " + "─" * 40)
        lines.append(f"Original   : {r['original']}")

        if r["case"] == 3:
            lines.append(f"AI Output  : {r['ai_out']}")
            lines.append(f"Supersonic : {r['supersonic_out']}")
            lines.append(f"Result     : ⚠ Engines disagree — manual review recommended")
        else:
            lines.append(f"Corrected  : {r['final_corrected']}")

        lines.append(f"Confidence : {r['confidence_label']}")
        lines.append("")
    return "\n".join(lines).strip()

# ============================================================
# SECTION 10 : FLASK ROUTES
# ============================================================

@app.route("/health", methods=["GET"])
def health():
    """Health check endpoint — used by test.py to confirm Flask is ready."""
    return jsonify({"status": "ok"})


@app.route("/analyze", methods=["POST"])
def analyze():
    """
    Main analysis endpoint.

    Expects JSON body :
    {
      "text"   : "Input text with one or more sentences.",
      "engine" : "supersonic" | "go_deep" | "hybimix" |
                 "go_deep_max" | "hybimix_max"
    }

    Returns JSON :
    {
      "status"           : "success",
      "engine"           : "...",
      "sentence_count"   : N,
      "formatted_output" : "Formatted result string for display"
    }
    """
    data   = request.get_json(force=True)
    text   = data.get("text", "").strip()
    engine = data.get("engine", "").strip()

    # Input validation
    if not text:
        return jsonify({"error": "No text provided."}), 400
    if engine not in ("supersonic", "go_deep", "hybimix",
                      "go_deep_max", "hybimix_max"):
        return jsonify({"error": f"Unknown engine : '{engine}'"}), 400

    # Split into sentences
    sentences = split_into_sentences(text)
    if not sentences:
        return jsonify({"error": "No valid sentences found in input."}), 400

    try:
        if engine == "supersonic":
            results          = run_supersonic(sentences)
            formatted_output = format_supersonic_output(results)

        elif engine == "go_deep":
            results          = run_go_deep(sentences)
            formatted_output = format_ai_output(results, "Go Deep (LLaMA)")

        elif engine == "hybimix":
            results          = run_hybimix(sentences)
            formatted_output = format_hybimix_output(results, "Hybimix")

        elif engine == "go_deep_max":
            results          = run_go_deep_max(sentences)
            formatted_output = format_ai_output(results, "Go Deep Max (Sarvam)")

        elif engine == "hybimix_max":
            results          = run_hybimix_max(sentences)
            formatted_output = format_hybimix_output(results, "Hybimix Max")

        return jsonify({
            "status"          : "success",
            "engine"          : engine,
            "sentence_count"  : len(sentences),
            "formatted_output": formatted_output
        })

    except FileNotFoundError as e:
        return jsonify({"error": str(e)}), 503

    except Exception as e:
        return jsonify({"error": f"Engine error : {str(e)}"}), 500


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000, debug=False, use_reloader=False)