import os
import json
import time
import threading
import uvicorn
from fastapi import FastAPI
from fastapi.responses import HTMLResponse
from unsloth import FastLanguageModel

# 1. IMPORT YOUR NATIVE CLIENT (No more broken OpenEnv imports!)
from envs.free_guy.client import FreeGuyEnv
from envs.free_guy.models import FreeGuyAction, FreeGuyObservation, FreeGuyState

# ==========================================
# 2. CLOUD MODEL LOADING
# ==========================================
print("🧠 Downloading & Loading Trained LoRA Model from Hugging Face...")
hf_token = os.getenv("HF_TOKEN") 

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "NikhilKhatri/FreeGuy-Combat-LoRA", # Your model!
    max_seq_length = 512,
    dtype = None,
    load_in_4bit = True,
    token = hf_token 
)
FastLanguageModel.for_inference(model)

# ==========================================
# 3. WEB SERVER & LOGGING SETUP
# ==========================================
app = FastAPI()
agent_logs = ["System Booting...", "Unreal Engine booting in background...", "Waiting for connection..."]

def log_msg(msg):
    global agent_logs
    print(msg)
    agent_logs.append(msg)
    if len(agent_logs) > 25:
        agent_logs.pop(0)

# ==========================================
# 4. THE BACKGROUND INFERENCE LOOP
# ==========================================
def background_ai_loop():
    log_msg("⏳ Waiting for Unreal Engine to boot and open port 8556...")
    
    # Wrap the entire connection in an infinite retry loop!
    while True:
        try:
            # Use your native FreeGuyEnv!
            with FreeGuyEnv(base_url="http://127.0.0.1:8556").sync() as env:
                log_msg("✅ Connected to Unreal Engine OpenEnv Server!")
                
                episode = 1
                while True: 
                    result = env.reset()
                    obs = getattr(result, 'observation', result)
                    ep_reward = 0.0
                    log_msg(f"🎮 --- STARTING EPISODE {episode} ---")
                    
                    while True:
                        enemy_state = "Idle"
                        if getattr(obs, 'CurrentOpponentAction', 0.0) == 2.0: enemy_state = "Flanking"
                        elif getattr(obs, 'CurrentOpponentAction', 0.0) == 3.0: enemy_state = "Attacking"
                        
                        dist = getattr(obs, 'DistanceFromEnemy', 0.0)
                        our_hp = getattr(obs, 'OurHealth', 0.0)
                        enemy_hp = getattr(obs, 'EnemyHealth', 0.0)
                        
                        obs_text = f"Distance {dist:.2f}m, Your Health {our_hp}, Enemy Health {enemy_hp}, Enemy {enemy_state}"
                        
                        prompt = f"<s>[INST] State: {obs_text}. What is your action? [/INST] "
                        inputs = tokenizer([prompt], return_tensors="pt").to("cuda")
                        outputs = model.generate(**inputs, max_new_tokens=40, use_cache=True)
                        llm_response = tokenizer.batch_decode(outputs[:, inputs.input_ids.shape[1]:], skip_special_tokens=True)[0]
                        
                        action_dict = {"W": 0.0, "A": 0.0, "S": 0.0, "D": 0.0, "Attack": 0.0}
                        try:
                            start = llm_response.find('{')
                            end = llm_response.rfind('}') + 1
                            if start != -1 and end != -1:
                                clean_json_str = llm_response[start:end]
                                parsed_json = json.loads(clean_json_str)
                                for key in action_dict.keys():
                                    raw_val = float(parsed_json.get(key, 0.0))
                                    action_dict[key] = 1.0 if raw_val > 0.0 else 0.0
                        except Exception:
                            pass 
                        
                        action = FreeGuyAction(**action_dict)
                        
                        for _ in range(30):
                            result = env.step(action)
                            obs = getattr(result, 'observation', result)
                            
                            current_reward = getattr(obs, 'reward', 0.0)
                            ep_reward += current_reward
                            
                            if getattr(result, 'done', False) or getattr(obs, 'terminated', False):
                                break
                                
                        log_msg(f"Action Taken: {action_dict} | Reward: {current_reward:.2f}")
                        
                        if getattr(result, 'done', False) or getattr(obs, 'terminated', False):
                            log_msg(f"💀 Episode {episode} Ended | Total Reward: {ep_reward:.2f}")
                            break
                    
                    episode += 1
                    time.sleep(2) 
                    
        except Exception as e:
            log_msg(f"⏳ Engine not ready yet, retrying in 5 seconds... ({e})")
            time.sleep(5) # Wait 5 seconds and try connecting again!

@app.on_event("startup")
def startup_event():
    thread = threading.Thread(target=background_ai_loop)
    thread.daemon = True
    thread.start()

# ==========================================
# 5. THE WEB UI FOR THE JUDGES
# ==========================================
@app.get("/")
def read_root():
    log_html = "<br>".join(agent_logs)
    html_content = f"""
    <html>
        <head>
            <title>Meta Hackathon: LLM Combat Agent</title>
            <style>
                body {{ background-color: #1e1e1e; color: #00ff00; font-family: monospace; padding: 20px; }}
                h1 {{ color: #ffffff; }}
                .terminal {{ background-color: #000; padding: 15px; border-radius: 5px; height: 600px; overflow-y: scroll; }}
            </style>
            <script>
                setTimeout(function(){{ window.location.reload(1); }}, 2000);
            </script>
        </head>
        <body>
            <h1>🤖 Live LLM Reinforcement Learning Agent</h1>
            <h3>Environment: Unreal Engine 5 | Model: Qwen LoRA</h3>
            <p>This space is running a headless 3D physics engine and a live LLM simultaneously.</p>
            <div class="terminal">
                {log_html}
            </div>
        </body>
    </html>
    """
    return HTMLResponse(content=html_content)

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)