import torch
import requests
from bs4 import BeautifulSoup
from duckduckgo_search import DDGS
from transformers import AutoModelForCausalLM, AutoTokenizer
import wikipedia

model_path = "cosmosai471/Luna-v2"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(
    model_path, torch_dtype=torch.bfloat16
).to("cuda" if torch.cuda.is_available() else "cpu")

def detect_mode(query):
    code_keywords = ["code", "program", "python", "javascript", "function", "script", "build", "html", "css"]
    creative_keywords = ["story", "write a story", "poem", "creative", "imagine", "novel", "dialogue"]

    if any(kw in query.lower() for kw in code_keywords):
        return "code"
    elif any(kw in query.lower() for kw in creative_keywords):
        return "creative"
    else:
        return "general"

def get_generation_params(query):
    mode = detect_mode(query)
    if mode == "code":
        return 0.3, 0.85
    elif mode == "creative":
        return 0.95, 0.95
    else:
        return 0.7, 0.9

def luna_generate(prompt, max_tokens=200, temperature=0.7, top_p=0.95):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    output = model.generate(
        **inputs,
        max_new_tokens=max_tokens,
        do_sample=True,
        temperature=temperature,
        top_p=top_p
    )
    return tokenizer.decode(output[0], skip_special_tokens=True)

def code_prompt_from_question(question):
    return f'''You are a helpful AI programmer. Your task is to generate complete and clean code with explanations.

Task: {question}

Requirements:
- Use functions where appropriate.
- Add comments explaining each part of the code.
- Follow best coding practices.
- Do not add extra text outside the code.

Code:
'''

def try_wikipedia(query):
    try:
        wikipedia.set_lang("en")
        page = wikipedia.page(query, auto_suggest=False)
        summary = wikipedia.summary(query, auto_suggest=False)
        return summary[:2000], page.url
    except:
        try:
            summary = wikipedia.summary(query, auto_suggest=True)
            return summary[:2000], None
        except:
            return None, None

def get_top_webpages(query, max_results=3):
    urls = []
    with DDGS() as ddgs:
        results = ddgs.text(f"current {query}", max_results=max_results)
        for result in results:
            urls.append(result["href"])
    return urls

def scrape_first_good_content(urls):
    for url in urls:
        try:
            res = requests.get(url, timeout=10)
            soup = BeautifulSoup(res.text, 'html.parser')
            text = soup.get_text(separator=" ", strip=True)
            if len(text) > 300:
                return text[:2000], url
        except:
            continue
    return None, None

def smart_luna_answer(user_question, max_tokens=512):
    temperature, top_p = get_generation_params(user_question)

    if detect_mode(user_question) == "code":
        prompt = code_prompt_from_question(user_question)
        code_response = luna_generate(prompt, max_tokens=max_tokens, temperature=temperature, top_p=top_p)

        if len(code_response.strip()) < 100 or "not sure" in code_response.lower():
            urls = get_top_webpages(user_question)
            context, url_used = scrape_first_good_content(urls)
            if context:
                enhanced_prompt = f'''You are a helpful AI programmer. Use the following web content as reference.

Context: {context}

Task: {user_question}

Generate well-structured code with inline comments and clear function design. Only return the code.

Code:
'''
                code_response = luna_generate(enhanced_prompt, max_tokens=max_tokens, temperature=temperature, top_p=top_p)
                return f"Luna (code + web-enhanced):\n{code_response.strip()}\n\n(Source: {url_used})"
        return f"Luna (code):\n{code_response.strip()}"

    base_prompt = f"User: {user_question}\nLuna:"
    base_answer = luna_generate(base_prompt, max_tokens=max_tokens, temperature=temperature, top_p=top_p)

    if any(kw in base_answer.lower() for kw in ["i don't know", "not sure", "as of", "unknown"]) or len(base_answer) < 20:
        wiki_summary, wiki_url = try_wikipedia(user_question)
        if wiki_summary:
            prompt = f'''Use the following Wikipedia info to answer:

Context: {wiki_summary}

Question: {user_question}
Answer:'''
            wiki_answer = luna_generate(prompt, max_tokens=max_tokens, temperature=temperature, top_p=top_p)
            return f"Luna (Wikipedia): {wiki_answer.strip()}\n\n(Source: {wiki_url or 'Wikipedia'})"

        urls = get_top_webpages(user_question)
        context, url_used = scrape_first_good_content(urls)
        if context:
            prompt = f'''Use the following context to answer:

Context: {context}

Question: {user_question}
Answer:'''
            web_answer = luna_generate(prompt, max_tokens=max_tokens, temperature=temperature, top_p=top_p)
            return f"Luna (web-enhanced): {web_answer.strip()}\n\n(Source: {url_used})"

        return "Luna: I couldn’t find good info online right now."

    return f"Luna: {base_answer.strip()}"