import os
import sys
import time
import json
import re
import shutil
import subprocess
import urllib.request
import urllib.error
import urllib.parse
from datetime import datetime
from http.server import HTTPServer, BaseHTTPRequestHandler

# ==========================================
# CONFIGURATION (Official Qwen 3.5 2B Instruct)
# ==========================================
#MODEL_REPO = "bartowski/Qwen_Qwen3.5-2B-GGUF"
#MODEL_FILE = "Qwen_Qwen3.5-2B-Q4_K_M.gguf"
MODEL_REPO = "bartowski/Qwen_Qwen3.5-0.8B-GGUF"
MODEL_FILE = "Qwen_Qwen3.5-0.8B-Q4_K_M.gguf"
MODEL_PATH = f"models/{MODEL_FILE}"

PROXY_HOST = "0.0.0.0"             
PROXY_PORT = 7860                  
INTERNAL_LLAMA_PORT = 8081         

SEARCH_TRIGGERS = [
    "search", "weather", "news", "latest", "current", "today", "tody",
    "google", "internet", "up to date", "who is", "what is happening", "date"
]

# ==========================================
# 1. AUTOMATIC MODEL DOWNLOADER
# ==========================================
os.makedirs("models", exist_ok=True)
if not os.path.exists(MODEL_PATH):
    print("⬇️ Downloading Gemma 4 E2B-IT model from Hugging Face...")
    from huggingface_hub import hf_hub_download
    downloaded = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
    shutil.copy2(downloaded, MODEL_PATH)
    print("✅ Model successfully cached locally.")

# ==========================================
# 2. ENGINE LAUNCH (Gemma 4 Optimized Flags)
# ==========================================
def launch_optimized_llama():
    possible_bins = [
        "./llama.cpp/build/bin/llama-server",
        "./llama-server",
        "llama-server"
    ]
    
    llama_bin = None
    for path in possible_bins:
        if shutil.which(path) or os.path.exists(path):
            llama_bin = path
            break
            
    if not llama_bin:
        print("❌ Error: Could not find 'llama-server' executable.")
        sys.exit(1)

    print(f"🚀 Found backend binary at: {llama_bin}")
    
    cmd = [
        llama_bin,
        "-m", MODEL_PATH,
        "-c", "4096",               # Perfect context size for your 16GB RAM footprint
        "-t", "2",                  # Matches your 2 physical vCPUs perfectly
        "-tb", "2",                 # Match batch processing to your hardware threads
        "--prio", "2",              # Sets high process priority so the OS handles text streams smoothly
        "--mlock",
        "--flash-attn", "on",       # Enforces active structural execution paths
        "--temp", "1.0",            # Cleaned duplicate; restored Gemma 4 baseline calibration
        "--top-p", "0.95",          
        "--top-k", "64",            
        "--batch-size", "256",      # Raised slightly to process prompts faster on initial load
        "--ubatch-size", "64",
        "--parallel", "1",
        "--repeat-penalty", "1.1",
        "--jinja",
        "--chat-template-kwargs", '{"enable_thinking":false}',
        "--port", str(INTERNAL_LLAMA_PORT),
        "--host", "127.0.0.1"
    ]
    
    process = subprocess.Popen(cmd, stdout=None, stderr=None)
    
    print("⏳ Warming up model weights...")
    for _ in range(30):
        if process.poll() is not None:
            print(f"❌ Core engine exited prematurely with code {process.returncode}.")
            raise RuntimeError("Internal llama-server failed to launch. Verify your llama.cpp build supports Gemma 4 architectures.")
            
        try:
            req = urllib.request.Request(f"http://127.0.0.1:{INTERNAL_LLAMA_PORT}/health")
            with urllib.request.urlopen(req, timeout=2) as r:
                if r.status == 200:
                    print(f"✅ Core inference engine live on port {INTERNAL_LLAMA_PORT}")
                    return process
        except:
            time.sleep(2)
            
    raise RuntimeError("Internal llama-server timed out during warmup initialization.")

# ==========================================
# 3. INTERNET SEARCH TOOL
# ==========================================
def execute_web_search(query):
    print(f"🌐 Sourcing live web data for: '{query}'")
    try:
        url = f"https://html.duckduckgo.com/html/?q={urllib.parse.quote(query)}"
        req = urllib.request.Request(
            url, 
            headers={'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64)'}
        )
        with urllib.request.urlopen(req, timeout=5) as response:
            html = response.read().decode('utf-8')
            snippets = re.findall(r'<a class="result__snippet"[^>]*>(.*?)</a>', html, re.DOTALL)[:3]
            clean_results = []
            for item in snippets:
                clean_text = re.sub(r'<[^>]+>', '', item).strip()
                if clean_text:
                    clean_results.append(clean_text)
            
            if clean_results:
                return "\n".join([f"- {txt}" for txt in clean_results])
    except Exception as e:
        print(f"⚠️ Internet lookup bypassed: {e}")
    return "No search results returned."

