import os
import asyncio
import logging
import json
import torch
import numpy as np
import httpx
import uvicorn
from contextlib import asynccontextmanager
from urllib.parse import urlparse
from bs4 import BeautifulSoup
from trafilatura import extract
from duckduckgo_search import DDGS
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig

# Enforce clean production log formats
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger("infinitygrm")

# Thread-safe global state mapping dictionary for engine assets
engine_state = {}

# Temporary in-memory session storage
# Replace with Redis later for persistent sessions
session_store = {}

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Securely handles neural weights allocation on Hugging Face hardware targets."""
    logger.info("Initializing neural weights and engine dependencies on L4 GPU...")

    device = "cuda" if torch.cuda.is_available() else "cpu"
    # L4 GPU easily hosts 3B models in float16 precision consuming only ~6-7GB VRAM
    dtype = torch.float16 if device == "cuda" else torch.float32
    model_name = "Qwen/Qwen2.5-3B-Instruct"

    tokenizer = AutoTokenizer.from_pretrained(model_name)
    tokenizer.pad_token = tokenizer.eos_token
    config = AutoConfig.from_pretrained(model_name)
    config.pad_token_id = tokenizer.eos_token_id

    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        config=config,
        torch_dtype=dtype,
        device_map="auto" if device == "cuda" else None
    ).eval()

    embed_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", device=device)

    engine_state["tokenizer"] = tokenizer
    engine_state["model"] = model
    engine_state["embed_model"] = embed_model
    engine_state["max_history"] = 15

    logger.info(f"Engine layers safely mapped to execution target: {device}")
    yield
    engine_state.clear()
    logger.info("Application context fully flushed.")

app = FastAPI(title="Infinity GRM Engine", lifespan=lifespan)

# --- PRODUCTION CORS SETUP FOR WORDPRESS ---
# Allows your WordPress site frontend to make asynchronous browser requests to Hugging Face
# Replace lines 72 to 81 completely with this clean configuration:
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class ChatRequest(BaseModel):
    message: str
    session_id: str

def fetch_or_create_session(session_id: str) -> dict:
    """Fetches user session from in-memory storage."""
    
    if session_id in session_store:
        return session_store[session_id]

    default_session = {
        "history": [],
        "last_entity": "",
        "last_web_query": "",
        "last_response": "",
        "last_web_context": []
    }

    session_store[session_id] = default_session
    return default_session

def save_session_to_redis(session_id: str, session: dict):
    """Saves session to in-memory storage."""
    session_store[session_id] = session

def update_history_buffer(session: dict, role: str, content: str):
    session["history"].append({"role": role, "content": content})
    if len(session["history"]) > engine_state["max_history"]:
        session["history"].pop(0)

def calculate_semantic_similarity(a: str, b: str) -> float:
    embedder = engine_state["embed_model"]
    embeddings = embedder.encode([a, b], normalize_embeddings=True)
    return float(np.dot(embeddings[0], embeddings[1]))

def checks_pronoun_reference(query: str) -> bool:
    pronoun_references = {"he", "him", "his", "she", "her", "hers", "they", "them", "their", "it", "its", "that", "those", "this", "these", "then", "next", "after"}
    return any(word in pronoun_references for word in query.lower().split())

def is_followup(query: str, session: dict) -> bool:
    if not session["last_entity"]:
        return False
    scores = []
    try:
        scores.append(calculate_semantic_similarity(query, session["last_entity"]))
    except Exception:
        pass
    try:
        if session["last_response"]:
            scores.append(calculate_semantic_similarity(query, session["last_response"][:500]))
    except Exception:
        pass
    best_match = max(scores, default=0)
    return (best_match > 0.45 or len(query.split()) <= 3 or checks_pronoun_reference(query))

def execute_llm_generation(prompt: str) -> str:
    tokenizer = engine_state["tokenizer"]
    model = engine_state["model"]
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        out = model.generate(
            **inputs,
            max_new_tokens=2048,
            temperature=0.15,
            top_p=0.9,
            repetition_penalty=1.1
        )
    return tokenizer.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True).strip()

def run_topic_extraction(query: str) -> str:
    prompt = f"Extract the primary subject.\nQuery: {query}\nReturn only the subject.\n"
    return execute_llm_generation(prompt).split("\n")[0].strip()

def rewrite_followup_query(query: str, session: dict) -> str:
    history = "\n".join([f"User: {x['user']}\nSearch: {x['search']}" for x in session["last_web_context"][-5:]])
    prompt = f"Entity: {session['last_entity']}\nRecent conversation:\n{history}\nUser follow-up: {query}\nRewrite as a standalone web search query. Return only the query.\n"
    return execute_llm_generation(prompt).split("\n")[0].strip()

def should_trigger_web_search(q: str, session: dict) -> bool:
    q = q.lower()
    search_triggers = ["who", "what", "when", "where", "why", "how", "founder", "ceo", "company", "price", "stock", "news", "latest", "code", "python", "html"]
    if any(trigger in q for trigger in search_triggers):
        return True
    return is_followup(q, session)

def isolated_blocking_ddg(query_string: str) -> list:
    try:
        with DDGS() as ddgs:
            results = list(ddgs.text(query_string, max_results=5))
            return [r.get("href") for r in results if r.get("href", "").startswith("http")]
    except Exception as e:
        logger.error(f"DuckDuckGo engine wrapper error: {str(e)}")
        return []

async def search_web_async(query: str) -> list:
    urls = await asyncio.to_thread(isolated_blocking_ddg, query)
    if not urls:
        urls.append(f"https://wikipedia.org{query.replace(' ', '_')}")
    return list(dict.fromkeys(urls))[:5]

async def scrape_target_url(client: httpx.AsyncClient, url: str) -> dict:
    try:
        r = await client.get(url, timeout=10, follow_redirects=True)
        if r.status_code != 200:
            return None
        text = extract(r.text)
        if not text:
            soup = BeautifulSoup(r.text, "html.parser")
            text = soup.get_text(" ", strip=True)
        if not text or len(text.split()) < 20:
            return None
        return {"url": url, "domain": urlparse(url).netloc, "text": text[:2000]}
    except Exception:
        return None

async def build_context_pipeline(query: str) -> tuple:
    target_urls = await search_web_async(query)
    async with httpx.AsyncClient(headers={"User-Agent": "Mozilla/5.0 Production Engine"}) as client:
        tasks = [scrape_target_url(client, u) for u in target_urls]
        scraped_results = await asyncio.gather(*tasks)
    context_blocks, extraction_sources = [], []
    for result in scraped_results:
        if result:
            context_blocks.append(f"[{result['domain']}] {result['text']}")
            extraction_sources.append(result["url"])
    return "\n\n".join(context_blocks), list(set(extraction_sources))

async def process_ask_orchestration(q: str, session: dict) -> tuple:
    use_web = should_trigger_web_search(q, session)
    search_query = q
    if use_web and is_followup(q, session) and session["last_entity"]:
        search_query = rewrite_followup_query(q, session)
    context, sources = "", []
    if use_web:
        if not is_followup(q, session):
            try:
                session["last_entity"] = run_topic_extraction(q)
            except Exception:
                session["last_entity"] = q
        context, sources = await build_context_pipeline(search_query)
        if len(context.strip()) < 150:
            context, sources = await build_context_pipeline(search_query + " wikipedia")
        session["last_web_context"].append({"user": q, "search": search_query})
        if len(session["last_web_context"]) > 5:
            session["last_web_context"].pop(0)
        session["last_web_query"] = search_query
        messages = [{"role": "system", "content": "You are a detailed research assistant. Provide a structured comprehensive answer without bullet points. Always use provided web context in your answer."}]
        messages.extend(session["history"][-15:])
        messages.append({"role": "user", "content": f"WEB CONTEXT:\n{context}\n\nQUESTION:\n{q}"})
    else:
        messages = [{"role": "system", "content": "You are a helpful assistant."}]
        messages.extend(session["history"][-6:])
        messages.append({"role": "user", "content": q})

    compiled_prompt = engine_state["tokenizer"].apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    response_string = execute_llm_generation(compiled_prompt)
    session["last_response"] = response_string
    update_history_buffer(session, "user", q)
    update_history_buffer(session, "assistant", response_string)
    return response_string, sources

# --- API ENDPOINTS ---
@app.post("/query")
async def chat_endpoint(request: ChatRequest):
    try:
        session = fetch_or_create_session(request.session_id)
        answer, sources = await process_ask_orchestration(request.message, session)
        save_session_to_redis(request.session_id, session)
        return {"response": answer, "sources": sources}

    except Exception as e:
        logger.error(f"Critical exception inside chat runtime pipeline: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail=str(e)
        )

@app.get("/health")
async def validation_heartbeat():
    return {"status": "healthy", "gpu_acceleration_active": torch.cuda.is_available()}

if not os.path.exists("static"):
    os.makedirs("static")
app.mount("/static", StaticFiles(directory="static"), name="static")

@app.get("/")
async def home_route_processor():
    return {"message": "Infinity GRM Engine Core is active."}

# --- HUGGING FACE DOCKER PORT BINDING ---
if __name__ == "__main__":
    # Hugging Face Spaces mandates listening exclusively on port 7860
    port = int(os.getenv("PORT", 7860))
    uvicorn.run("main:app", host="0.0.0.0", port=port)