# Libraries

import boto3
import os
from io import StringIO
from pathlib import Path
from dotenv import load_dotenv
import paramiko
import json
import time
import logging


# Variables

logger = logging.getLogger(__name__)


# Functions 

def get_public_ip(instance_id, ec2_client):
    logger.info('Getting public ip')
    reservations = ec2_client.describe_instances(InstanceIds=[instance_id]).get("Reservations", [])
    for reservation in reservations:
        for instance in reservation.get("Instances", []):
            return instance.get("PublicIpAddress")
    return None

def fetch_result(ec2_ip, pkey):
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    ssh.connect(hostname=ec2_ip, username='ec2-user', pkey=pkey)
    sftp = ssh.open_sftp()
    sftp.get('/home/ec2-user/result.json', 'result.json')
    sftp.close()
    ssh.close()

    with open("result.json", "r") as f:
        return json.load(f)

def run_remot_script(ec2_ip, env_path, script_path, pkey, streamlit_handler=None):
    

    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())

    ssh.connect(hostname=ec2_ip, username='ec2-user', pkey=pkey)

    logger.info('Putting parameters into sftp')
    sftp = ssh.open_sftp()
    sftp.put('tmp/params.json', '/home/ec2-user/params.json')
    sftp.close()


    logger.info(f'Executing script {script_path}')

    full_command = (
    'cd /home/ec2-user/set6pool_web && '
    'source /home/ec2-user/boost_env/bin/activate && '
    'PYTHONPATH=. nohup python3.11 future_to_options_batch/fto_batch.py > /home/ec2-user/fto_batch_log.txt 2>&1 &'
    )

    stdin, stdout, stderr = ssh.exec_command(full_command)
    last_log_content = ''
    while True:
        try:
            sftp = ssh.open_sftp()
            try:
                with sftp.open('/home/ec2-user/fto_batch_log.txt', 'r') as remote_log:
                    log_content = remote_log.read().decode().splitlines()
                    if log_content[-7:] != last_log_content:
                        streamlit_handler.clear()
                        logger.info("\n".join(log_content[-7:])) 
                    last_log_content = log_content[-7:]
            except FileNotFoundError:
                pass  # el log aún no existe
        
            # Revisar flag de finalización
            sftp.stat('/home/ec2-user/fto_batch_result.txt')
            logger.info("Remote Script Completed")
            ssh.exec_command('rm /home/ec2-user/fto_batch_result.txt')
            sftp.close()
            break

        except FileNotFoundError:
            logger.info("Waiting answer from EC2 process")
            time.sleep(2)
        except Exception as e:
            print(f"⚠️ Error reading EC2 log: {e}")
            time.sleep(1)

    print('Returning results')
    output = stdout.read().decode()

    ssh.close()
    return output