from dataclasses import dataclass, is_dataclass, fields
from string import Template
from typing import Any, Generic, Mapping, TypeVar, cast
from functools import singledispatch


from .lazy import Lazy


T = TypeVar("T")


@singledispatch
def substitute(template, env: Mapping[str, str]):
    dtype = type(template)
    if is_dataclass(dtype):
        args = {
            f.name: substitute(getattr(template, f.name), env)
            for f in fields(dtype)
            if f.name[0] != "_"
        }
        return dtype(**args)

    return template
    # raise TypeError(f"Can't perform string substitution on object of type: {dtype}")


@substitute.register
def _(template: str, env: Mapping[str, str]) -> str:
    return Template(template).safe_substitute(env)


@substitute.register
def _(template: list, env: Mapping[str, str]) -> list:
    return [substitute(x, env) for x in template]

@substitute.register
def _(_template: None, _) -> None:
    return None


@singledispatch
def gather_args(template: Any) -> set[str]:
    dtype = type(template)
    if is_dataclass(dtype):
        args = (
            gather_args(getattr(template, f.name))
            for f in fields(dtype)
            if f.name[0] != "_"
        )
        return set().union(*args)

    return set()
    # raise TypeError(f"Can't perform string substitution on object of type: {dtype}")


@gather_args.register
def _(template: str) -> set[str]:
    return set(Template(template).get_identifiers())


@gather_args.register
def _(template: list) -> set[str]:
    return set().union(*map(gather_args, template))


@gather_args.register
def _(_template: None) -> set[str]:
    return set()

from dataclasses import dataclass
from typing import Iterable, Optional
import pytest
from brei.task import TemplateVariable, Variable
from brei.template_strings import gather_args, substitute
from brei.lazy import LazyDB


class Environment(LazyDB[Variable, TemplateVariable]):
    def __setitem__(self, k: str, v: str):
        self.add(TemplateVariable([Variable(k)], [], v))

    def __getitem__(self, k: str) -> str:
        return self.index[Variable(k)].result

    def __contains__(self, k: str) -> bool:
        return Variable(k) in self.index

    def items(self) -> Iterable[str]:
        return (k.name for k in self.index if isinstance(k, Variable))

    @property
    def environment(self):
        return self


@pytest.mark.asyncio
async def test_template_string():
    env = Environment()
    env["x"] = "Hello, ${y}!"
    env["y"] = "World"
    env["z"] = "print('${x}')"
    await env.run(Variable("z"), db=env)
    assert env["x"] == "Hello, World!"
    assert env["z"] == "print('Hello, World!')"


@dataclass
class MyData:
    some_list: list[str]
    some_prop: str
    some_none: Optional[str] = None


def test_template_dtype():
    data = MyData(
        some_list = ["${x} bar", "bar ${x} bar"],
        some_prop = "bar ${x}"
    )

    assert gather_args(data) == set("x")
    subst =  substitute(data, {"x": "foo"})
    assert subst.some_list == ["foo bar", "bar foo bar"]
    assert subst.some_prop == "bar foo"
    assert subst.some_none is None

from dataclasses import dataclass


@dataclass
class Runner:
    command: str
    args: list[str]


DEFAULT_RUNNERS: dict[str, Runner] = {
    "python": Runner("python", ["${script}"]),
    "bash": Runner("bash", ["${script}"]),
}

from __future__ import annotations
import asyncio
from contextlib import contextmanager, nullcontext
from copy import copy
from dataclasses import dataclass, field
from pathlib import Path
import re
import string
from tempfile import NamedTemporaryFile
from typing import Any, Optional, TextIO
from asyncio import create_subprocess_exec
from textwrap import indent
import shlex

from .result import TaskFailure
from .lazy import MissingDependency, Lazy, LazyDB, Phony
from .utility import stat
from .logging import logger
from .template_strings import gather_args, substitute
from .runner import Runner, DEFAULT_RUNNERS


log = logger()


@dataclass
class Variable:
    name: str

    def __hash__(self):
        return hash(f"var({self.name})")

    def __str__(self):
        return f"var({self.name})"


