import os
import json
import numpy as np
import pandas as pd
import time
import warnings
from dotenv import load_dotenv
import google.generativeai as genai

# Load API Key from .env
load_dotenv()
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")

warnings.filterwarnings("ignore")

from server.battery_environment import BatteryEnvironment
from models import BatteryAction, BatteryObservation

def grade_trajectory(total_reward, mode):
    max_theoretical = {"easy": 5000.0, "medium": 2000.0, "hard": 500.0}
    target = max_theoretical.get(mode, 1000.0)
    score = max(0.0, min(1.0, total_reward / target))
    return score

def rule_based_action(obs, mode="medium", step=0):
    target_soc = 15.0 if mode == "hard" else 10.0
    if step > (96 - 20):
        if obs.soc < target_soc - 0.5: return BatteryAction(market_choice=0.0, p_fraction=-1.0)
        elif obs.soc > target_soc + 0.5: return BatteryAction(market_choice=0.0, p_fraction=0.5)
    
    # Better FCR threshold (50 in Easy mode, 25 in Medium/Hard mode)
    fcr_threshold = 40.0 # 0.5 * fcr_price > 40 means fcr_price > 80
    if obs.fcr_price * 0.5 > fcr_threshold: 
        return BatteryAction(market_choice=0.8, p_fraction=0.0)
    
    if obs.energy_price < 40.0: return BatteryAction(market_choice=0.0, p_fraction=-1.0)
    elif obs.energy_price > 120.0: return BatteryAction(market_choice=0.0, p_fraction=1.0)
    return BatteryAction(market_choice=0.0, p_fraction=0.0)

def run_evaluation(mode, agent_type="rule_based", limit_steps=96):
    env = BatteryEnvironment(mode=mode)
    obs = env.reset()
    total_reward, done, step_count = 0.0, False, 0
    print(f"[{agent_type}] Task: {mode}... ", end="", flush=True)
    
    if agent_type == "gemini":
        if not GEMINI_API_KEY: return None
        genai.configure(api_key=GEMINI_API_KEY)
        model = genai.GenerativeModel("gemini-1.5-flash")

    while not done and step_count < limit_steps:
        if agent_type == "rule_based":
            action = rule_based_action(obs, mode, step_count)
        elif agent_type == "gemini":
            try:
                time.sleep(1.0) # Reduced sleep
                prompt = f"Battery agent. Return JSON: {{'market_choice': [0,1], 'p_fraction': [-1,1]}}. E:{obs.energy_price}, F:{obs.fcr_price}, S:{obs.soc}"
                response = model.generate_content(prompt, generation_config={"response_mime_type": "application/json"})
                action_json = json.loads(response.text)
                action = BatteryAction(market_choice=float(action_json.get("market_choice", 0.0)), p_fraction=float(action_json.get("p_fraction", 0.0)))
            except: action = BatteryAction(market_choice=0.0, p_fraction=0.0)
        else: action = BatteryAction(market_choice=0.0, p_fraction=0.0)

        obs = env.step(action)
        total_reward += obs.reward or 0.0
        done, step_count = obs.done, step_count + 1
        
    score = grade_trajectory(total_reward, mode)
    print(f"Done. Score: {score:.2f}")
    return score

if __name__ == "__main__":
    print("\nProject Baselines (Rule-Based Algorithm):")
    for task in ["easy", "medium", "hard"]:
        run_evaluation(task, "rule_based")
    
    print("\nGemini 1.5 Flash (Sampling 5 steps):")
    for task in ["easy", "medium", "hard"]:
        run_evaluation(task, "gemini", limit_steps=5)