import torch
import json
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import openai

SYSTEM_PROMPT = (
    "You are an advanced AI model specialized in extracting aspects and determining their sentiment polarity from customer reviews.\n\n"
    "Instructions:\n"
    "1. Extract only the aspects (nouns) mentioned in the review.\n"
    "2. Assign a sentiment to each aspect: \"positive\", \"negative\", or \"neutral\".\n"
    "3. Return aspects in the same language as they appear.\n"
    "4. An aspect must be a noun that refers to a specific item or service the user described.\n"
    "5. Ignore adjectives, general ideas, and vague topics.\n"
    "6. Do NOT translate, explain, or add extra text.\n"
    "7. The output must be just a valid JSON list with 'aspect' and 'sentiment'. Start with `[` and stop at `]`.\n"
    "8. Do NOT output the instructions, review, or any text — only one output JSON list.\n"
    "9. Just one output and one review."
)


MODEL_OPTIONS = {
    "Araberta": {
        "base": "asmashayea/absa-araberta",
        "adapter": "asmashayea/absa-araberta"
    },
    "mT5": {
        "base": "google/mt5-base",
        "adapter": "asmashayea/mt4-absa"
    },
    "mBART": {
        "base": "facebook/mbart-large-50-many-to-many-mmt",
        "adapter": "asmashayea/mbart-absa"
    },
    "GPT3.5": {"base": "openai/gpt-3.5-turbo",
               "model_id": "ft:gpt-3.5-turbo-0125:asma:gpt-3-5-turbo-absa:Bb6gmwkE"},
    "GPT4o": {"base": "openai/gpt-4o",
              "model_id": "ft:gpt-4o-mini-2024-07-18:asma:gpt4-finetune-absa:BazoEjnp"},
    "ALLaM": {
        "base": "ALLaM-AI/ALLaM-7B-Instruct-preview",
        "adapter": "asmashayea/allam-absa"
    },
    "DeepSeek": {
        "base": "deepseek-ai/deepseek-llm-7b-chat",
        "adapter": "asmashayea/deepseek-absa"
    }
}


cached_models = {}

# ✅ Reusable for both mT5 + mBART
def load_mt5_bart(model_key):
    base_id = MODEL_OPTIONS[model_key]["base"]
    adapter_id = MODEL_OPTIONS[model_key]["adapter"]

    tokenizer = AutoTokenizer.from_pretrained(adapter_id)
    base_model = AutoModelForSeq2SeqLM.from_pretrained(base_id)
    peft_model = PeftModel.from_pretrained(base_model, adapter_id)
    peft_model.eval()

    cached_models[model_key] = (tokenizer, peft_model)
    return tokenizer, peft_model

def infer_t5_bart(text, model_choice):

    tokenizer, peft_model = load_mt5_bart(model_choice)
    prompt = SYSTEM_PROMPT + f"\n\nReview: {text}"
    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True).to(peft_model.device)

    with torch.no_grad():
        outputs = peft_model.generate(
            **inputs,
            max_new_tokens=256,
            num_beams=4,
            do_sample=False,
            temperature=0.0,
            early_stopping=True,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
    decoded = decoded.replace('<extra_id_0>', '').replace('</s>', '').strip()

    try:
        return json.loads(decoded)
    except json.JSONDecodeError:
        return {"raw_output": decoded, "error": "Invalid JSON"}



OPENAI_API_KEY = "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"
openai.api_key = OPENAI_API_KEY

def infer_gpt_absa(text, model_key):
    MODEL_ID = MODEL_OPTIONS[model_key]["model_id"]
    try:
        response = openai.chat.completions.create(
            model=MODEL_ID,
            messages=[
                {
                    "role": "system",
                    "content": SYSTEM_PROMPT
                },
                {
                    "role": "user",
                    "content": text
                }
            ],
            temperature=0
        )
        decoded = response.choices[0].message.content.strip()
        return json.loads(decoded)
    except Exception as e:
        return {"error": str(e)}



def infer_allam(review_text):
    tokenizer, model = cached_models.get("ALLaM") or load_allam()

    prompt = tokenizer.apply_chat_template(
        [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": review_text}
        ],
        tokenize=False
    )

    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512).to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=128,
            do_sample=False,
            temperature=0.0,
            pad_token_id=tokenizer.eos_token_id
        )

    decoded = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True).strip()
    try:
        parsed = json.loads(decoded)
        return parsed
    except Exception as e:
        return {"error": str(e), "raw": decoded}