def str_to_target(s: str) -> Path | Phony | Variable:
    if s[0] == "#":
        return Phony(s[1:])
    elif m := re.match(r"var\(([^\s\(\)]+)\)", s):
        return Variable(m.group(1))
    else:
        return Path(s)


def is_oneliner(s: str) -> bool:
    return len(s.splitlines()) == 1


@dataclass
class Task(Lazy[Path | Phony | Variable, str | None]):
    name: Optional[str] = None
    runner: Optional[str] = None
    path: Optional[Path] = None
    script: Optional[str] = None
    stdin: Optional[Path | Variable] = None
    stdout: Optional[Path | Variable] = None
    description: Optional[str] = None
    force: bool = False

    @property
    def target_paths(self):
        return (p for p in self.creates if isinstance(p, Path))

    @property
    def dependency_paths(self):
        return (p for p in self.requires if isinstance(p, Path))

    def __str__(self):
        tgts = ", ".join(str(t) for t in self.creates)
        deps = ", ".join(str(t) for t in self.requires)
        if self.script is not None:
            src = indent(self.script, prefix=" ▎ ", predicate=lambda _: True)
        elif self.path is not None:
            src = str(self.path)
        else:
            src = " - "
        name = f"{self.name}: " if self.name else ""
        return name + f"[{tgts}] <- [{deps}]\n" + src

    def __post_init__(self):
        if self.name and Phony(self.name) not in self.creates:
            self.creates.append(Phony(self.name))
        if self.stdin and self.stdin not in self.requires:
            self.requires.append(self.stdin)
        if self.path and self.path not in self.requires:
            self.requires.append(self.path)
        if self.stdout and self.stdout not in self.creates:
            self.creates.append(self.stdout)

    def always_run(self) -> bool:
        return self.force or len(list(self.target_paths)) == 0

    def needs_run(self) -> bool:
        if any(not p.exists() for p in self.target_paths):
            return True
        target_stats = [stat(p) for p in self.target_paths]
        dep_stats = [stat(p) for p in self.dependency_paths]
        if any(t < d for t in target_stats for d in dep_stats):
            return True
        return False

    @contextmanager
    def get_script_path(self):
        if self.path is not None:
            tmpfile = None
            path = self.path
        elif self.script is not None:
            tmpfile = NamedTemporaryFile("w")
            tmpfile.write(self.script)
            tmpfile.flush()
            path = Path(tmpfile.name)
        else:
            raise ValueError("A `Rule` can have either `path` or `script` defined.")

        yield path

        if tmpfile is not None:
            tmpfile.close()

    @contextmanager
    def get_stdout(self):
        match self.stdout:
            case Variable(x):
                yield asyncio.subprocess.PIPE
            case x if isinstance(x, Path):
                stdout = open(x, "w")
                yield stdout
                stdout.close()
            case _:
                yield None

    async def run(self, *, db: TaskDB):
        if not self.always_run() and not self.needs_run() and not db.force_run:
            tgts = " ".join(f"`{t}`" for t in self.target_paths)
            log.info(f"Targets {tgts} already up-to-date.")
            return

        log.debug(f"{self}")
        if (self.path is None and self.script is None):
            return

        targets = " ".join(f"`{t}`" for t in self.creates)
        short_note = self.description or (f"#{self.name}" if self.name else None) \
            or f"creating {targets}"
        log.info(f"[green]{short_note}[/]", extra={'markup': True})

        stdin: TextIO | int | None = None
        match self.stdin:
            case Variable(x):
                stdin = asyncio.subprocess.PIPE
                input_data = db.environment[x].encode()
            case x if isinstance(x, Path):
                stdin = open(x, "r")
                input_data = None
            case _:
                stdin = None
                input_data = None

        if self.runner is None and self.script is not None:
            if not is_oneliner(self.script):
                assert self.stdin is None
            with self.get_stdout() as stdout:
                stdout_data = b""
                for line in self.script.splitlines():
                    async with db.throttle or nullcontext():
                        proc = await create_subprocess_exec(
                            *shlex.split(line),
                            stdin=stdin,
                            stdout=stdout,
                            stderr=asyncio.subprocess.PIPE,
                        )
                        stdout_data_part, stderr_data = await proc.communicate(input_data)
                        log.debug(f"return-code {proc.returncode}")
                    if stdout_data_part:
                        stdout_data += stdout_data_part
                    if stderr_data:
                        log.info(f"[gold1]{short_note}[/] %s", stderr_data.decode().rstrip(), extra={"markup": True})

        elif self.runner is not None:
            with self.get_script_path() as path, self.get_stdout() as stdout:
                runner = db.runners[self.runner]
                args = [string.Template(arg).substitute(script=path) for arg in runner.args]
                async with db.throttle or nullcontext():
                    proc = await create_subprocess_exec(
                        runner.command,
                        *args,
                        stdin=stdin,
                        stdout=stdout,
                        stderr=asyncio.subprocess.PIPE,
                    )
                    stdout_data, stderr_data = await proc.communicate(input_data)
                    log.debug(f"return-code {proc.returncode}")

            if stderr_data:
                log.info(f"[gold1]{short_note}[/] %s", stderr_data.decode().rstrip(), extra={"markup": True})

        else:
            return

        if self.needs_run():
            raise TaskFailure("Task didn't achieve goals.")

        return stdout_data.decode().strip() if stdout_data else None


