import logging
import sys
from functools import lru_cache
from pathlib import Path
from threading import Lock

import jedi
import jedi.api.classes
import jedi.api.environment
import jedi.cache
from pydantic import computed_field

from py_agent.models.base import FrozenModel

logger = logging.getLogger(__name__)


class JediMetadata(FrozenModel):
    name: str
    full_name: str
    module_name: str
    file_path: Path | None
    is_core: bool
    type: str
    signatures: list[str]

    @computed_field  # type: ignore[prop-decorator]
    @property
    def class_name(self) -> str | None:
        if ".".join([self.module_name, self.name]) != self.full_name:
            return self.full_name.split(".")[-2]
        return None


class JediContainer:
    @lru_cache(maxsize=128)
    @staticmethod
    def get_jedi_by_absolute_path(
        absolute_file_path: Path,
    ) -> tuple["jedi.Script", "jedi.api.environment.Environment"]:
        project_path = JediContainer.find_root_project(absolute_file_path=absolute_file_path)
        env = JediContainer._find_env(project_path)
        script = jedi.Script(
            path=absolute_file_path,
            project=jedi.Project(project_path),
            environment=env,
        )
        return script, env

    @staticmethod
    def _find_env(project_path: Path) -> "jedi.api.environment.Environment":
        venv = jedi.find_virtualenvs(project_path.rglob("*"), use_environment_vars=False)
        try:
            if result := next(venv):
                return result
        except StopIteration:
            pass
        return jedi.api.environment.get_default_environment()

    @lru_cache(maxsize=128)
    @staticmethod
    def find_root_project(absolute_file_path: Path) -> Path:
        topmost = absolute_file_path
        for parent in absolute_file_path.parents:
            if (parent / "__init__.py").exists():
                topmost = parent
            else:
                break
        return topmost.parent

    @staticmethod
    def clear_all_caches() -> None:
        JediContainer.get_jedi_by_absolute_path.cache_clear()
        JediContainer.find_root_project.cache_clear()
        jedi.cache.clear_time_caches(delete_all=True)


class ImportJediExtractor:
    lock = Lock()

    def __init__(self, absolute_file_path: str = "", project_path: str = ""):
        self.absolute_file_path: Path = Path(absolute_file_path)
        self.project_path: Path = JediContainer.find_root_project(
            absolute_file_path=Path(self.absolute_file_path)
        )
        self.script: jedi.Script | None = None
        self.env: jedi.api.environment.Environment | None = None

    def _populate_script(self) -> None:
        if self.script:
            return
        script, env = JediContainer.get_jedi_by_absolute_path(
            absolute_file_path=self.absolute_file_path
        )
        self.env = env
        self.script = script

    def get_env_path(self) -> Path:
        self._populate_script()
        return Path(self.env.path)  # type: ignore[union-attr]

    def _is_core(self, name: "jedi.api.classes.Name") -> bool:
        self._populate_script()
        if name.in_builtin_module():
            return True
        if name.module_path is None:
            return False
        return name.module_path.is_relative_to(self.env.path) or name.module_path.is_relative_to(  # type: ignore[union-attr, no-any-return]
            Path(sys.base_prefix)
        )

    def get_by_position(self, line: int, col: int) -> JediMetadata | None:
        self._populate_script()
        with ImportJediExtractor.lock:
            inferred_names = self.script.infer(line, col)  # type: ignore[union-attr]
            if not inferred_names:
                return None
            top_result: jedi.api.classes.Name = inferred_names[0]
            try:
                return JediMetadata(
                    name=top_result.name,
                    full_name=top_result.full_name,
                    module_name=top_result.module_name,
                    file_path=top_result.module_path,
                    is_core=self._is_core(top_result),
                    type=top_result.type,
                    signatures=[x.to_string() for x in top_result.get_signatures()],
                )
            except Exception:
                return None

    def get_jedi_metadata(self, name: str) -> JediMetadata | None:
        self._populate_script()
        with ImportJediExtractor.lock:
            names = self.script.search(name)  # type: ignore[union-attr]
            if not names:
                return None
            unresolved_name: jedi.api.classes.Name = names[0]
            if unresolved_name.in_builtin_module():
                return JediMetadata(
                    name=unresolved_name.name,
                    full_name=name,
                    module_name=unresolved_name.module_name,
                    file_path=unresolved_name.module_path,
                    is_core=True,
                    type=unresolved_name.type,
                    signatures=[],
                )
            resolved_names = unresolved_name.infer()
            if not resolved_names:
                return JediMetadata(
                    name=unresolved_name.name,
                    full_name=name,
                    module_name=unresolved_name.module_name,
                    file_path=unresolved_name.module_path,
                    is_core=self._is_core(unresolved_name),
                    type=unresolved_name.type,
                    signatures=[x.to_string() for x in unresolved_name.get_signatures()],
                )
            resolved_name: jedi.api.classes.Name = resolved_names[0]
            return JediMetadata(
                name=resolved_name.name,
                full_name=name,
                module_name=resolved_name.module_name,
                file_path=resolved_name.module_path,
                is_core=self._is_core(resolved_name),
                type=resolved_name.type,
                signatures=[x.to_string() for x in resolved_name.get_signatures()],
            )

    def get_absolute_import_path(
        self, function_name: str, class_name: str | None = None
    ) -> str | None:
        self._populate_script()
        full_name = function_name
        if class_name:
            full_name = ".".join([class_name, function_name])
        with ImportJediExtractor.lock:
            names = list(self.script.search(full_name))  # type: ignore[union-attr]
        if not names:
            return None
        return str(names[0].module_name) if names[0].module_name else None
