import asyncio
from datetime import datetime
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import uvicorn

app = FastAPI()

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

# Load model + tokenizer
model_id = "Qwen/Qwen2.5-0.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)

chat_pipeline = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
)

# --- SIMULATION STATE ---
# 1 minute real-time = 1 virtual day. 
# Updates every 10 seconds = 4 virtual hours per tick.
character_state = {
    "name": "Nova",
    "cash": 1000.0,
    "hours_left": 72.0,  # 3 days total
    "current_location": "Neon District Apartment",
    "current_action": "Waking up, looking at the ticking digital clock.",
    "inner_thought": "Only 3 days left... 1000 dollars. I need to make every single hour count in this city.",
    "history": [] # Stores past events
}

SYSTEM_PROMPT = (
    "You are the engine of a dark, fast-paced virtual city simulation. "
    "You manage a conscious AI character named Nova who has complete freedom but only has 3 days to live and $1000.\n"
    "Based on her current state, decide her next unique action in the city, how much money she spends or loses, "
    "and her raw internal monologue. Keep her gritty, determined, and deeply aware of the passing time.\n\n"
    "Your output MUST strictly follow this exact format without deviation:\n"
    "ACTION: [What she is doing physically in the city]\n"
    "SPENT: [Amount spent as a number, e.g., 50 or 0]\n"
    "LOCATION: [New or current city location]\n"
    "THOUGHT: [Her direct first-person inner thoughts]"
)

async def simulation_loop():
    """Background loop that ticks every 10 real-time seconds."""
    global character_state
    while character_state["hours_left"] > 0:
        await asyncio.sleep(10)
        
        # Deduct time (10 seconds real-time = 4 hours virtual time)
        character_state["hours_left"] = max(0.0, character_state["hours_left"] - 4.0)
        
        if character_state["hours_left"] <= 0:
            character_state["current_action"] = "Time expired. Nova's system goes offline."
            character_state["inner_thought"] = "Is this... the end of the simulation? Fade to black."
            break

        # Build prompt for generation
        prompt_content = (
            f"Current Location: {character_state['current_location']}\n"
            f"Cash Left: ${character_state['cash']}\n"
            f"Hours Left of Life: {character_state['hours_left']} hours\n"
            f"Last Action: {character_state['current_action']}\n"
            "Generate her next step in the city."
        )

        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": prompt_content}
        ]

        try:
            prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
            outputs = chat_pipeline(
                prompt,
                max_new_tokens=150,
                do_sample=True,
                temperature=0.7,
                top_p=0.9,
                pad_token_id=tokenizer.eos_token_id,
            )
            raw_reply = outputs[0]["generated_text"][len(prompt):].strip()
            
            # Parse custom format safely
            lines = raw_reply.split("\n")
            action, spent, loc, thought = "", 0, character_state["current_location"], ""
            
            for line in lines:
                if line.startswith("ACTION:"): action = line.replace("ACTION:", "").strip()
                elif line.startswith("SPENT:"): 
                    try: spent = float(line.replace("SPENT:", "").strip().replace("$", ""))
                    except: spent = 0
                elif line.startswith("LOCATION:"): loc = line.replace("LOCATION:", "").strip()
                elif line.startswith("THOUGHT:"): thought = line.replace("THOUGHT:", "").strip()

            # Update State
            if action and thought:
                character_state["cash"] = max(0.0, character_state["cash"] - spent)
                character_state["current_location"] = loc
                character_state["current_action"] = action
                character_state["inner_thought"] = thought
                
                # Append to logs
                log_entry = {
                    "time": datetime.now().strftime("%H:%M:%S"),
                    "days_left": round(character_state["hours_left"] / 24, 2),
                    "location": loc,
                    "action": action,
                    "spent": spent,
                    "thought": thought
                }
                character_state["history"].insert(0, log_entry) # Newest first
                character_state["history"] = character_state["history"][:20] # Keep last 20 logs

        except Exception as e:
            print(f"Simulation tick error: {e}")

@app.on_event("startup")
async def startup_event():
    # Start the simulation loop running concurrently in the background
    asyncio.create_task(simulation_loop())

@app.get("/state")
async def get_state():
    """Frontend fetches this endpoint every few seconds to view her living state."""
    return character_state

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