@dataclass
class TaskProxy:
    """Input to create an actual Task.

    An actual `Task` has no template variables remaining and untyped strings
    are replaced with `Path` `Variable` or `Phony` objects.
    """
    creates: list[str] = field(default_factory=list)
    requires: list[str] = field(default_factory=list)
    name: Optional[str] = None
    runner: Optional[str] = None
    path: Optional[str] = None
    script: Optional[str] = None
    stdin: Optional[str] = None
    stdout: Optional[str] = None
    description: Optional[str] = None
    force: bool = False

    @property
    def all_targets(self):
        return (
            self.creates
            + ([self.stdout] if self.stdout else [])
            + ([f"#{self.name}"] if self.name else [])
        )

    @property
    def all_dependencies(self):
        return (
            self.requires
            + ([self.stdin] if self.stdin else [])
            + ([self.path] if self.path else [])
        )


@dataclass
class TemplateVariable(Lazy[Variable, str]):
    template: str

    def __post_init__(self):
        self.requires += [Variable(arg) for arg in gather_args(self.template)]

    async def run(self, *, db) -> str:
        return substitute(self.template, db.environment)


@dataclass
class TemplateTask(Lazy[Path | Phony | Variable, Task]):
    template: TaskProxy

    def __post_init__(self):
        assert not gather_args(self.template.creates)
        self.creates += [str_to_target(t) for t in self.template.all_targets]
        self.requires += [Variable(arg) for arg in gather_args(self.template)]

    async def run(self, *, db):
        proxy = substitute(self.template, db.environment)
        tgts = [str_to_target(t) for t in proxy.all_targets]
        deps = [str_to_target(t) for t in proxy.all_dependencies]
        path = Path(proxy.path) if proxy.path else None
        stdin = str_to_target(proxy.stdin) if proxy.stdin else None
        stdout = str_to_target(proxy.stdout) if proxy.stdout else None
        assert not isinstance(stdin, Phony) and not isinstance(stdout, Phony)
        task = Task(
            tgts,
            deps,
            proxy.name,
            proxy.runner,
            path,
            proxy.script,
            stdin,
            stdout,
            proxy.description,
            proxy.force
        )
        return task


@dataclass
class TaskDB(LazyDB[Path | Variable | Phony, Task | TemplateTask | TemplateVariable]):
    runners: dict[str, Runner] = field(default_factory=lambda: copy(DEFAULT_RUNNERS))
    throttle: Optional[asyncio.Semaphore] = None
    force_run: bool = False

    def on_missing(self, t: Path | Phony | Variable):
        if isinstance(t, Path) and t.exists():
            return Task([t], [])
        raise MissingDependency()

    def is_resolvable(self, s: Any) -> bool:
        return all(v in self.index for v in map(Variable, gather_args(s)))

    async def resolve_object(self, s: Any) -> Any:
        vars = gather_args(s)
        await asyncio.gather(*(self.run(Variable(v), db=self) for v in vars))
        result = substitute(s, self.environment)
        log.debug(f"substituting {s} => {result}")
        return result

    @property
    def environment(self):
        return Environment(self)


