import ast
import json
import re
import subprocess
from pathlib import Path
from shutil import copy2, rmtree
from tempfile import TemporaryDirectory

from ..config import ValidatedConfig
from ..format import format_pyi_tree
from ..fs import Kind, Walker, remove, walk_dir
from ..source import AstConfig, Source


def generate_copy(inp: Path, out: Path):
    # in copy mode we just copy the pyi file as-is

    out.parent.mkdir(exist_ok=True, parents=True)
    copy2(inp, out)


def run_pyright(tmp: Path, base_path: Path, module_name: str, config: ValidatedConfig) -> Path:
    cfg = tmp / "config.json"
    cfg.write_text(
        json.dumps(
            {
                "extraPaths": [
                    str(base_path),
                ],
                # "include": [],
            }
        )
    )

    major, minor = config.get_python_version()
    subprocess.run(
        [
            "pyright",
            "--createstub",
            module_name,
            # "--verbose",
            "--pythonversion",
            f"{major}.{minor}",
            "-p",
            str(cfg),
        ],
        cwd=tmp,
        check=True,
    )

    return tmp / "typings"


def generate_pyright(inp: Path, out: Path, is_file: bool, config: ValidatedConfig):
    # otherwise we invoke pyright with the right base directory to do a src import from
    with TemporaryDirectory() as tmpdir:
        tmp = Path(tmpdir)

        # determine which module name to import from which directory
        base_path = inp.parent
        module_name = inp.stem

        # call pyright
        pyright_out = run_pyright(tmp, base_path, module_name, config)

        if is_file:
            # identify where pyright put the output, and normalize it
            gen = pyright_out / module_name / "__init__.pyi"
            normalize_out_file(gen, config)

            # copy to output
            out.parent.mkdir(parents=True, exist_ok=True)
            gen.rename(out)
        else:
            # identify where pyright put the output, and normalize it
            gen = tmp / "typings" / module_name
            for child in walk_dir(gen):
                if child.is_file() and child.suffix == ".pyi":
                    normalize_out_file(child, config)

            # copy to output
            out.parent.mkdir(parents=True, exist_ok=True)
            rmtree(out, ignore_errors=True)
            gen.rename(out)


def generate(inp_root: Path, out_root: Path, config: ValidatedConfig):
    # pylint: disable=duplicate-code

    walker = Walker(inp_root, out_root)

    for event in walker.walk():
        inp = event.inp_path
        out = event.out_path

        match event.kind:
            case Kind.ROOT:
                print(f"Clean {event.out_rel_pattern}")
                remove(out)

                if inp.is_dir() or (inp.suffix in (".py", ".pyi")):
                    print(f"Stub {event.inp_rel_pattern} -> {event.out_rel_pattern}")
                    generate_pyright(inp, out, inp.is_file(), config)
            case Kind.COPY:
                print(f"Copy {event.out_rel_pattern}")
                generate_copy(inp, out)
            case Kind.STUB:
                pass

    format_pyi_tree(walker, config)


def normalize(content: str, config: ValidatedConfig) -> str:
    """
    Remove generic docstring inserted by pyright, and restore position of existing docstring
    """

    source = Source(content, Path("<pyright output>"), AstConfig(feature_version=config.get_python_version()))
    parsed = source.parse_module()

    prefix = True
    lines: list[str] = []
    for stmt in parsed.body:
        if (
            prefix
            and isinstance(stmt, ast.Expr)
            and isinstance(stmt.value, ast.Constant)
            and isinstance(stmt.value.value, str)
        ):
            if stmt.value.value.strip() == "This type stub file was generated by pyright.":
                # skip this
                pass
            else:
                lines.insert(0, source.unparse_original_source(stmt))
                lines.insert(1, "")
                prefix = False
        elif isinstance(stmt, ast.Import | ast.ImportFrom):
            lines.append(source.unparse_original_source(stmt))
        else:
            if lines and lines[-1] != "":
                lines.append("")

            # keep the rest of the file as it is
            start, _ = source.unparse_offsets(stmt)
            for line in content[start:].splitlines():
                lines.append(line)
            break

    content = "\n".join(lines).strip() + "\n"

    content = re.sub(r" *# -> .*\n", "\n", content)

    return content


def normalize_out_file(path: Path, config: ValidatedConfig):
    content = path.read_text()
    content = normalize(content, config)
    path.write_text(content)
