from __future__ import annotations
import os
import json
import hashlib
import importlib.util
import inspect
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple
from jinja2 import Environment, BaseLoader, select_autoescape
from markupsafe import Markup


class LayoutEngine:
    """
    Single source of truth for:
      - Jinja environment (Caspian delimiters)
      - Template loading (file + relative-to-py)
      - Rendering (string + page + layout)
      - Nested layout system (Next.js style layout discovery and wrapping)
    """

    def __init__(self) -> None:
        self.env = Environment(
            loader=BaseLoader(),
            autoescape=select_autoescape(["html", "xml"]),
            variable_start_string="[[",
            variable_end_string="]]",
            comment_start_string="[#",
            comment_end_string="#]",
            block_start_string="[%",   # NOTE: matches your existing config
            block_end_string="%]",
        )
        self.env.filters["json"] = self._to_safe_json
        self.env.filters["dump"] = self._to_safe_json

        # Legacy alias expected elsewhere
        self.string_env = self.env

    # ---------------------------------------------------------------------
    # Jinja helpers
    # ---------------------------------------------------------------------

    @staticmethod
    def _to_safe_json(value: Any) -> Markup:
        """Safely serialize to JSON for HTML/JS embedding."""
        s = json.dumps(value, ensure_ascii=False, default=str)
        escaped = (
            s.replace("</", "<\\/")
             .replace("<!--", "<\\!--")
             .replace("\u2028", "\\u2028")
             .replace("\u2029", "\\u2029")
        )
        return Markup(escaped)

    # ---------------------------------------------------------------------
    # Core template loading & rendering (moved from render.py)
    # ---------------------------------------------------------------------

    def load_template_file(self, file_path: str) -> str:
        """Load template file directly by path."""
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"Template not found: {file_path}")
        with open(file_path, "r", encoding="utf-8") as f:
            return f.read()

    def load_template(self, py_file: str, template_name: Optional[str] = None) -> str:
        """
        Load HTML template relative to a Python file.

        Args:
            py_file: Pass __file__ from calling module
            template_name: Template filename (auto-detected if None)
        """
        dir_path = os.path.dirname(os.path.abspath(py_file))

        if template_name is None:
            template_name = os.path.basename(py_file).replace(".py", ".html")

        html_path = os.path.join(dir_path, template_name)
        return self.load_template_file(html_path)

    def render(
        self,
        py_file: str,
        context: Optional[Dict[str, Any]] = None,
        template_name: Optional[str] = None,
    ) -> str:
        """Unified render function for pages and layouts."""
        context = context or {}
        template_str = self.load_template(py_file, template_name)
        template = self.env.from_string(template_str)
        return template.render(**context)

    def render_page(self, py_file: str, context: Optional[Dict[str, Any]] = None) -> str:
        """Render index.html with context."""
        return self.render(py_file, context, "index.html")

    def render_layout(self, py_file: str, context: Optional[Dict[str, Any]] = None) -> str:
        """Render layout.html with context."""
        if context is None:
            # Optimization: No context? Just load the raw file (preserves all tags)
            return self.load_template(py_file, "layout.html")

        # Preserve [[ children ]] tag if caller didn't provide children
        if "children" not in context:
            context["children"] = Markup("[[ children ]]")

        return self.render(py_file, context, "layout.html")

    def render_string(self, template_str: str, context: Optional[Dict[str, Any]] = None) -> str:
        """Render a template string directly."""
        context = context or {}
        template = self.env.from_string(template_str)
        return template.render(**context)

    # ---------------------------------------------------------------------
    # Layout system (existing behavior, now using this engine as source)
    # ---------------------------------------------------------------------

    @dataclass
    class LayoutInfo:
        dir_path: str
        has_py: bool
        has_html: bool

        @property
        def layout_id(self) -> str:
            return hashlib.md5(self.dir_path.encode()).hexdigest()[:8]

        @property
        def py_path(self) -> str:
            return os.path.join(self.dir_path, "layout.py")

        @property
        def html_path(self) -> str:
            return os.path.join(self.dir_path, "layout.html")

    @dataclass
    class LayoutResult:
        content: str
        props: Dict[str, Any] = field(default_factory=dict)

    def discover_layouts(self, route_dir: str) -> List["LayoutEngine.LayoutInfo"]:
        """Find all layouts from root to route (Next.js style)."""
        route_dir = route_dir.replace("\\", "/")
        rel_path = os.path.relpath(route_dir, "src/app").replace("\\", "/")

        layouts: List[LayoutEngine.LayoutInfo] = []

        def check_dir(dir_path: str) -> Optional[LayoutEngine.LayoutInfo]:
            dir_path = dir_path.replace("\\", "/")
            html_exists = os.path.exists(os.path.join(dir_path, "layout.html"))
            py_exists = os.path.exists(os.path.join(dir_path, "layout.py"))

            if html_exists or py_exists:
                return LayoutEngine.LayoutInfo(
                    dir_path=dir_path, has_py=py_exists, has_html=html_exists
                )
            return None

        # Root layout
        root = check_dir("src/app")
        if root:
            layouts.append(root)

        # Path segments
        current = "src/app"
        for segment in rel_path.split("/"):
            if segment in (".", "src", "app", ""):
                continue
            current = os.path.join(current, segment).replace("\\", "/")
            found = check_dir(current)
            if found:
                layouts.append(found)

        return layouts

    def load_layout_module(self, py_path: str):
        """Dynamically load layout.py module."""
        if not os.path.exists(py_path):
            return None

        module_name = f"layout_{hashlib.md5(py_path.encode()).hexdigest()}"
        spec = importlib.util.spec_from_file_location(module_name, py_path)

        if spec and spec.loader:
            module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(module)
            return module
        return None

    def resolve_layout(self, layout: "LayoutEngine.LayoutInfo", context: Dict[str, Any]) -> "LayoutEngine.LayoutResult":
        """
        Resolve layout content:
        1. layout.py exists → execute it (MUST have layout() function)
        2. layout.html only → load it directly
        3. Neither → passthrough
        """
        passthrough = "[[ children | safe ]]"

        # Priority 1: layout.py
        if layout.has_py:
            module = self.load_layout_module(layout.py_path)

            if module and not hasattr(module, "layout"):
                raise AttributeError(
                    f"The file '{layout.py_path}' is missing the required 'def layout():' function."
                )

            if module:
                # Auto-extract module-level variables (Single Source of Truth)
                module_props: Dict[str, Any] = {}

                if hasattr(module, "title"):
                    module_props["title"] = getattr(module, "title")
                if hasattr(module, "description"):
                    module_props["description"] = getattr(
                        module, "description")

                if hasattr(module, "metadata") and isinstance(module.metadata, dict):
                    module_props.update(module.metadata)

                sig = inspect.signature(module.layout)
                result = module.layout(
                    context) if sig.parameters else module.layout()

                if result is None:
                    if layout.has_html:
                        return LayoutEngine.LayoutResult(self.load_template_file(layout.html_path), module_props)
                    return LayoutEngine.LayoutResult(passthrough, module_props)

                # Case A: User returns string (HTML)
                if isinstance(result, str):
                    return LayoutEngine.LayoutResult(result, module_props)

                # Case B: User returns tuple (HTML, props)
                if isinstance(result, tuple) and len(result) == 2:
                    combined_props = {**module_props, **(result[1] or {})}
                    return LayoutEngine.LayoutResult(result[0], combined_props)

                # Case C: User returns dict (Context for html file)
                if isinstance(result, dict):
                    content = self.load_template_file(
                        layout.html_path) if layout.has_html else passthrough
                    return LayoutEngine.LayoutResult(content, {**module_props, **result})

            if layout.has_html:
                return LayoutEngine.LayoutResult(self.load_template_file(layout.html_path))
            return LayoutEngine.LayoutResult(passthrough)

        # Priority 2: layout.html only
        if layout.has_html:
            return LayoutEngine.LayoutResult(self.load_template_file(layout.html_path))

        return LayoutEngine.LayoutResult(passthrough)

    def load_layout(self, py_or_html_path: str) -> str:
        """Load layout or page HTML file."""
        if py_or_html_path.endswith(".html"):
            return self.load_template_file(py_or_html_path)

        dir_path = os.path.dirname(os.path.abspath(py_or_html_path))
        base_name = os.path.splitext(os.path.basename(py_or_html_path))[0]
        html_path = os.path.join(dir_path, f"{base_name}.html")
        
        return self.load_template_file(html_path)

    def get_loading_files(self) -> str:
        """Get loading indicator HTML."""
        try:
            from .loading import get_loading_files as _get_loading
            return _get_loading()
        except ImportError:
            return ""

    def _inject_meta(self, html: str, layout_id: str) -> str:
        """Inject root layout meta tag."""
        meta = f'<meta name="pp-root-layout" content="{layout_id}">'
        if "<head>" in html:
            return html.replace("<head>", f"<head>\n    {meta}")
        return html

    def render_with_nested_layouts(
        self,
        children: str,
        route_dir: str,
        title: Optional[str] = None,
        description: Optional[str] = None,
        context_data: Optional[Dict[str, Any]] = None,
        transform_fn=None,
    ) -> Tuple[str, str]:
        """
        Render page wrapped in nested layouts with strict precedence:
          Page > Layout > Default
        """
        context = context_data or {}

        layouts = self.discover_layouts(route_dir)
        root_id = layouts[0].layout_id if layouts else "default"

        if not layouts:
            html = children + self.get_loading_files()
            if transform_fn:
                html = transform_fn(html)
            return html, root_id

        current_html = children
        accumulated_props: Dict[str, Any] = {}

        # Render leaf → root
        for i, layout in enumerate(reversed(layouts)):
            is_root = (i == len(layouts) - 1)

            result = self.resolve_layout(layout, context)
            layer_content = result.content

            # Child layouts override parent layouts
            accumulated_props = {**result.props, **accumulated_props}

            # Page > Layout > Default
            effective_title = title if title is not None else accumulated_props.get(
                "title", "App")
            effective_desc = description if description is not None else accumulated_props.get(
                "description", "")

            layer_context = {
                **accumulated_props,
                "title": effective_title,
                "description": effective_desc,
            }

            if is_root:
                current_html += self.get_loading_files()
                layer_content = self._inject_meta(layer_content, root_id)

            template = self.env.from_string(layer_content)
            current_html = template.render(
                children=Markup(current_html),
                layout=layer_context,
                **context,
            )

        if transform_fn:
            current_html = transform_fn(current_html)

        return current_html, root_id