class Environment:
    def __init__(self, db):
        self.db = db

    def __contains__(self, k: str):
        return Variable(k) in self.db.index

    def items(self):
        return (k.name for k in self.db.index if isinstance(k, Variable))

    def __getitem__(self, k: str):
        return self.db.index[Variable(k)].result


class Template(TaskProxy):
    """A `Template` can receive the same arguments as a TaskProxy. The difference is 
    that any template variables in the Template can be substituted with a `TemplateCall`."""
    def call(self, args: dict[str, Any]) -> TaskProxy:
        return substitute(self, args)

from __future__ import annotations
import asyncio
from copy import copy
from itertools import chain, product, repeat
from dataclasses import dataclass, field
from enum import Enum
import json
from pathlib import Path

import tomllib


from .template_strings import gather_args
from .logging import logger
from .errors import HelpfulUserError, UserError

from .utility import construct, read_from_file
from .task import (
    Variable,
    TaskDB,
    Template,
    Runner,
    TaskProxy,
    TemplateTask,
    TemplateVariable,
)


log = logger()


@dataclass
class MissingInclude(UserError):
    path: Path

    def __str__(self):
        return f"Include `{self.path}` not found."


@dataclass
class MissingTemplate(UserError):
    name: str

    def __str__(self):
        return f"Template `{self.name}` not found."


class Join(Enum):
    INNER = 1
    OUTER = 2


@dataclass
class TemplateCall:
    """Calls a template with a set of arguments.

    Members:

      - template: name of the template.
      - args: arguments to the call.
      - collect: name of the phony target by which to collect all generated targets.
      - join: `inner` or `outer` join.
    """
    template: str
    args: dict[str, str | list[str]]
    collect: str | None = None
    join: Join = Join.INNER

    @property
    def all_args(self):
        if all(isinstance(v, str) for v in self.args.values()):
            yield self.args
            return

        if self.join == Join.INNER:
            for v in zip(
                *map(
                    lambda x: repeat(x) if isinstance(x, str) else x,
                    self.args.values(),
                )
            ):
                yield dict(zip(self.args.keys(), v))

        else:  # cartesian product
            for v in product(
                *map(lambda x: [x] if isinstance(x, str) else x, self.args.values())
            ):
                yield dict(zip(self.args.keys(), v))


@dataclass
class Program:
    """A Brei program.

    Members:

      - task: list of tasks.
      - environment: variables.
      - template: set of templates.
      - call: list of calls to templates.
      - include: list of includes.
      - runner: extra configured task runners.
    """
    task: list[TaskProxy] = field(default_factory=list)
    environment: dict[str, str] = field(default_factory=dict)
    template: dict[str, Template] = field(default_factory=dict)
    call: list[TemplateCall] = field(default_factory=list)
    include: list[str] = field(default_factory=list)
    runner: dict[str, Runner] = field(default_factory=dict)

    @staticmethod
    def read(path: Path, section: str | None = None) -> Program:
        return read_from_file(Program, path, section)


def tasks_from_call(template: Template, call: TemplateCall) -> list[TaskProxy]:
    tasks = [template.call(args) for args in call.all_args]
    if call.collect:
        targets = list(chain.from_iterable(t.all_targets for t in tasks))
        collection = TaskProxy([], targets, name=call.collect)
        return tasks + [collection]
    else:
        return tasks


async def resolve_delayed(db: TaskDB, tasks: list[TaskProxy]) -> list[TaskProxy]:
    """Resolve `tasks` (substituting variables in targets).

    Returns: list of unresolvable tasks.
    """
    async def resolve(task: TaskProxy) -> TaskProxy | None:
        if not db.is_resolvable(task.all_targets):
            return task
        tt = await db.resolve_object(task)
        db.add(TemplateTask([], [], tt))
        return None

    return [t for t in await asyncio.gather(*map(resolve, tasks)) if t]


