from pathlib import Path
from typing import Literal,Any,Optional
from tqdm import tqdm
import yaml
import logging
from logging.handlers import RotatingFileHandler

from .node import Node
from .dag import DAG
from .utils import executor

class Flow:
    @classmethod
    def from_config(cls, config, node_cls_map: dict[str, type]):
        flow = cls(**config["flow"])

        for node_cfg in config.get("nodes", []):
            node_cfg = dict(node_cfg)
            node_cls_name = node_cfg.pop("cls")
            node_cfg["cls"] = node_cls_map[node_cls_name]
            flow.add_node(**node_cfg)

        for edge_cfg in config.get("edges", []):
            edge_cfg = dict(edge_cfg)
            edge_cfg["ports"] = [tuple(p) for p in edge_cfg["ports"]]
            flow.add_edge(**edge_cfg)

        return flow

    @classmethod
    def from_yaml(cls, config_path: str, node_cls_map: dict[str,type]):
        
        with open(config_path, "r") as f:
            config = yaml.safe_load(f)
        return cls.from_config(config, node_cls_map)

    def __init__(self, work_path: str):
        self.work_path = Path(work_path)
        self.work_path.mkdir(parents=True, exist_ok=True)

        self.log_path = self.work_path / "logs"
        self.log_path.mkdir(parents=True, exist_ok=True)
        self.stores_path = self.work_path / "stores"
        self.stores_path.mkdir(parents=True, exist_ok=True)

        self.nodes: dict[str,Node] = {}
        self.edges: dict[tuple[str,str],tuple[list,list]] = {}
        self.dag = DAG()

        self.init_logger(self.log_path)
        self.build_context()

    
    def init_logger(self,log_path: Path):
        self.logger = logging.getLogger(f"{__name__}.{id(self)}")
        self.logger.setLevel(logging.DEBUG)
        self.logger.propagate = False

        formatter = logging.Formatter(
            '%(asctime)s [%(levelname)s] [%(filename)s:%(lineno)d] %(message)s',
            datefmt='%Y-%m-%d %H:%M:%S'
        )

        console_handler = logging.StreamHandler()
        console_handler.setFormatter(formatter)
        console_handler.setLevel(logging.INFO)
        self.logger.addHandler(console_handler)

        log_path.mkdir(parents=True, exist_ok=True)
        log_file_path = log_path / "app.log"
        file_handler = RotatingFileHandler(log_file_path, maxBytes=1*1024*1024, backupCount=3, encoding='utf-8')
        file_handler.setFormatter(formatter)
        file_handler.setLevel(logging.DEBUG)
        self.logger.addHandler(file_handler)

    def build_context(self):
        self.context = {
            "flow":{
                "dag":self.dag,
                "logger":self.logger,
                "work_path":self.work_path,
                "stores_path":self.stores_path,
                "log_path":self.log_path
            }
        }

    def add_node(self,                
                name: str,
                cls: type,
                init_input: Optional[dict[str, Any]] = None,
                run_input: Optional[dict[str, Any]] = None,
                run_output: Optional[dict[str, Any]] = None,
                database_type: str = "cache",
                ):
        
        node = Node(name=name,
                    cls=cls,
                    init_input=init_input,
                    run_input=run_input,
                    run_output=run_output,
                    stores_path=self.stores_path,
                    database_type=database_type,
                    context=self.context,
                    )
        self.nodes[name] = node
        self.dag.add_node(name)
    
    def add_edge(self, src_name: str, dst_name: str, ports: list[tuple[str,str]]):        
        self.edges[(src_name,dst_name)] = ports
        self.dag.add_edge(src_name,dst_name)

    def build_run_context(self, name: str, task: int | str):
        run_context = {"task":task,"node":name}
        return run_context

    def find_run_input(self, name: str, task: int | str):
        run_input = {}
        predecessors = self.dag.find_predecessors(name)
        for predecessor in predecessors:
            predecessor_output = self.find_run_output(predecessor,task)
            if predecessor_output is None:
                continue
            edge = self.edges[(predecessor,name)]
            run_input_ = {port[1]: predecessor_output[port[0]] for port in edge}
            run_input.update(run_input_)
        return run_input

    def find_run_output(self, name: str, task: int | str):
        node = self.nodes[name]
        return node.find_run_output(task)

    def find_run_completed_tasks(self, name: str) -> list[int | str]:
        node = self.nodes[name]
        return node.find_run_completed_tasks()

    def run_task_chain(self, task: int | str, chain: list[str]):
        for name in chain:
            node = self.nodes[name]
            if node.requires_run(task):
                run_context = self.build_run_context(name,task)
                run_input = self.find_run_input(name,task)
                node.run(run_context,run_input)
            self.pbar.set_postfix({"node":name,"task":task})
            self.pbar.update(1)

    def prepare_resume_nodes(self,resume_nodes: list[str], graph: dict[str,list[str]]):
        subgraph = DAG.get_successor_subgraph(graph,resume_nodes)
        resume_nodes = DAG.get_graph_nodes(subgraph)
        for name in resume_nodes:
            node=self.nodes[name]
            node.bump_version()

    def run(self,
            tasks: str | int | list[int | str],
            nodes:  list[str] = None,
            resume_nodes: list[str] = None,
            mode: Literal["task","stage"] = "stage",
            num_workers: int = 1
            ):
        
        if isinstance(tasks,str):
            tasks = self.nodes[tasks].generate_task()

        if isinstance(tasks,int):
            tasks = list(range(tasks))
        
        if nodes is None:
            graph=self.dag.graph
            chain = DAG.topological_sort(graph)
        else:
            graph = DAG.get_predecessor_subgraph(graph, nodes)
            chain = DAG.topological_sort(graph)

        if resume_nodes is not None:
            resume_nodes = [node for node in resume_nodes if node in chain]
            self.prepare_resume_nodes(resume_nodes,graph)

        num_task = len(tasks)
        num_node = len(chain)
        total_task = num_task * num_node
        self.logger.info(f"chain:{chain}")
        self.logger.info(f"mode:{mode} num_node:{num_node} num_task:{num_task} num_workers:{num_workers}")
        self.logger.info(f"Start running")

        self.pbar = tqdm(total = total_task,desc="Running")
        if mode == "task":
            for name in chain:
                self.nodes[name].load()
            kwargs_list = [{"task":task,"chain":chain} for task in tasks]
            executor(self.run_task_chain,kwargs_list,num_workers)
            for name in chain:
                self.nodes[name].unload()
        elif mode == "stage":
            for name in chain:
                self.nodes[name].load()
                kwargs_list = [{"task":task,"chain":[name]} for task in tasks]
                executor(self.run_task_chain,kwargs_list,num_workers)
                self.nodes[name].unload()
        else:
            raise ValueError(f"mode {mode} not supported")
        self.pbar.set_description("done")
        self.pbar.close()
        self.logger.info(f"Finish running")