# ==========================================
# 4. PASSTHROUGH PROXY GATEWAY WITH THINKING INJECTION
# ==========================================
class OpenAICompatibilityProxy(BaseHTTPRequestHandler):
    def log_message(self, format, *args):
        pass  

    def handle_proxy(self, method, data=None):
        url = f"http://127.0.0.1:{INTERNAL_LLAMA_PORT}{self.path}"
        req_headers = {}
        for k, v in self.headers.items():
            if k.lower() not in ['host', 'content-length']:
                req_headers[k] = v

        req = urllib.request.Request(url, data=data, headers=req_headers, method=method)
        
        try:
            with urllib.request.urlopen(req) as response:
                self.send_response(response.status)
                for key, val in response.getheaders():
                    if key.lower() not in ['transfer-encoding', 'connection', 'content-length', 'content-encoding']:
                        self.send_header(key, val)
                self.end_headers()
                
                while True:
                    chunk = response.read(2048)
                    if not chunk:
                        break
                    self.wfile.write(chunk)
                    self.wfile.flush()
                    
        except urllib.error.HTTPError as e:
            self.send_response(e.code)
            for key, val in e.headers.items():
                if key.lower() not in ['transfer-encoding', 'connection', 'content-length']:
                    self.send_header(key, val)
            self.end_headers()
            self.wfile.write(e.read())
        except Exception as e:
            self.send_response(500)
            self.end_headers()
            self.wfile.write(json.dumps({"error": str(e)}).encode('utf-8'))

    def do_GET(self):
        self.handle_proxy("GET")

    def do_POST(self):
        content_length = int(self.headers.get('Content-Length', 0))
        body_bytes = self.rfile.read(content_length) if content_length > 0 else b''
        
        forward_data = body_bytes
        intercept_paths = ["/v1/chat/completions", "/chat/completions", "/completion"]
        
        if content_length > 0 and self.path in intercept_paths:
            try:
                body = json.loads(body_bytes.decode('utf-8'))
                
                last_user_message = ""
                if "messages" in body:
                    for msg in reversed(body.get("messages", [])):
                        if msg.get("role") == "user":
                            last_user_message = msg.get("content", "")
                            break
                elif "prompt" in body:
                    last_user_message = body.get("prompt", "")

                current_time_str = datetime.now().strftime("%A, %B %d, %Y")
                
                # Injected <|think|> trigger tells Gemma 4 to activate reasoning chains natively
                base_system_context = f"<|think|>\n[SYSTEM INFO: Current real-world date is {current_time_str}. User Location: Chennai, India.]"

                web_context = ""
                if last_user_message and any(t in last_user_message.lower() for t in SEARCH_TRIGGERS):
                    web_context = execute_web_search(last_user_message)
                    base_system_context += f"\n[Live Web Results found for query:\n{web_context}]"
                
                if "messages" in body:
                    if body["messages"] and body["messages"][0].get("role") == "system":
                        body["messages"][0]["content"] = base_system_context + "\n\n" + body["messages"][0]["content"]
                    else:
                        body["messages"].insert(0, {"role": "system", "content": base_system_context})
                elif "prompt" in body:
                    body["prompt"] = f"{base_system_context}\n\nUser: {body['prompt']}"
                    
                forward_data = json.dumps(body).encode('utf-8')
            except Exception as e:
                print(f"⚠️ Text preprocessing bypassed: {e}")
                forward_data = body_bytes

        self.handle_proxy("POST", data=forward_data)

def run_server():
    server = HTTPServer((PROXY_HOST, PROXY_PORT), OpenAICompatibilityProxy)
    print(f"🪐 High-Accuracy Interface active at: http://{PROXY_HOST}:{PROXY_PORT}")
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\n🛑 Closing application layers safely...")
    finally:
        server.server_close()

if __name__ == "__main__":
    backend_proc = None
    try:
        backend_proc = launch_optimized_llama()
        run_server()
    finally:
        if backend_proc:
            backend_proc.terminate()
            backend_proc.wait()