import argparse
import datetime
from pathlib import Path
from typing import Optional
from typing import Tuple

import pydantic
import yaml

import autoeft
import autoeft.base.basis as eft_basis
import autoeft.base.model as ir_model
from autoeft import exceptions
from autoeft import utils


class MetaData(pydantic.BaseModel):
    timestamp: datetime.datetime
    eft_name: str
    dimension: int
    args: dict

    @pydantic.validator("args")
    def validate_args(cls, v):
        for k in "model", "output", "generators":
            if k in v:
                v[k] = Path(v[k])
        return v


class BasisFile:
    basis_path: Path
    model_path: Path
    extra_path: Path

    def __init__(self, basis_path: Path):
        self.basis_path = basis_path
        self.model_path = basis_path / "model.json"
        self.extra_path = basis_path / ".autoeft"

    def __bool__(self):
        return (
            self.basis_path.is_dir()
            and self.model_path.is_file()
            and self.extra_path.is_file()
        )

    def get_model(self) -> ir_model.Model:
        try:
            return ir_model.Model.parse_json(self.model_path.read_text())
        except FileNotFoundError as e:
            errmsg = f"No '{self.model_path.name}' in basis directory."
            raise exceptions.BasisNotFoundError(errmsg) from e

    def set_model(self, model: ir_model.Model):
        self.basis_path.mkdir(parents=True, exist_ok=True)
        data = model.json(indent=2)
        self.model_path.write_text(data)

    def get_extras(self):
        try:
            return [
                MetaData(**data) for data in yaml.safe_load(self.extra_path.read_text())
            ]
        except FileNotFoundError as e:
            errmsg = f"No '{self.extra_path.name}' in basis directory."
            raise exceptions.BasisNotFoundError(errmsg) from e
        except pydantic.ValidationError as e:
            errmsg = f"{e}\nCannot parse '{self.extra_path.name}'."
            raise exceptions.BasisNotFoundError(errmsg) from e

    def get_extra(self):
        return self.get_extras()[-1]

    def set_extra(
        self,
        eft_name: str,
        dimension: int,
        args: argparse.Namespace,
    ):
        exclude = {"main_func", "quiet", "verbose", "dry_run", "name"}
        args_dict = {k: v for k, v in vars(args).items() if k not in exclude}
        meta = MetaData.construct(
            timestamp=datetime.datetime.now().isoformat(timespec="seconds"),
            eft_name=eft_name,
            dimension=dimension,
            args=args_dict,
        )
        self.basis_path.mkdir(parents=True, exist_ok=True)
        with self.extra_path.open("a") as extra_file:
            extra_file.write(
                yaml.dump([meta.dict()], sort_keys=False, Dumper=utils.YamlDumper),
            )

    def add_operator_info(self, operator_info: eft_basis.OperatorInfoPermutation):
        self.basis_path.mkdir(parents=True, exist_ok=True)
        operator_path = self.basis_path / operator_info.path()
        operator_path.parent.mkdir(parents=True, exist_ok=True)
        data = (
            f"# '{operator_path.name}' generated by "
            f"{autoeft.program} {autoeft.version}\n"
        )
        data += operator_info.yaml()
        tmp_operator_path = operator_path.with_suffix(".tmp")
        tmp_operator_path.write_text(data)
        tmp_operator_path.rename(operator_path)

    def get_operator_info(self, rel_path: Path) -> eft_basis.OperatorInfoPermutation:
        model = self.get_model()
        operator_path = self.basis_path / rel_path
        try:
            return eft_basis.OperatorInfoPermutation(
                model=model,
                **yaml.safe_load(operator_path.read_text()),
            )
        except FileNotFoundError as e:
            errmsg = f"No '{rel_path}' in basis directory."
            raise exceptions.BasisNotFoundError(errmsg) from e

    def get_operator_raw_numbers(self, rel_path: Path) -> Optional[Tuple[int, int]]:
        operator_path = self.basis_path / rel_path
        if operator_path.is_file():
            data = yaml.safe_load(operator_path.read_text())
            return data["n_terms"], data["n_operators"]
        return None

    def get_basis(self) -> eft_basis.Basis:
        def _sort(op_info: eft_basis.OperatorInfo):
            family = tuple(-nh for nh in op_info.op_type.family)
            op_type = tuple((f.name, m) for f, m in op_info.op_type)
            return op_info.N, (-op_info.nl, -op_info.nr), family, op_type

        model = self.get_model()
        operators = {}
        for operator_path in self.basis_path.rglob("*.yml"):
            op_info = eft_basis.OperatorInfoPermutation(
                model=model,
                **yaml.safe_load(operator_path.read_text()),
            )
            operators[repr(op_info)] = op_info
        operators_sorted = sorted(operators.items(), key=lambda item: _sort(item[1]))
        operator_basis = eft_basis.Basis(model)
        operator_basis.update(operators_sorted)
        return operator_basis
