import datetime
import logging
import os
import psycopg2
import redis
from typing import Tuple

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# Environment variables (or default values)
POSTGRES_HOST = os.getenv("POSTGRES_HOST", "postgresql-genesis-u50607.vm.elestio.app")
POSTGRES_PORT = int(os.getenv("POSTGRES_PORT", "25432"))
POSTGRES_USER = os.getenv("POSTGRES_USER", "aiva")
POSTGRES_PASSWORD = os.getenv("POSTGRES_PASSWORD", "your_postgres_password")  # Replace with actual password
POSTGRES_DB = os.getenv("POSTGRES_DB", "aiva_db")

REDIS_HOST = os.getenv("REDIS_HOST", "redis-genesis-u50607.vm.elestio.app")
REDIS_PORT = int(os.getenv("REDIS_PORT", "26379"))

AIVA_COST_PER_DAY = 119.0
THEORETICAL_MAX_STORIES_PER_DAY = 380.0

class UtilizationTracker:
    """
    Tracks AIVA's utilization metrics and stores them in a database.
    """

    def __init__(self):
        """
        Initializes the UtilizationTracker, connecting to the database and Redis.
        """
        self.postgres_conn = self._connect_postgres()
        self.redis_client = self._connect_redis()

    def _connect_postgres(self):
        """
        Establishes a connection to the PostgreSQL database.

        Returns:
            psycopg2.extensions.connection: The database connection object.
        Raises:
            Exception: If the connection fails.
        """
        try:
            conn = psycopg2.connect(
                host=POSTGRES_HOST,
                port=POSTGRES_PORT,
                user=POSTGRES_USER,
                password=POSTGRES_PASSWORD,
                dbname=POSTGRES_DB
            )
            logging.info("Connected to PostgreSQL database.")
            return conn
        except Exception as e:
            logging.error(f"Failed to connect to PostgreSQL: {e}")
            raise

    def _connect_redis(self):
        """
        Establishes a connection to the Redis server.

        Returns:
            redis.Redis: The Redis client object.
        Raises:
            Exception: If the connection fails.
        """
        try:
            redis_client = redis.Redis(host=REDIS_HOST, port=REDIS_PORT, decode_responses=True)
            redis_client.ping()  # Check connection
            logging.info("Connected to Redis.")
            return redis_client
        except Exception as e:
            logging.error(f"Failed to connect to Redis: {e}")
            raise

    def log_utilization(self, tokens_generated: int) -> None:
        """
        Logs AIVA's utilization metrics, including requests processed and tokens generated.

        Args:
            tokens_generated (int): The number of tokens generated by AIVA for the request.
        """
        try:
            # Increment the number of requests processed today
            requests_today = self.redis_client.incr("aiva:requests_today")

            # Increment the total number of tokens generated today
            self.redis_client.incrby("aiva:tokens_generated_today", tokens_generated)

            logging.info(f"Logged utilization: Requests Today = {requests_today}, Tokens Generated = {tokens_generated}")

            self._store_utilization_data(requests_today)

        except Exception as e:
            logging.error(f"Failed to log utilization: {e}")

    def _store_utilization_data(self, requests_today: int) -> None:
        """
        Stores the utilization data in the PostgreSQL database.

        Args:
            requests_today (int): The number of requests processed today.
        """
        try:
            tokens_generated_today = int(self.redis_client.get("aiva:tokens_generated_today") or 0)
            uptime_percent = self._calculate_uptime_percent()
            cost_per_story = self._calculate_cost_per_story(requests_today)

            theoretical_max = THEORETICAL_MAX_STORIES_PER_DAY  #stories/day
            utilization_vs_capacity = (requests_today / theoretical_max) * 100

            current_time = datetime.datetime.now(datetime.timezone.utc)

            with self.postgres_conn.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO aiva_utilization (timestamp, requests_today, tokens_generated, uptime_percent, cost_per_story, utilization_vs_capacity)
                    VALUES (%s, %s, %s, %s, %s, %s);
                    """,
                    (current_time, requests_today, tokens_generated_today, uptime_percent, cost_per_story, utilization_vs_capacity),
                )
                self.postgres_conn.commit()

            logging.info("Utilization data stored in PostgreSQL.")

        except Exception as e:
            logging.error(f"Failed to store utilization data in PostgreSQL: {e}")
            self.postgres_conn.rollback()  # Rollback in case of error

    def _calculate_uptime_percent(self) -> float:
        """
        Calculates the uptime percentage of AIVA.  This is a placeholder and should be
        replaced with a real uptime monitoring solution.

        Returns:
            float: The uptime percentage.
        """
        # Placeholder: Assuming 99.9% uptime for now
        return 99.9

    def _calculate_cost_per_story(self, stories_processed: int) -> float:
        """
        Calculates the cost per story processed by AIVA.

        Args:
            stories_processed (int): The number of stories processed.

        Returns:
            float: The cost per story.
        """
        if stories_processed == 0:
            return 0.0  # Avoid division by zero
        return AIVA_COST_PER_DAY / stories_processed

    def reset_daily_metrics(self) -> None:
        """
        Resets the daily metrics (requests_today, tokens_generated_today) in Redis.
        This should be called at the start of each day.
        """
        try:
            self.redis_client.set("aiva:requests_today", 0)
            self.redis_client.set("aiva:tokens_generated_today", 0)
            logging.info("Daily metrics reset in Redis.")
        except Exception as e:
            logging.error(f"Failed to reset daily metrics in Redis: {e}")

    def close(self) -> None:
        """
        Closes the database connection.
        """
        if self.postgres_conn:
            self.postgres_conn.close()
            logging.info("PostgreSQL connection closed.")

        if self.redis_client:
            self.redis_client.close()
            logging.info("Redis connection closed.")

# Example usage (for testing - remove in production):
if __name__ == "__main__":
    tracker = UtilizationTracker()
    try:
        # Simulate AIVA processing requests
        tracker.log_utilization(tokens_generated=100)
        tracker.log_utilization(tokens_generated=150)
        tracker.log_utilization(tokens_generated=200)

        # Reset daily metrics (for testing purposes)
        # tracker.reset_daily_metrics()

        # Simulate more requests after reset
        # tracker.log_utilization(tokens_generated=120)

    finally:
        tracker.close() # Close connections on exit