async def resolve_tasks(program: Program) -> TaskDB:
    """Resolve a program. A resolved program has all of its includes and
    template calls done, so that only tasks remains. In order to resolve
    a program, some tasks may need to be run. Variables that appear in
    the `creates` field of a task (aka targets), will be resolved eagerly.

    Returns: TaskDB instance.
    """
    db = TaskDB()
    template_index = dict()

    async def go(program: Program):
        for var, template in program.environment.items():
            db.add(TemplateVariable([Variable(var)], [], template))

        task_templates = copy(program.task)
        template_index.update(program.template)
        delayed_calls: list[TemplateCall] = []
        delayed_templates: list[TaskProxy] = []

        db.runners.update(program.runner)

        for c in program.call:
            if c.template not in template_index:
                log.debug(
                    "template `%s` not available, waiting for includes to resolve",
                    c.template,
                )
                delayed_calls.append(c)
                continue

            task_templates.extend(tasks_from_call(template_index[c.template], c))

        for tt in task_templates:
            # we could check for resolvability here, but I don't like the
            # idea that order then matters. this way the rule is:
            # > if a task has a templated target, those variables should be
            # > resolvable after all other tasks were added, seeing that the
            # > task to resolve these variables can't have templated targets
            # > themselves.
            if gather_args(tt.all_targets):
                delayed_templates.append(tt)
            else:
                db.add(TemplateTask([], [], tt))

        delayed_templates = await resolve_delayed(db, delayed_templates)

        for inc in program.include:
            incp = Path(await db.resolve_object(inc))
            if incp in db.index:
                await db.run(incp, db=db)
            if not incp.exists():
                raise MissingInclude(incp)

            prg = Program.read(incp)
            await go(prg)

        for c in delayed_calls:
            if c.template not in template_index:
                log.debug(
                    "template `%s` still not available, now this is an error", c.template
                )
                raise MissingTemplate(c.template)

            for tt in tasks_from_call(template_index[c.template], c):
                if gather_args(tt.creates):
                    delayed_templates.append(tt)
                else:
                    db.add(TemplateTask([], [], tt))

        delayed_templates = await resolve_delayed(db, delayed_templates)
        if delayed_templates:
            unresolvable = [p for t in delayed_templates for p in t.creates if not db.is_resolvable(p)]
            raise UserError(f"Task has unresolvable targets: {unresolvable}")

        return db

    return await go(program)

from __future__ import annotations

from dataclasses import dataclass

from pathlib import Path
from datetime import datetime

import os

from .construct import construct, FromStr, read_from_file 


def normal_relative(path: Path) -> Path:
    return path.resolve()  # .relative_to(Path.cwd())


@dataclass
class FileStat:
    path: Path
    modified: datetime

    @staticmethod
    def from_path(path: Path):
        stat = os.stat(path)
        return FileStat(path, datetime.fromtimestamp(stat.st_mtime))

    def __lt__(self, other: FileStat) -> bool:
        return self.modified < other.modified


def stat(path: Path) -> FileStat:
    path = normal_relative(path)
    return FileStat.from_path(path)

from dataclasses import dataclass
import time
from contextlib import asynccontextmanager

@dataclass
class Elapsed:
    elapsed: float | None = None

@asynccontextmanager
async def timer():
    e = Elapsed()
    t = time.perf_counter()
    yield e
    e.elapsed = time.perf_counter() - t

import typing
from typing import Any, Self, Union, TypeVar, TypeGuard, Type, Optional, cast
import types

import tomllib
import json

from dataclasses import is_dataclass
from enum import Enum
from pathlib import Path

from .errors import HelpfulUserError, InputError

T = TypeVar("T")


def isgeneric(annot):
    return typing.get_origin(annot) and hasattr(annot, "__args__")


