import socket
import os
import struct
from typing import Generator, Dict, Optional, AsyncGenerator
from pathlib import Path

from alayaflow.execution.executors.base_executor import BaseExecutor
from alayaflow.execution.env_manager import EnvManager
from alayaflow.workflow.workflow_manager import WorkflowManager
from alayaflow.common.config import settings


class UvExecutor(BaseExecutor):
    def __init__(
        self,
        workflow_manager: WorkflowManager,
        env_manager: Optional[EnvManager] = None
    ):
        self.workflow_manager = workflow_manager
        self.env_mgr = env_manager if env_manager is not None else EnvManager()
    
    def _get_free_port(self):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.bind(('', 0))
            return s.getsockname()[1]
    
    def execute_stream(
        self,
        workflow_id: str,
        version: str,
        inputs: dict,
        context: dict,
        metadata: dict = None
    ) -> Generator[Dict, None, None]:
        raise NotImplementedError("uv executor not supported yet")

        # # 1. 确定工作流目录和版本
        # workflow_info = self.workflow_manager.get_workflow_info(workflow_id, version)
        # if not workflow_info:
        #     yield {"error": f"工作流 {workflow_id}_{version} 不存在"}
        #     return
        
        # wf_dir = str(workflow_info.wf_dir)
        
        # # 2. 读取工作流依赖
        # requirements = self.workflow_manager.get_workflow_requirements(workflow_id, version)
        
        # # 3. 使用 uv 准备环境 (自动管理 python 版本)
        # python_exe = self.env_mgr.ensure_env(workflow_id, version, requirements, use_venv=True)

        # 4. 准备 Socket Runner
        # if settings.is_dev():
        #     # dev 模式：从项目根目录找
        #     runner_script_path = os.path.join(settings.alayaflow_root, 'src/alayaflow/execution/workflow_runner.py')
        # else:
        #     # dev-uneditable/prod 模式：从安装包中找
        #     runner_script_path = str(Path(__file__).parent.parent / 'workflow_runner.py')
        
        # # 5. 启动 Server Socket
        # server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        # port = self._get_free_port()
        # server_socket.bind(('127.0.0.1', port))
        # server_socket.listen(1)
        
        # # 6. 启动子进程
        # import subprocess
        # import json
        # process = subprocess.Popen(
        #     [python_exe, runner_script_path, str(port), wf_dir],
        #     stdin=subprocess.PIPE,
        #     stdout=subprocess.PIPE,
        #     stderr=subprocess.PIPE,
        #     text=True
        # )
        
        # process.stdin.write(json.dumps(init_args))
        # process.stdin.write('\n')
        # process.stdin.write(json.dumps(config))
        # process.stdin.write('\n')
        # process.stdin.write(json.dumps(input_data))
        # process.stdin.write('\n')
        # process.stdin.close()

        # # 7. 等待连接
        # server_socket.settimeout(10.0)  # 首次 uv 下载 python 可能稍慢，增加超时
        # conn = None
        # try:
        #     conn, addr = server_socket.accept()
            
        #     # 9. 读取数据流
        #     while True:
        #         header = conn.recv(4)
        #         if not header:
        #             break
        #         msg_len = struct.unpack('!I', header)[0]
                
        #         chunks = []
        #         bytes_recd = 0
        #         while bytes_recd < msg_len:
        #             chunk = conn.recv(min(msg_len - bytes_recd, 4096))
        #             if not chunk:
        #                 raise RuntimeError("Socket connection broken")
        #             chunks.append(chunk)
        #             bytes_recd += len(chunk)
                
        #         full_msg = b''.join(chunks)
        #         full_msg_json = json.loads(full_msg.decode('utf-8'))
        #         yield full_msg_json

        # except socket.timeout:
        #     yield {"error": "Runner connection timeout (Maybe UV is downloading python?)"}
        # except Exception as e:
        #     yield {"error": f"Socket Error: {e}"}
        # finally:
        #     if conn:
        #         conn.close()
        #     server_socket.close()
            
        #     logs = process.stdout.read()
        #     if logs:
        #         print(f"--- [Clean Runner Logs] ---\n{logs}---------------------------")
            
        #     stderr = process.stderr.read()
        #     if stderr:
        #         print(f"[Stderr] {stderr}")


    def execute_stream_async(
        self,
        workflow_id: str,
        version: str,
        inputs: dict,
        context: dict,
        metadata: dict = None
    ) -> AsyncGenerator[Dict, None]:
        raise NotImplementedError("uv executor not supported yet")