# -----------------------------------------------------------------------------
# Singleton + module-level backward-compatible API
# -----------------------------------------------------------------------------

_engine = LayoutEngine()

# Jinja env exports
env = _engine.env
string_env = _engine.string_env

# Render exports (moved here from render.py)


def load_template_file(file_path: str) -> str:
    return _engine.load_template_file(file_path)


def load_template(py_file: str, template_name: Optional[str] = None) -> str:
    return _engine.load_template(py_file, template_name)


def render(py_file: str, context: Optional[Dict[str, Any]] = None, template_name: Optional[str] = None) -> str:
    return _engine.render(py_file, context, template_name)


def render_page(py_file: str, context: Optional[Dict[str, Any]] = None) -> str:
    return _engine.render_page(py_file, context)


def render_layout(py_file: str, context: Optional[Dict[str, Any]] = None) -> str:
    return _engine.render_layout(py_file, context)


def render_string(template_str: str, context: Optional[Dict[str, Any]] = None) -> str:
    return _engine.render_string(template_str, context)


# Layout system exports
LayoutInfo = LayoutEngine.LayoutInfo
LayoutResult = LayoutEngine.LayoutResult


def discover_layouts(route_dir: str) -> List[LayoutInfo]:
    return _engine.discover_layouts(route_dir)


