import os
import sys
import time
import socket
import subprocess
import threading
import logging
import uuid
import requests
from datetime import datetime, timezone
from dotenv import load_dotenv

# Load environment variables from .env
load_dotenv()

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(threadName)s - %(message)s",
    handlers=[
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger("scheduler_agent")

# Configuration
API_URL = os.getenv("API_URL", "http://127.0.0.1:8000/v1")
ENVIRONMENT = os.getenv("ENV", "local")
SERVER_GROUP = os.getenv("SERVER_GROUP", "content-studio-prod")
SERVER_KEY = os.getenv("SERVER_KEY", socket.gethostname())
PRODUCT_KEY = os.getenv("PRODUCT_KEY", "content-studio")
POLL_INTERVAL = 30

# Thread-safe tracker for running jobs
running_jobs = set()
running_jobs_lock = threading.Lock()

def execute_job(job):
    job_key = job["jobKey"]
    execution_id = str(uuid.uuid4())

    
    with running_jobs_lock:
        if job_key in running_jobs:
            logger.warning(f"Job {job_key} is already running. Skipping execution trigger.")
            return
        running_jobs.add(job_key)
        
    logger.info(f"Starting job {job_key} (Execution ID: {execution_id})")
    started_at = datetime.now(timezone.utc)
    
    interpreter = job["interpreter"]
    script_path = job["scriptPath"]
    arguments = job.get("arguments", [])
    cwd = job.get("workingDirectory", ".")
    timeout = job.get("timeoutSeconds", 3600)
    local_log_path = job.get("localLogPath")
    
    cmd = [interpreter, script_path] + arguments
    
    # Try creating directory for local logs if specified
    log_file = None
    if local_log_path:
        try:
            os.makedirs(os.path.dirname(local_log_path), exist_ok=True)
            log_file = open(local_log_path, "a", encoding="utf-8")
            log_file.write(f"\n--- Centralized Scheduler Run: {job_key} at {started_at.isoformat()} (Exec ID: {execution_id}) ---\n")
            log_file.flush()
        except Exception as e:
            logger.error(f"Failed to open local log file {local_log_path}: {e}")

    stdout_data = ""
    exit_code = -1
    status = "failure"
    
    try:
        # Run process and capture all output
        process = subprocess.Popen(
            cmd,
            cwd=cwd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True
        )
        
        try:
            stdout_data, _ = process.communicate(timeout=timeout)
            exit_code = process.returncode
            status = "success" if exit_code == 0 else "failure"
        except subprocess.TimeoutExpired:
            process.kill()
            stdout_data, _ = process.communicate()  # retrieve remaining output
            exit_code = -2
            status = "failure"
            stdout_data += f"\n[ERROR] Execution timed out after {timeout} seconds and process was killed."
            
    except Exception as e:
        logger.exception(f"Failed to spawn subprocess for {job_key}")
        stdout_data = f"Exception spawning process: {str(e)}"
        status = "failure"
        exit_code = -3
        
    ended_at = datetime.now(timezone.utc)
    duration = (ended_at - started_at).total_seconds()
    
    # Write captured output to local log file
    if log_file:
        try:
            log_file.write(stdout_data)
            log_file.write(f"\n--- Centralized Scheduler End: {job_key} at {ended_at.isoformat()} (Exit Code: {exit_code}) ---\n")
            log_file.close()
        except Exception as e:
            logger.error(f"Failed to write end logs to {local_log_path}: {e}")

    logger.info(f"Finished job {job_key} in {duration:.2f} seconds. Status: {status}, Exit Code: {exit_code}")

    # Upload execution result to backend
    payload = {
        "executionId": execution_id,
        "jobKey": job_key,
        "environment": ENVIRONMENT,
        "serverKey": SERVER_KEY,
        "status": status,
        "logs": stdout_data,
        "startedAt": started_at.isoformat(),
        "endedAt": ended_at.isoformat(),
        "durationSeconds": duration,
        "exitCode": exit_code
    }
    
    upload_url = f"{API_URL}/scheduler/execution-result"
    if PRODUCT_KEY:
        upload_url += f"?productKey={PRODUCT_KEY}"
        
    try:
        resp = requests.post(upload_url, json=payload, timeout=10)
        if resp.status_code == 200:
            logger.info(f"Successfully uploaded execution result for {job_key}")
        else:
            logger.error(f"Failed to upload execution result for {job_key}. Status: {resp.status_code} - {resp.text}")
    except Exception as e:
        logger.error(f"Failed to connect to backend to upload execution result: {e}")

    # Clean up track list
    with running_jobs_lock:
        running_jobs.remove(job_key)

def poll_due_jobs():
    url = f"{API_URL}/scheduler/due-jobs"
    params = {
        "serverGroup": SERVER_GROUP,
        "environment": ENVIRONMENT
    }
    if PRODUCT_KEY:
        params["productKey"] = PRODUCT_KEY
    try:
        resp = requests.get(url, params=params, timeout=10)
        if resp.status_code == 200:
            due_jobs = resp.json().get("due_jobs", [])
            if due_jobs:
                logger.info(f"Retrieved {len(due_jobs)} due job(s) from scheduler")
                for job in due_jobs:
                    # Spawn threading job execution
                    t = threading.Thread(
                        target=execute_job,
                        args=(job,),
                        name=f"Worker-{job['jobKey']}"
                    )
                    t.daemon = True
                    t.start()
        else:
            logger.error(f"Failed to poll due jobs: {resp.status_code} - {resp.text}")
    except Exception as e:
        logger.error(f"Error polling due jobs: {e}")

def main():
    logger.info(f"Starting Scheduler Agent. Server Key: {SERVER_KEY}, Group: {SERVER_GROUP}, Environment: {ENVIRONMENT}")
    logger.info(f"API Target: {API_URL}")
    
    while True:
        try:
            poll_due_jobs()
        except Exception as e:
            logger.error(f"Unhandled error in main agent loop: {e}")
        time.sleep(POLL_INTERVAL)

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        logger.info("Scheduler Agent stopped by user.")
        sys.exit(0)