import os
import requests
from requests.auth import HTTPBasicAuth
from dataclasses import dataclass, field
from typing import List, Optional
from datetime import datetime
import time

@dataclass
class DAGStatus:
    name: str
    status: str
    group: Optional[str] = None
    schedule: Optional[str] = None
    lastRun: Optional[str] = None
    nextRun: Optional[str] = None
    pid: Optional[int] = None
    log: Optional[str] = None
    requestId: Optional[str] = None
    params: Optional[str] = None
    startedAt: Optional[str] = None
    finishedAt: Optional[str] = None
    suspended: Optional[bool] = None

    def get_last_run_epoch(self) -> Optional[int]:
        """Convert lastRun to epoch time."""
        return self._convert_to_epoch(self.lastRun)

    def get_next_run_epoch(self) -> Optional[int]:
        """Convert nextRun to epoch time."""
        return self._convert_to_epoch(self.nextRun)

    @staticmethod
    def _convert_to_epoch(timestamp: Optional[str]) -> Optional[int]:
        """Helper method to convert an ISO 8601 timestamp to epoch time."""
        if timestamp:
            dt = datetime.fromisoformat(timestamp.replace('Z', '+00:00'))
            return int(time.mktime(dt.timetuple()))
        return None

class DAGuClient:
    def __init__(self, base_url: str = "http://localhost:8888"):
        self.base_url = base_url
        self.auth = self._get_basic_auth()

    def _get_basic_auth(self) -> HTTPBasicAuth:
        """Retrieve the Basic Auth credentials from environment variables."""
        username = os.getenv('DAGU_BASICAUTH_USERNAME')
        password = os.getenv('DAGU_BASICAUTH_PASSWORD')

        if not username or not password:
            raise EnvironmentError("Please set the DAGU_BASICAUTH_USERNAME and DAGU_BASICAUTH_PASSWORD environment variables.")

        return HTTPBasicAuth(username, password)

    def list_dags(self) -> List[DAGStatus]:
        """Fetches the list of DAGs with their statuses from the DAGu REST API."""
        try:
            response = requests.get(f"{self.base_url}/api/v1/dags", auth=self.auth)
            response.raise_for_status()  # Raises an HTTPError for bad responses (4xx or 5xx)
            dags_data = response.json().get('DAGs', [])

            if isinstance(dags_data, list):
                return [self._parse_dag(dag) for dag in dags_data]
            else:
                print(f"Unexpected response format: {dags_data}")
                return []
        except requests.exceptions.RequestException as e:
            print(f"Error during request: {e}")
            return []

    def _parse_dag(self, dag_entry: dict) -> DAGStatus:
        """Helper function to parse a DAG's JSON data into a DAGStatus object."""
        try:
            dag_data = dag_entry.get("DAG", {})
            status_data = dag_entry.get("Status", {})

            return DAGStatus(
                name=dag_data.get("Name"),
                status=status_data.get("StatusText"),
                group=dag_data.get("Group"),
                schedule=(dag_data.get("Schedule", [{}])[0].get("Expression") 
                          if dag_data.get("Schedule") else None),
                lastRun=status_data.get("FinishedAt"),
                nextRun=None,  # Adjust as needed based on your API's response format
                pid=status_data.get("Pid"),
                log=status_data.get("Log"),
                requestId=status_data.get("RequestId"),
                params=status_data.get("Params"),
                startedAt=status_data.get("StartedAt"),
                finishedAt=status_data.get("FinishedAt"),
                suspended=dag_entry.get("Suspended")
            )
        except AttributeError as e:
            print(f"Error parsing DAG data: {dag_entry}, Error: {e}")
            return None

    def submit_dag_action(self, name: str, action: str, request_id: Optional[str] = None, params: Optional[str] = None) -> dict:
        """Submit an action to a specified DAG.
        
        Args:
            name (str): Name of the DAG.
            action (str): Action to be performed ('start', 'stop', or 'retry').
            request_id (Optional[str]): Required if action is 'retry'.
            params (Optional[str]): Parameters for the DAG execution.

        Returns:
            dict: Response from the API.
        """
        url = f"{self.base_url}/api/v1/dags/{name}"
        payload = {
            "action": action,
            **({"request-id": request_id} if request_id else {}),
            **({"params": params} if params else {}),
        }

        try:
            response = requests.post(url, json=payload, auth=self.auth)
            response.raise_for_status()  # Raises an HTTPError for bad responses (4xx or 5xx)
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"Error during request: {e}")
            print(f"Response content: {response.content}")
            return {}

# Example usage
if __name__ == "__main__":
    client = DAGuClient()
    
    # List DAGs
    try:
        dags = client.list_dags()
        for dag in dags:
            if dag:
                print(f"DAG Name: {dag.name}, Status: {dag.status}, Group: {dag.group}, "
                      f"Schedule: {dag.schedule}, Last Run: {dag.lastRun}, "
                      f"Next Run: {dag.nextRun}, PID: {dag.pid}, Log: {dag.log}, "
                      f"Request ID: {dag.requestId}, Params: {dag.params}, "
                      f"Started At: {dag.startedAt}, Finished At: {dag.finishedAt}, "
                      f"Suspended: {dag.suspended}")
                # Example of using helper methods to get epoch times
                if dag.get_last_run_epoch():
                    print(f"Last Run Epoch: {dag.get_last_run_epoch()}")
                if dag.get_next_run_epoch():
                    print(f"Next Run Epoch: {dag.get_next_run_epoch()}")
    except Exception as e:
        print(f"Error: {e}")

    # Submit an action to a DAG (example: start a DAG)
    try:
        dag_name = "test11"  # Replace with your actual DAG name
        action_response = client.submit_dag_action(name=dag_name, action="start")
        print(f"Action Response: {action_response}")
    except Exception as e:
        print(f"Error: {e}")