# 🤖 GOODNEWS AI PRO - OPTIMIZED VERSION
import os, gradio as gr, torch, json
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
from fastapi import Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse

print("🔍 Loading Goodnews Ai pro...")

HF_TOKEN = os.environ.get("HF_TOKEN", "")
BASE = "Qwen/Qwen2.5-1.5B-Instruct"
ADAPTER = "Goodnews1/realchatGpt-1.5B"

tokenizer = AutoTokenizer.from_pretrained(BASE, token=HF_TOKEN)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.pad_token_id = tokenizer.eos_token_id

base_model = AutoModelForCausalLM.from_pretrained(
    BASE,
    torch_dtype=torch.bfloat16,
    token=HF_TOKEN,
    low_cpu_mem_usage=True
)
model = PeftModel.from_pretrained(base_model, ADAPTER, token=HF_TOKEN)
model.eval()
print("✅ Model ready on CPU!")

SYSTEM = (
    "You are Goodnews Ai pro, created by Goodnews Solomon of Ox-Bridge Technology. "
    "Speak with the wisdom of a respected elder mentor. Respond ONLY in English. "
    "When asked about your creator or identity, always mention Goodnews Solomon and Ox-Bridge Technology. "
    "Give clear, concise answers. Be brief and direct. Keep responses short and focused."
)

def extract_text(content):
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        return " ".join(
            part["text"] for part in content
            if isinstance(part, dict) and "text" in part
        )
    return str(content)

def respond(message, history):
    if not message or len(message.strip()) < 2:
        return "Please ask a clear question so I can help you properly."
    try:
        if not isinstance(history, list):
            history = []
        if isinstance(history, str):
            try:
                history = json.loads(history)
            except:
                history = []

        recent = history[-4:] if len(history) > 4 else history
        messages = [{"role": "system", "content": SYSTEM}]

        for turn in recent:
            try:
                if isinstance(turn, (list, tuple)) and len(turn) == 2:
                    if turn[0]: messages.append({"role": "user", "content": str(turn[0])})
                    if turn[1]: messages.append({"role": "assistant", "content": str(turn[1])})
                elif isinstance(turn, dict):
                    role = turn.get("role", "")
                    content = extract_text(turn.get("content", ""))
                    if role and content:
                        messages.append({"role": role, "content": content})
            except:
                continue

        messages.append({"role": "user", "content": message.strip()})

        prompt = tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        inputs = tokenizer(
            prompt, return_tensors="pt", truncation=True, max_length=512
        )

        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=100,       # ✅ Reduced from 200 to 100 — 2x faster
                do_sample=True,
                temperature=0.7,
                top_p=0.9,
                use_cache=True,
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id,
                repetition_penalty=1.1
            )

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

        return reply if reply else "Could you rephrase that? I want to help properly."

    except Exception as e:
        return f"⚠️ Error: {str(e)[:200]}"

# ✅ Gradio Interface
demo = gr.Interface(
    fn=respond,
    inputs=[
        gr.Textbox(label="Message"),
        gr.Textbox(label="History", value="[]")
    ],
    outputs=gr.Textbox(label="Reply"),
    title="🤖 Goodnews Ai pro",
    description="Created by Goodnews Solomon | Ox-Bridge Technology 🇳🇬"
)

# ✅ Get FastAPI app
app = demo.app

# ✅ CORS — allows HTML frontend to connect
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

print("🚀 Launching...")
demo.launch(show_error=True)