class FromStr:
    @classmethod
    def from_str(cls, _: str) -> Self:
        raise NotImplementedError()


def construct(annot: Any, json: Any) -> Any:
    try:
        return _construct(annot, json)
    except (AssertionError, ValueError) as e:
        raise InputError(annot, json) from e


def is_object_type(dtype: Type[Any]) -> TypeGuard[Type[dict[str, Any]]]:
    return (
        isgeneric(dtype)
        and typing.get_origin(dtype) is dict
        and typing.get_args(dtype)[0] is str
    )


def is_optional_type(dtype: Type[Any]) -> TypeGuard[Type[Optional[Any]]]:
    return (
        isgeneric(dtype)
        and typing.get_origin(dtype) is Union
        and typing.get_args(dtype)[1] is types.NoneType
    )


def _construct(annot: Type[T], json: Any) -> T:
    """Construct an object from a given type from a JSON stream.

    The `annot` type should be one of: str, int, list[T], Optional[T],
    or a dataclass, and the JSON data should match exactly the given
    definitions in the dataclass hierarchy.
    """
    if annot is bool:
        assert isinstance(json, bool)
        return cast(T, json)
    if annot is str:
        assert isinstance(json, str)
        return cast(T, json)
    if annot is int:
        assert isinstance(json, int)
        return cast(T, json)
    if is_object_type(annot):
        assert isinstance(json, dict)
        return cast(
            T, {k: construct(typing.get_args(annot)[1], v) for k, v in json.items()}
        )
    if annot is Any:
        return cast(T, json)
    # if annot is dict or isgeneric(annot) and typing.get_origin(annot) is dict:
    #    assert isinstance(json, dict)
    #    return json
    if annot is Path and isinstance(json, str):
        return cast(T, Path(json))
    if isgeneric(annot) and typing.get_origin(annot) is list:
        assert isinstance(json, list)
        return cast(T, [construct(typing.get_args(annot)[0], item) for item in json])
    if is_optional_type(annot):
        if json is None:
            return cast(T, None)
        else:
            return cast(T, construct(typing.get_args(annot)[0], json))
    if isgeneric(annot) and typing.get_origin(annot) is types.UnionType:
        for dtype in typing.get_args(annot):
            try:
                return cast(T, _construct(dtype, json))
            except ValueError:
                continue
            except AssertionError:
                continue
        raise ValueError("None of the choices in type union match data.")
    if type(annot) is type and issubclass(annot, FromStr) and isinstance(json, str):
        return cast(T, annot.from_str(json))
    if is_dataclass(annot):
        assert isinstance(json, dict)
        arg_annot = typing.get_type_hints(annot)
        # assert all(k in json for k in arg_annot)
        args = {k: construct(arg_annot[k], json[k]) for k in json}
        return cast(T, annot(**args))
    if isinstance(json, str) and isinstance(annot, type) and issubclass(annot, Enum):
        options = {opt.name.lower(): opt for opt in annot}
        assert json.lower() in options
        return cast(T, options[json.lower()])
    raise ValueError(f"Couldn't construct {annot} from {repr(json)}")


def read_from_file(data_type: Type[T], path: Path, section: Optional[str] = None) -> T:
    """Read a config from given `path` in given `section`. The path should refer to
    a TOML or JSON file that should decode to a `Config` object. If `section` is given, only
    that section is decoded to a `Config` object. The `section` string may contain
    periods to indicate deeper nesting.

    Example:

    ```python
    read_from_file(Config, Path("./pyproject.toml"), "tool.loom")
    ```
    """
    if not path.exists():
        raise HelpfulUserError(f"File not found: {path}")
    with open(path, "rb") as f:
        if path.suffix == ".toml":
            data: Any = tomllib.load(f)
        elif path.suffix == ".json":
            data = json.load(f)
        else:
            raise HelpfulUserError(f"Unrecognized file format: {path}")

    try:
        if section is not None:
            for s in section.split("."):
                data = data[s]
    except KeyError as e:
        raise HelpfulUserError(
            f"Data file `{path}` should contain section `{section}`."
        ) from e

    return construct(data_type, data)