def load_layout_module(py_path: str):
    return _engine.load_layout_module(py_path)


def resolve_layout(layout: LayoutInfo, context: Dict[str, Any]) -> LayoutResult:
    return _engine.resolve_layout(layout, context)


def get_loading_files() -> str:
    return _engine.get_loading_files()


def render_with_nested_layouts(
    children: str,
    route_dir: str,
    title: Optional[str] = None,
    description: Optional[str] = None,
    context_data: Optional[Dict[str, Any]] = None,
    transform_fn=None,
) -> Tuple[str, str]:
    return _engine.render_with_nested_layouts(
        children=children,
        route_dir=route_dir,
        title=title,
        description=description,
        context_data=context_data,
        transform_fn=transform_fn,
    )


# Back-compat alias (your old __all__ referenced this name)
render_with_layouts = render_with_nested_layouts


def load_layout(py_or_html_path: str) -> str:
    return _engine.load_layout(py_or_html_path)


def load_page(py_file_path):
    """Load HTML file next to the Python file"""
    dir_path = os.path.dirname(py_file_path)
    html_path = os.path.join(dir_path, 'index.html')
    with open(html_path, 'r', encoding='utf-8') as f:
        return f.read()


__all__ = [
    # Engine (optional export)
    "LayoutEngine",

    # Jinja env
    "env",
    "string_env",

    # Render (single truth)
    "load_template_file",
    "load_template",
    "render",
    "render_page",
    "render_layout",
    "render_string",

    # Layout system
    "LayoutInfo",
    "LayoutResult",
    "discover_layouts",
    "resolve_layout",
    "load_layout_module",
    "render_with_nested_layouts",
    "render_with_layouts",
    "load_layout",
]
