import os
import json
import numpy as np
from dotenv import load_dotenv
import google.generativeai as genai
from server.battery_environment import BatteryEnvironment
from models import BatteryAction

# Load .env
load_dotenv()
API_KEY = os.getenv("GEMINI_API_KEY")

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 run_baseline(mode, model_name="gemini-1.5-flash"):
    if not API_KEY:
        print("GEMINI_API_KEY not found in .env. Falling back to Rule-Based logic.")
        return 0.0

    genai.configure(api_key=API_KEY)
    model = genai.GenerativeModel(model_name)
    
    env = BatteryEnvironment(mode=mode)
    obs = env.reset()
    
    total_reward = 0.0
    done = False
    step_count = 0
    max_steps = 96
    
    print(f"\n--- Starting Task: {mode.upper()} ---")
    
    while not done and step_count < max_steps:
        prompt = (f"You are a battery management agent. Maximize profit. "
                  f"Return action as JSON: {{'market_choice': [0.0-1.0], 'p_fraction': [-1.0-1.0]}}. "
                  f"Current State: energy={obs.energy_price}, fcl={obs.fcr_price}, soc={obs.soc}")
        
        try:
            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 Exception as e:
            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"Task '{mode}' completed. Total Reward: {total_reward:.2f}. Score: {score:.2f} / 1.00")
    return score

if __name__ == "__main__":
    scores = {}
    for task_mode in ["easy", "medium", "hard"]:
        score = run_baseline(task_mode)
        scores[task_mode] = score
        
    print("\n=== Final Gemini Baseline Scores ===")
    for task_mode, score in scores.items():
        print(f"{task_mode.capitalize()}: {score:.2f}")