from pathlib import Path
import re
from typing import Sequence, Optional, Dict
from jinja2 import Environment, FileSystemLoader
import importlib.util
from dataclasses import dataclass
import logging
from hashlib import md5


@dataclass
class PromptTemplate:
    name: str
    variables: set[str]
    content: str
    source_path: Path


class PromptCompiler:
    def __init__(
        self,
        working_dir: Path = Path("."),
        file_extension: str = "mdx",
        output_dir: Optional[Path] = None,
        output_file: str = "prompts.py",
    ):
        """
        Initialize the compiler.

        Args:
            working_dir: Directory to start searching from (default: current directory)
            file_extension: File extension to look for (default: mdx)
            output_dir: Where to put generated files (default: working_dir/__prompts__)
            output_file: Name of the generated Python file (default: prompts.py)
        """
        self.working_dir = working_dir
        self.file_extension = file_extension.lstrip(".")
        self.output_dir = output_dir or (working_dir / "__prompts__")
        self.output_file = output_file
        self.output_path = self.output_dir / self.output_file
        self.output_dir.mkdir(exist_ok=True, parents=True)
        self._cache: dict[Path, str] = {}

        # Setup Jinja2 environment with custom delimiters
        self.env = Environment(
            loader=FileSystemLoader(working_dir),
            variable_start_string="<var",
            variable_end_string=">",
            comment_start_string="<!--",
            comment_end_string="-->",
        )

    def compile_all(self) -> None:
        """Batch compile all template files"""
        try:
            templates = self._collect_templates()
            if not templates:
                logging.warning(
                    f"No .{self.file_extension} files found in {self.working_dir}"
                )
                return
            self._generate_module(templates)
        except Exception as e:
            logging.error(f"Failed to compile templates: {e}")

    def compile_file(self, path: Path) -> None:
        """Compile single template file"""
        if path.suffix != f".{self.file_extension}":
            return

        if not self._should_compile(path):
            logging.debug(f"Skipping unchanged file: {path}")
            return

        try:
            # Get all existing templates
            existing_templates = self._collect_templates()

            # Update or add the changed template
            new_template = self._parse_template(path)

            # Replace or add the template
            templates = [
                new_template if t.name == new_template.name else t
                for t in existing_templates
            ]
            if new_template.name not in {t.name for t in existing_templates}:
                templates.append(new_template)

            self._generate_module(templates)
        except Exception as e:
            logging.error(f"Failed to compile {path}: {e}")

    def _should_compile(self, path: Path) -> bool:
        """Check if file needs recompilation based on content hash"""
        content = path.read_text()
        content_hash = md5(content.encode()).hexdigest()
        if self._cache.get(path) == content_hash:
            return False
        self._cache[path] = content_hash
        return True

    def _collect_templates(self) -> Sequence[PromptTemplate]:
        """Find and parse all template files recursively"""
        return [
            self._parse_template(p)
            for p in self.working_dir.rglob(f"*.{self.file_extension}")
        ]

    def _parse_template(self, path: Path) -> PromptTemplate:
        """Parse template file and extract variables"""
        content = path.read_text()
        # Use original variable extraction pattern
        pattern = re.compile(r"<var\s+(\w+)\s*>")
        variables = set(pattern.findall(content))

        # Use relative path for name to avoid collisions
        relative_path = path.relative_to(self.working_dir)
        name = (
            str(relative_path.with_suffix(""))
            .replace("/", "_")
            .replace("\\", "_")
        )

        return PromptTemplate(
            name=name,
            variables=variables,
            content=content,
            source_path=path,
        )

    def _generate_module(self, templates: Sequence[PromptTemplate]) -> None:
        """Generate Python module with template functions"""
        code = [
            """from jinja2 import Environment, FileSystemLoader
from pathlib import Path

_env = Environment(
    loader=FileSystemLoader(Path(__file__).parent.parent),
    variable_start_string='{',
    variable_end_string='}',
    comment_start_string='{/*',
    comment_end_string='*/}'
)
"""
        ]

        # Sort templates by name for consistent output
        sorted_templates = sorted(templates, key=lambda t: t.name)

        for template in sorted_templates:
            params = ", ".join(
                f"{var}: str" for var in sorted(template.variables)
            )
            template_args = ", ".join(
                f"{var}={var}" for var in sorted(template.variables)
            )

            code.append(f"""
def {template.name}({params}) -> str:
    \"\"\"
    Generated from: {template.source_path.relative_to(self.working_dir)}
    
    Template:
    {template.content}
    \"\"\"
    template = _env.get_template("{template.source_path.relative_to(self.working_dir)}")
    return template.render({template_args})
""")

        self.output_path.write_text("\n".join(code))
        logging.info(f"Generated module at {self.output_path}")
        self._reload_module(self.output_path)

    @staticmethod
    def _reload_module(module_path: Path) -> None:
        """Force reload of generated module"""
        try:
            spec = importlib.util.spec_from_file_location(
                module_path.stem, module_path
            )
            if spec and spec.loader:
                module = importlib.util.module_from_spec(spec)
                spec.loader.exec_module(module)
        except Exception as e:
            logging.error(f"Failed to reload module: {e}")
