import os
import json
import time
import random

import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

# ======================================================
# AmCCM v1.0 — Adaptive Memory of Contextual Creativity Model
# Local Qwen2.5-1.5B-Instruct (no external API, CPU-friendly)
# ======================================================

MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
MEMORY_FILE = "amccm_memory_v1.json"
MAX_NEW_TOKENS = 256

# ----------------- Model load (once) -----------------
device = 0 if torch.cuda.is_available() else -1

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
)
generator = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device=device,
)


# ======================================================
# Utility: memory
# ======================================================
def load_memory():
    if os.path.exists(MEMORY_FILE):
        try:
            with open(MEMORY_FILE, "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception:
            pass
    return {"log": []}


def save_memory(mem):
    try:
        with open(MEMORY_FILE, "w", encoding="utf-8") as f:
            json.dump(mem, f, indent=2, ensure_ascii=False)
    except Exception:
        pass


def append_log(question, answer, mode, intent, certainty_label):
    mem = load_memory()
    mem["log"].append(
        {
            "ts": time.time(),
            "mode": mode,
            "intent": intent,
            "certainty": certainty_label,
            "question": question,
            "answer_preview": answer[:400],
        }
    )
    mem["log"] = mem["log"][-80:]
    save_memory(mem)


# ======================================================
# Core LLM call (local Qwen) — plain text prompt
# ======================================================
def call_llm(prompt: str, max_new_tokens=MAX_NEW_TOKENS, temperature=0.7, top_p=0.9):
    try:
        out = generator(
            prompt,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=temperature,
            top_p=top_p,
            pad_token_id=tokenizer.eos_token_id,
        )
        full = out[0]["generated_text"]
        completion = full[len(prompt) :].strip()
        return completion
    except Exception as e:
        return f"[AmCCM ERROR] {e}"


# ======================================================
# Context Awareness Core (CAC) — intent + literal detection
# ======================================================
def classify_intent(text: str) -> str:
    t = text.lower()

    fact_triggers = [
        "what is",
        "when did",
        "who is",
        "where is",
        "how many",
        "explain",
        "define",
        "definition",
        "histor",
        "date of",
        "formula",
        "law of",
        "how does",
        "how do",
    ]
    creative_triggers = [
        "story",
        "backstory",
        "character",
        "poem",
        "song",
        "lyrics",
        "rap",
        "novel",
        "scene",
        "fanfic",
        "fan fiction",
        "roleplay",
        "role play",
        "write a story",
        "write me a story",
    ]

    fact_hits = any(k in t for k in fact_triggers)
    creative_hits = any(k in t for k in creative_triggers)

    if fact_hits and not creative_hits:
        return "fact"
    if creative_hits and not fact_hits:
        return "creative"
    if fact_hits and creative_hits:
        return "mixed"

    if "why" in t or "tell me about" in t:
        return "mixed"

    return "open"  # fallback


def detect_literal_mode(text: str) -> bool:
    t = text.lower()
    triggers = [
        "no metaphors",
        "without metaphors",
        "bez metafor",
        "bez prirovnaní",
        "literally",
        "literal explanation",
        "purely factual",
        "zrozumiteľne",
        "jednoducho vysvetli",
    ]
    return any(k in t for k in triggers)


# ======================================================
# Prompt builder (adds minimal chat context)
# ======================================================
def build_prompt(user_message: str, history, system_mode: str) -> str:
    """
    Simple instruction-style prompt for Qwen in pure text mode.
    """
    sys_desc = (
        "You are AmCCM v1.0, a careful assistant that:\n"
        "- Tries to be accurate on factual questions\n"
        "- Is creative when asked for stories or fiction\n"
        "- Explicitly marks when parts of an answer may be speculative or uncertain\n"
        "- Uses clear, direct language by default\n\n"
        f"Current behavior profile: {system_mode}\n\n"
    )

    convo = sys_desc + "Conversation so far:\n"
    for u, a in history:
        convo += f"User: {u}\nAssistant: {a}\n"
    convo += f"User: {user_message}\nAssistant:"
    return convo


# ======================================================
# Hallucination Intercept Layer (HIL)
# ======================================================
def evaluate_certainty(question: str, answer: str) -> str:
    """
    Returns: 'SAFE', 'UNSURE', or 'RISKY'
    """
    eval_prompt = (
        "You are an AI that judges how reliable an answer is.\n\n"
        "Read the user question and the assistant answer.\n"
        "Decide if the answer is:\n"
        "- SAFE: likely factual and mostly correct\n"
        "- UNSURE: partially speculative or incomplete\n"
        "- RISKY: likely hallucinated or mostly made up\n\n"
        "Respond with exactly one word: SAFE, UNSURE, or RISKY.\n\n"
        f"Question:\n{question}\n\nAnswer:\n{answer}\n\nLabel:"
    )

    raw = call_llm(eval_prompt, max_new_tokens=8, temperature=0.1, top_p=0.9)
    label = raw.strip().upper()

    if "RISK" in label:
        return "RISKY"
    if "UNSURE" in label:
        return "UNSURE"
    if "SAFE" in label:
        return "SAFE"
    return "UNSURE"


def attach_uncertainty_notice(answer: str, label: str) -> str:
    if label == "SAFE":
        return answer
    if label == "UNSURE":
        return (
            answer
            + "\n\n[AmCCM Notice] Some parts of this answer may be uncertain or approximate. "
              "Do not treat this as guaranteed factual."
        )
    if label == "RISKY":
        return (
            answer
            + "\n\n[AmCCM Notice] This answer may go beyond reliable training data and could be incorrect or hallucinated, "
              "even if it sounds confident."
        )
    return answer


# ======================================================
# Creativity control (simple temp/top-p steering)
# ======================================================
def creativity_profile(intent: str, literal_mode: bool, ui_mode: str):
    """
    Returns (temperature, top_p) based on intent + literal/creative flags + UI override.
    ui_mode: 'Auto', 'Factual', 'Creative', 'Balanced'
    """
    if ui_mode == "Factual":
        return 0.35, 0.9
    if ui_mode == "Creative":
        return 0.95, 0.96
    if ui_mode == "Balanced":
        return 0.7, 0.92

    # Auto:
    if literal_mode:
        return 0.3, 0.9

    if intent == "fact":
        return 0.4, 0.9
    if intent == "creative":
        return 0.9, 0.95
    if intent == "mixed":
        return 0.65, 0.93

    # open / fallback
    return 0.7, 0.92


# ======================================================
# Literal explainer mode (no metaphors, no fiction)
# ======================================================
def literal_explainer(question: str) -> str:
    prompt = (
        "Explain the following as clearly and literally as possible.\n"
        "Use simple, factual language only.\n"
        "No metaphors, no analogies, no fictional stories.\n\n"
        f"Question:\n{question}\n\nAnswer:"
    )
    return call_llm(prompt, max_new_tokens=220, temperature=0.3, top_p=0.9)


# ======================================================
# Main AmCCM answering pipeline
# ======================================================
def amccm_answer(message: str, history, ui_mode: str) -> str:
    """
    Core pipeline:
    1) Classify intent (fact / creative / mixed / open)
    2) Detect literal mode
    3) Build prompt with minimal conversation
    4) Generate base answer
    5) If factual or mixed → run hallucination intercept + attach notice
    6) Log interaction
    """
    intent = classify_intent(message)
    literal_mode = detect_literal_mode(message)

    if literal_mode and ui_mode != "Creative":
        answer = literal_explainer(message)
        certainty = "UNSURE"
        append_log(message, answer, ui_mode, intent, certainty)
        return answer

    if len(message.strip().split()) <= 3 and message.strip().lower() in {
        "hi",
        "hello",
        "hey",
        "yo",
        "čau",
        "čaute",
        "ahoj",
    }:
        answer = "Ahoj, som AmCCM. Čo chceš skúsiť?"
        certainty = "SAFE"
        append_log(message, answer, ui_mode, intent, certainty)
        return answer

    temp, top_p = creativity_profile(intent, literal_mode, ui_mode)

    sys_mode = f"UI mode: {ui_mode}, intent: {intent}, literal_mode: {literal_mode}"
    prompt = build_prompt(message, history, sys_mode)

    base_answer = call_llm(
        prompt,
        max_new_tokens=MAX_NEW_TOKENS,
        temperature=temp,
        top_p=top_p,
    )

    if intent in {"fact", "mixed"} and ui_mode != "Creative":
        certainty = evaluate_certainty(message, base_answer)
        final_answer = attach_uncertainty_notice(base_answer, certainty)
    else:
        certainty = "SAFE"
        final_answer = base_answer

    append_log(message, final_answer, ui_mode, intent, certainty)
    return final_answer


# ======================================================
# Gradio UI wiring
# ======================================================
def ui_response(message, history, mode):
    return amccm_answer(message, history, mode)


mode_dropdown = gr.Dropdown(
    choices=["Auto", "Factual", "Creative", "Balanced"],
    value="Auto",
    label="AmCCM Mode",
)

demo = gr.ChatInterface(
    fn=ui_response,
    additional_inputs=[mode_dropdown],
    title="AmCCM v1.0 — Adaptive Memory of Contextual Creativity Model",
    description=(
        "Local Qwen2.5-1.5B-Instruct assistant with:\n"
        "- Intent-aware creativity control\n"
        "- Literal mode when requested (no metaphors)\n"
        "- Hallucination Intercept Layer that marks uncertain / risky answers\n"
        "- Lightweight memory log of interactions"
    ),
)

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