def load_allam():
    base_model = AutoModelForCausalLM.from_pretrained(
        MODEL_OPTIONS["ALLaM"]["base"],
        device_map="auto",
        torch_dtype=torch.float16,
        trust_remote_code=True
    )
    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_OPTIONS["ALLaM"]["adapter"],
        trust_remote_code=True
    )
    model = PeftModel.from_pretrained(base_model, MODEL_OPTIONS["ALLaM"]["adapter"])

    cached_models["ALLaM"] = (tokenizer, model)
    return tokenizer, model





def load_allam():
    base = AutoModelForCausalLM.from_pretrained(
        MODEL_OPTIONS["ALLaM"]["base"],
        torch_dtype=torch.float16,
        trust_remote_code=True
    )
    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_OPTIONS["ALLaM"]["adapter"], trust_remote_code=True
    )
    model = PeftModel.from_pretrained(base, MODEL_OPTIONS["ALLaM"]["adapter"])
    cached_models["ALLaM"] = (tokenizer, model)
    return tokenizer, model

def infer_allam(review):
    if "ALLaM" not in cached_models:
        tokenizer, model = load_allam()
    else:
        tokenizer, model = cached_models["ALLaM"]

    prompt = tokenizer.apply_chat_template([
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": review}
    ], tokenize=False)

    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    with torch.no_grad():
        output = model.generate(**inputs, max_new_tokens=256)
    decoded = tokenizer.decode(output[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)

    try:
        return json.loads(decoded)
    except:
        return decoded




def build_deepseek_prompt(review_text, output=""):
    return f"""<|system|>
You are an advanced AI model specialized in extracting aspects and determining their sentiment polarity from customer reviews.


Instructions:
1. Extract only the aspects (nouns) mentioned in the review.
2. Assign a sentiment to each aspect: "positive", "negative", or "neutral".
3. Return aspects in the same language as they appear.
4. An aspect must be a noun that refers to a specific item or service the user described.
5. Ignore adjectives, general ideas, and vague topics.
6. Do NOT translate, explain, or add extra text.
7. The output must be just a valid JSON list with 'aspect' and 'sentiment'. Start with `[` and stop at `]`.
8. Do NOT output the instructions, review, or any text — only one output JSON list.
9. Just one output and one review.
<|user|>
{review_text}
<|assistant|>
{output}"""  # ✅ include the output here


def load_deepseek():
    base = AutoModelForCausalLM.from_pretrained(
        MODEL_OPTIONS["DeepSeek"]["base"],
        torch_dtype=torch.float16,
        trust_remote_code=True
    )
    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_OPTIONS["DeepSeek"]["adapter"], trust_remote_code=True
    )
    model = PeftModel.from_pretrained(base, MODEL_OPTIONS["DeepSeek"]["adapter"])
    cached_models["DeepSeek"] = (tokenizer, model)
    return tokenizer, model

def infer_deepseek(review):
    if "DeepSeek" not in cached_models:
        tokenizer, model = load_deepseek()
    else:
        tokenizer, model = cached_models["DeepSeek"]

    prompt = build_deepseek_prompt(review)
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512).to(model.device)

    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=128,
            do_sample=False,
            temperature=0.0,
            pad_token_id=tokenizer.eos_token_id
        )

    decoded = tokenizer.decode(
        output[0][inputs["input_ids"].shape[1]:],
        skip_special_tokens=True
    ).strip()

    try:
        return json.loads(decoded)
    except Exception as e:
        print(f"❌ DeepSeek JSON parse error: {e}")
        return decoded