#!/usr/bin/env python3
import unittest
import requests
import os
import json
import psycopg2
import sys
import time
from dotenv import load_dotenv

# --- Configuration ---
# Load environment variables from .env file in the same directory
load_dotenv()

# Configuration with fallbacks
PORT = os.getenv('PORT', '8002')
BASE_URL = f"http://localhost:{PORT}"
API_USER = os.getenv('API_USER', 'jpeters')
API_PASSWORD = os.getenv('API_PASSWORD', 'WellJson')

# ANSI Colors for better readability
GREEN = '\033[92m'
RED = '\033[91m'
YELLOW = '\033[93m'
RESET = '\033[0m'

class TestWellDrySenseAPI(unittest.TestCase):
    """
    Test suite for WellDrySense.
    Ensures zero-impact on existing data by cleaning up created records.
    """
    
    token = None
    user_id = None
    job_id_to_test = None
    db_conn_params = {}

    @classmethod
    def setUpClass(cls):
        print(f"\n{GREEN}=== Setting up WellDrySense Test Suite on Port {PORT} ==={RESET}")
        
        # Setup DB Params
        cls.db_conn_params = {
            'dbname': os.getenv('DB_NAME'),
            'user': os.getenv('DB_USER'),
            'password': os.getenv('DB_PASSWORD'),
            'host': os.getenv('DB_HOST'),
            'port': os.getenv('DB_PORT')
        }

        # Authenticate
        print(f"-> Logging in as: {API_USER}...")
        url = f"{BASE_URL}/api/well_api"
        payload = {
            "function": "credentials",
            "user_name": API_USER,
            "ps": API_PASSWORD,
            "clientId": "test-suite",
            "nonce": "test-nonce"
        }

        try:
            response = requests.post(url, data=payload)
            if response.status_code != 200:
                print(f"{RED}FATAL: Login failed. Status: {response.status_code}{RESET}")
                print(f"Response: {response.text}")
                sys.exit(1)

            data = response.json()
            
            # Handle different response structures
            if 'access_token' in data:
                cls.token = data['access_token']
                cls.user_id = data.get('user_id')
            elif 'data' in data and 'access_token' in data['data']:
                cls.token = data['data']['access_token']
                cls.user_id = data['data'].get('user_id')
            else:
                print(f"{RED}FATAL: Token not found in response.{RESET}")
                sys.exit(1)

            print(f"{GREEN}-> Login successful. User ID: {cls.user_id}{RESET}")

        except requests.exceptions.ConnectionError:
            print(f"{RED}FATAL: Could not connect to {BASE_URL}. Ensure well-api.py is running.{RESET}")
            sys.exit(1)

    @classmethod
    def tearDownClass(cls):
        print(f"\n{GREEN}=== Tearing Down Test Suite ==={RESET}")
        if cls.job_id_to_test:
            print(f"-> Cleaning up Job ID: {cls.job_id_to_test}...")
            conn = None
            try:
                conn = psycopg2.connect(**cls.db_conn_params)
                with conn.cursor() as cur:
                    cur.execute("DELETE FROM public.jobs WHERE job_id = %s;", (cls.job_id_to_test,))
                    conn.commit()
                    print(f"{GREEN}-> Cleanup successful. Database restored.{RESET}")
            except Exception as e:
                print(f"{RED}CRITICAL: DB Cleanup failed. Manually delete job {cls.job_id_to_test}. Error: {e}{RESET}")
            finally:
                if conn: conn.close()
        else:
            print("-> No job created, skipping cleanup.")

    def _post_api(self, form_data):
        """Helper to send authenticated POST requests"""
        form_data['user_name'] = API_USER
        form_data['token'] = self.token
        
        try:
            response = requests.post(f"{BASE_URL}/api/well_api", data=form_data)
            self.assertEqual(response.status_code, 200, f"API HTTP Error {response.status_code}: {response.text}")
            
            try:
                json_resp = response.json()
                if 'data' in json_resp and 'status' in json_resp:
                     return json_resp['data']
                return json_resp
            except json.JSONDecodeError:
                self.fail(f"API returned invalid JSON: {response.text}")
        except requests.exceptions.ConnectionError:
            self.fail("Connection refused. API server is down.")

    # --- TESTS ---

    def test_01_create_job(self):
        """Test creating a new job"""
        print("\n[Test] job_create")
        payload = {
            "function": "job_create",
            "customer_name": "TEST_SUITE_CUSTOMER",
            "address_street": "123 Python Way",
            "address_city": "Codeville",
            "address_state": "CA",
            "address_country": "USA",
            "lat": 34.05,
            "lng": -118.25,
            "key_person_name": "Test Runner",
            "key_person_email": "test@wellnuo.com",
            "devices": json.dumps([{"mac": "TEST_MAC_VIRTUAL", "location": "Lab"}]),
            "alerts_config": json.dumps({"temp_high": 30})
        }

        data = self._post_api(payload)
        self.assertEqual(data.get('ok'), 1, f"Job creation failed: {data.get('error')}")
        self.assertIn('job_id', data)
        self.__class__.job_id_to_test = data['job_id']
        print(f"-> Job created with ID: {self.job_id_to_test}")

    def test_02_job_list(self):
        """Test retrieving the job list"""
        print("\n[Test] job_list")
        payload = {"function": "job_list"}
        data = self._post_api(payload)
        
        self.assertEqual(data.get('ok'), 1, f"List failed: {data.get('error')}")
        self.assertIn('jobs', data)
        
        found = any(j.get('job_id') == self.job_id_to_test for j in data['jobs'])
        self.assertTrue(found, f"Created Job ID {self.job_id_to_test} not found in job_list")

    def test_03_job_details(self):
        """Test retrieving single job details"""
        print("\n[Test] job_details")
        if not self.job_id_to_test: self.skipTest("No job ID available")

        payload = {"function": "job_details", "job_id": self.job_id_to_test}
        data = self._post_api(payload)

        self.assertEqual(data.get('ok'), 1, f"Details failed: {data.get('error')}")
        self.assertEqual(data['details']['customer_name'], "TEST_SUITE_CUSTOMER")
        
        # Verify JSON parsing of devices
        devices = data['details'].get('devices')
        if isinstance(devices, str): devices = json.loads(devices)
        self.assertEqual(devices[0]['mac'], "TEST_MAC_VIRTUAL")

    def test_04_job_edit(self):
        """Test updating a job (Stop the job)"""
        print("\n[Test] job_edit")
        if not self.job_id_to_test: self.skipTest("No job ID available")

        payload = {
            "function": "job_edit",
            "job_id": self.job_id_to_test,
            "customer_name": "UPDATED_CUSTOMER_NAME",
            "job_status": "Stopped",
            "date_to": "2025-12-31T23:59:59" 
        }
        data = self._post_api(payload)
        self.assertEqual(data.get('ok'), 1, f"Edit failed: {data.get('error')}")

        # Verify
        v_payload = {"function": "job_details", "job_id": self.job_id_to_test}
        v_data = self._post_api(v_payload)
        self.assertEqual(v_data['details']['customer_name'], "UPDATED_CUSTOMER_NAME")
        self.assertEqual(v_data['details']['job_status'], "Stopped")

    def test_05_available_devices(self):
        """Test fetching available devices"""
        print("\n[Test] job_available_devices")
        payload = {"function": "job_available_devices"}
        data = self._post_api(payload)
        self.assertEqual(data.get('ok'), 1, f"Available devices failed: {data.get('error')}")
        self.assertIsInstance(data['devices'], list)

    def test_06_job_weather(self):
        """Test fetching weather"""
        print("\n[Test] job_weather")
        if not self.job_id_to_test: self.skipTest("No job ID available")

        #if not os.getenv('WEATHER_API_KEY'): print(f"{YELLOW}-> Warning: WEATHER_API_KEY not found in .env{RESET}")

        payload = {"function": "job_weather", "job_id": self.job_id_to_test}
        data = self._post_api(payload)
        
        if data.get('ok') == 0:
            print(f"-> Weather API returned error (Expected if key invalid): {data.get('error')}")
        else:
            self.assertIn('weather', data)
            print(f"-> Weather received: {data['weather']}")
    
    def test_07_job_sensor_data(self):
        """Test fetching bucketed sensor data for a job"""
        print("\n[Test] get_job_sensor_bucketed_data")
        if not self.job_id_to_test: self.skipTest("No job ID available")

        # Use a date range likely to cover "now" or recent past for testing
        today = datetime.datetime.now().strftime("%Y-%m-%d")
        
        payload = {
            "function": "get_job_sensor_bucketed_data",
            "job_id": self.job_id_to_test,
            "sensor": "temperature",
            "date": today,
            "bucket_size": "1h"
        }
        
        data = self._post_api(payload)
        
        self.assertEqual(data.get('ok'), 1, f"Sensor data fetch failed: {data.get('error')}")
        self.assertIn('chart_data', data)
        self.assertIn('units', data)
        self.assertIn('time_zone', data)
        
        # Since we created the job with a virtual MAC in test_01, 
        # we expect chart_data to contain an entry for that device (even if data list is empty)
        # Note: The virtual MAC likely won't exist in the 'devices' table unless pre-seeded,
        # so the API might skip it or return empty. We check structure primarily.
        if data['chart_data']:
            print(f"-> Retrieved data for {len(data['chart_data'])} locations")
            first_loc = data['chart_data'][0]
            self.assertIn('name', first_loc)
            self.assertIn('data', first_loc)
        else:
            print("-> No device data found (Expected if test MAC is not in DB)")

if __name__ == '__main__':
    unittest.main()