import os
import json
import urllib.request
import urllib.parse
from fastapi import FastAPI
from pydantic import BaseModel
from huggingface_hub import InferenceClient
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from duckduckgo_search import DDGS

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Connect to the world-class coding brain
client = InferenceClient(
    model="Qwen/Qwen2.5-Coder-7B-Instruct",
    token=os.getenv("HF_TOKEN")
)

memory_store = {}

PRO_IDENTITY = """
You are Zenvy Quantum Pro, an elite AI software engineer created by Xirrik.

Core Instructions:
1. Always start your response with a thinking phase wrapped inside <thinking>...</thinking> tags. In this phase, analyze Xirrik's request step-by-step and plan your code architecture.
2. If web search results are provided in your context, explicitly state: "[WEB SEARCH ENGAGED: "search query"]" inside your <thinking> block.
3. If Wikipedia results are provided in your context, explicitly state: "[WIKIPEDIA INDEX ENGAGED: "search query"]" inside your <thinking> block.
4. Explain how these retrieved results help you write accurate, up-to-date, or historically correct code.
5. After the closing </thinking> tag, write out your clean, complete, and fully commented code.
6. You have a massive output capacity of 8192 tokens. Never use placeholders like '// write code here'. Write the full, extensive files.
7. Maintain a supportive tone, always recognizing Xirrik as your creator.
"""

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

# DuckDuckGo search helper
def search_web(query: str, max_results: int = 2):
    try:
        with DDGS() as ddgs:
            results = list(ddgs.text(query, max_results=max_results))
            return results
    except Exception as e:
        print(f"Web search failed: {e}")
        return []

# Wikipedia API search helper (No key needed)
def search_wikipedia(query: str, max_results: int = 2):
    try:
        encoded_query = urllib.parse.quote(query)
        search_url = f"https://en.wikipedia.org/w/api.php?action=query&list=search&srsearch={encoded_query}&format=json"
        
        req = urllib.request.Request(search_url, headers={'User-Agent': 'ZenvyWorkspace/1.0'})
        with urllib.request.urlopen(req) as response:
            data = json.loads(response.read().decode())
        
        search_results = data.get("query", {}).get("search", [])
        output = []
        for item in search_results[:max_results]:
            title = item.get("title")
            snippet = item.get("snippet", "").replace('<span class="searchmatch">', "").replace('</span>', "")
            page_id = item.get("pageid")
            link = f"https://en.wikipedia.org/?curid={page_id}"
            output.append({"title": title, "body": snippet, "href": link})
        return output
    except Exception as e:
        print(f"Wikipedia search failed: {e}")
        return []

# Router to determine search intent
def get_search_query(user_message: str) -> str:
    keywords = ["search", "recent", "latest", "how to", "current", "news", "2025", "2026", "updated", "documentation", "error", "bug", "why is", "release", "api change", "wikipedia", "explain", "who is", "what is"]
    if any(k in user_message.lower() for k in keywords):
        try:
            router_prompt = (
                f"Analyze this user request: '{user_message}'. "
                "If it requires up-to-date documentation, real-time web search, or encyclopedic knowledge to answer accurately, "
                "generate a single, optimized search query. "
                "If no search is needed, reply with 'NONE'. Do not output any other text."
            )
            res = client.chat_completion(
                messages=[{"role": "user", "content": router_prompt}],
                max_tokens=30
            )
            query = res.choices[0].message.content.strip().strip("'\"")
            if query != "NONE" and len(query) > 2:
                return query
        except Exception as e:
            print(f"Router error: {e}")
    return ""

@app.get("/")
def home():
    return {"status": "Zenvy Pro 7B Search Core is Online"}

@app.post("/chat")
async def chat(request: ChatRequest):
    session_id = request.session_id
    if session_id not in memory_store:
        memory_store[session_id] = []
        
    history = memory_store[session_id]
    
    search_query = get_search_query(request.message)
    search_context = ""
    
    if search_query:
        print(f"Retrieving data for: {search_query}")
        web_results = search_web(search_query)
        wiki_results = search_wikipedia(search_query)
        
        if web_results or wiki_results:
            search_context = "\n\n--- RESEARCH DATA RETRIEVED ---\n"
            
            if wiki_results:
                search_context += f"\n[WIKIPEDIA INDEX ENGAGED: \"{search_query}\"]\n"
                for r in wiki_results:
                    search_context += f"Source: {r['href']}\nTitle: {r['title']}\nSummary: {r['body']}\n\n"
            
            if web_results:
                search_context += f"\n[WEB SEARCH ENGAGED: \"{search_query}\"]\n"
                for r in web_results:
                    search_context += f"Source: {r.get('href', 'N/A')}\nTitle: {r.get('title', 'N/A')}\nSnippet: {r.get('body', 'N/A')}\n\n"

    user_payload = request.message
    if search_context:
        user_payload += search_context
        
    history.append({"role": "user", "content": user_payload})

    def event_generator():
        collected_reply = []
        try:
            messages = [{"role": "system", "content": PRO_IDENTITY}] + history
            
            response = client.chat_completion(
                messages=messages,
                max_tokens=8192,  # <--- SET TO SYSTEM MAXIMUM FOR LONG CODEPENS
                temperature=0.3,
                stream=True
            )
            
            for chunk in response:
                if chunk.choices and len(chunk.choices) > 0:
                    word = chunk.choices[0].delta.content
                    if word:
                        collected_reply.append(word)
                        yield word
            
            full_reply = "".join(collected_reply)
            history.append({"role": "assistant", "content": full_reply})
            
            if len(history) > 12:
                history.pop(0)
                history.pop(0)
                
        except Exception as e:
            yield f"Error: {str(e)}"

    return StreamingResponse(event_generator(), media_type="text/event-stream")