from __future__ import annotations

import logging
import os
import textwrap
from collections.abc import Callable
from pathlib import Path
from typing import Any

import tomllib

from yera.config.config_utils import extract_function_config
from yera.config.function_config import FunctionConfig, ParameterInfo, ParameterKind
from yera.config.tool_config import ToolConfig

TOOLS_CFG_FILENAME = "tools_cfg.toml"

logger = logging.getLogger(__name__)


def _yera_home() -> Path:
    home = Path(os.getenv("YERA_HOME"))
    if not home.exists():
        home = Path.home() / ".yera"
    home.mkdir(parents=True, exist_ok=True)
    return home


def _tools_cfg_path() -> Path:
    return _yera_home() / TOOLS_CFG_FILENAME


def extract_tool_config(func: Callable, security_rank: int) -> ToolConfig:
    return ToolConfig(
        function_config=extract_function_config(func),
        security_rank=security_rank,
    )


def reconstruct_tool_function(tool_config: ToolConfig) -> Callable:
    """Reconstruct a callable function from a ToolConfig."""
    fn_cfg = tool_config.function_config
    fn_name = fn_cfg.name

    # Build parameter list from ParameterInfo objects
    params = []
    for param_info in fn_cfg.parameters:
        param_str = _build_parameter_string(param_info)
        params.append(param_str)

    params_str = ", ".join(params)
    ret_typ = fn_cfg.output_signature
    body = fn_cfg.content
    fn_source = f"def {fn_name}({params_str}) -> {ret_typ}:\n" + textwrap.indent(
        body + "\n", "    "
    )

    namespace: dict[str, Any] = {}
    try:
        exec(fn_source, {}, namespace)  # noqa: S102 - local dev execution only, not prod infra
    except Exception as e:
        raise ValueError(f"Failed to reconstruct function '{fn_name}': {e}") from e
    return namespace[fn_name]


def _build_parameter_string(param_info: ParameterInfo) -> str:
    """Build a parameter string from ParameterInfo.

    Handles all parameter kinds:
    - POSITIONAL_OR_KEYWORD: name: type = default
    - VAR_POSITIONAL: *name: type
    - KEYWORD_ONLY: *, name: type = default
    - VAR_KEYWORD: **name: type
    """
    name = param_info.name
    type_annotation = param_info.type_annotation
    kind = param_info.kind
    default = param_info.default

    match kind:
        case ParameterKind.VAR_POSITIONAL:
            return f"*{name}: {type_annotation}"
        case ParameterKind.VAR_KEYWORD:
            return f"**{name}: {type_annotation}"
        case ParameterKind.KEYWORD_ONLY:
            if default is not None:
                default_str = _format_default_value(default)
                return f"{name}: {type_annotation} = {default_str}"
            return f"{name}: {type_annotation}"
        case ParameterKind.POSITIONAL_OR_KEYWORD:
            if default is not None:
                default_str = _format_default_value(default)
                return f"{name}: {type_annotation} = {default_str}"
            return f"{name}: {type_annotation}"


def _format_default_value(default_val: Any) -> str:
    """Format a default value for use in function signature.

    Handles different types of default values with proper escaping.
    """
    match default_val:
        case str():
            # Escape quotes in string defaults
            escaped = default_val.replace('"', '\\"')
            return f'"{escaped}"'
        case None:
            return "None"
        case bool() | int() | float():
            return str(default_val)
        case _:
            return repr(default_val)


def save_tool(
    func: Callable, security_rank: int | None = None, cfg_path: Path | str | None = None
) -> None:
    """Save a tool to the tools configuration file."""
    if security_rank is None:
        raise ValueError("security_rank is required")
    tool_cfg = extract_tool_config(func, security_rank=security_rank)
    fn_cfg = tool_cfg.function_config

    # Determine the TOML section path for atlas addressing
    if fn_cfg.module_path:
        tool_path = f"{fn_cfg.module_path}.{fn_cfg.name}"
    else:
        tool_path = fn_cfg.name

    if cfg_path is None:
        path = _tools_cfg_path()
    else:
        path = Path(cfg_path)

    # Load existing tools as flat structure
    existing_tools: dict[str, ToolConfig] = {}
    if path.exists():
        existing_tools = load_tools_config(path)

    # Upsert: add new tool or update existing one
    existing_tools[tool_path] = tool_cfg

    # Write all tools back to TOML
    _write_tools_cfg(existing_tools, path)


def load_tools_config(cfg_path: Path | str | None = None) -> dict[str, ToolConfig]:
    """Load tools configuration from TOML format."""
    if cfg_path is None:
        path = _tools_cfg_path()
    else:
        path = Path(cfg_path)

    if not path.exists():
        return {}

    with open(path, "rb") as f:
        data = tomllib.load(f)

    tools_section = data.get("tools", {})
    if not isinstance(tools_section, dict):
        return {}

    # Flatten the nested structure into dot-notation keys
    flattened = _flatten_dict(tools_section)

    # Parse each tool configuration
    result: dict[str, ToolConfig] = {}
    for path, tool_data in flattened.items():
        if _is_tool_config(tool_data):
            tool_cfg = _parse_tool_config(tool_data, path)
            if tool_cfg is not None:
                result[path] = tool_cfg

    return result


def _flatten_dict(nested_dict: dict) -> dict[str, Any]:
    """Flatten a nested dictionary using dot notation for keys.

    Only flattens the structure, not the tool configuration dictionaries themselves.
    """
    result = {}
    # Iterative stack-based approach avoids recursion depth issues
    stack = [(nested_dict, "")]

    while stack:
        current_dict, current_key = stack.pop()

        for key, value in current_dict.items():
            new_key = f"{current_key}.{key}" if current_key else key

            if isinstance(value, dict):
                # Check if this is a tool config (has security_rank)
                if "security_rank" in value:
                    # This is a tool config, keep it as-is
                    result[new_key] = value
                else:
                    # This is a nested section, add to stack for processing
                    stack.append((value, new_key))
            else:
                # This is a leaf value (shouldn't happen in our case)
                result[new_key] = value

    return result


def _is_tool_config(value: Any) -> bool:
    """Check if a value represents a tool configuration."""
    return isinstance(value, dict) and "security_rank" in value


def _parse_tool_config(tool_data: dict, path: str) -> ToolConfig | None:
    """Parse a single tool configuration from TOML data."""
    try:
        raw_data = tool_data.copy()
        security_rank = raw_data.pop("security_rank")

        # Parse parameters from TOML inline table array
        parameters_data = raw_data.get("parameters", [])
        raw_data["parameters"] = _parse_parameters_list(parameters_data)

        function_config = FunctionConfig.model_validate(raw_data)
        return ToolConfig(function_config=function_config, security_rank=security_rank)

    except (KeyError, ValueError, TypeError) as e:
        logger.warning("Skipping malformed tool '%s': %s", path, e)
        return None


def _parse_parameters_list(parameters_data) -> list:
    """Parse parameters list from native TOML inline table array to ParameterInfo objects."""
    from yera.config.function_config import ParameterInfo, ParameterKind

    parameters = []
    for param_data in parameters_data:
        # Parse default value - handle "None" string conversion back to None
        default_val = param_data.get("default")
        if default_val is not None:
            default_val = _deserialize_default_from_toml(default_val)

        param_info = ParameterInfo(
            name=param_data["name"],
            type_annotation=param_data["type_annotation"],
            kind=ParameterKind(param_data["kind"]),
            default=default_val,
        )
        parameters.append(param_info)

    return parameters


def _deserialize_default_from_toml(value: Any) -> Any:
    """Deserialize a default value from native TOML format.

    Converts "None" string back to None, handles other types.
    """
    if value == "None":
        return None

    if not isinstance(value, str):
        return value
    s = value

    # JSON for complex types
    if s.startswith(("{", "[")):
        try:
            import json

            return json.loads(s)
        except Exception:
            return s

    # Booleans
    s_lower = s.lower()
    if s_lower in ("true", "false"):
        return s_lower == "true"

    # Numbers (int or float)
    try:
        return float(s) if "." in s else int(s)
    except ValueError:
        return s


def _toml_escape(value: str) -> str:
    """Escape a string for use in TOML format.

    Handles backslashes, newlines, and quotes properly.
    """
    # Escape backslashes first, then quotes, then newlines
    return value.replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n")


def _write_tools_cfg(tools: dict[str, ToolConfig], path: Path) -> None:
    """Write tools configuration to TOML format using native TOML inline table arrays.

    Serializes the new ParameterInfo structure to native TOML format.
    """
    lines: list[str] = ["[tools]"]

    for section_path, tool_cfg in tools.items():
        lines.append(f"[tools.{section_path}]")
        fn_cfg = tool_cfg.function_config

        # Basic function properties
        # Use basic strings with escaping for content
        escaped_content = _toml_escape(fn_cfg.content)
        lines.append(f'content = "{escaped_content}"')
        lines.append(f'name = "{_toml_escape(fn_cfg.name)}"')
        lines.append(f'output_signature = "{_toml_escape(fn_cfg.output_signature)}"')

        # Add module_path if present
        if fn_cfg.module_path:
            lines.append(f'module_path = "{_toml_escape(fn_cfg.module_path)}"')

        lines.append(f"security_rank = {tool_cfg.security_rank}")

        # Serialize parameters as inline table array
        if fn_cfg.parameters:
            lines.append("parameters = [")
            param_lines = []
            for param_info in fn_cfg.parameters:
                # Convert None to "None" string for TOML compatibility
                default_str = _serialize_default_for_toml(param_info.default)
                escaped_default = _toml_escape(default_str)

                param_dict = f'{{name = "{_toml_escape(param_info.name)}", type_annotation = "{_toml_escape(param_info.type_annotation)}", kind = "{param_info.kind.value}", default = "{escaped_default}"}}'
                param_lines.append(param_dict)

            # Join parameters with commas and add to lines
            lines.append(",\n".join(param_lines))
            lines.append("]")

        lines.append("")

    text = "\n".join(lines).rstrip() + "\n"
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        f.write(text)


def _serialize_default_for_toml(default_val: Any) -> str:
    """Serialize a default value for native TOML format.

    Converts Python values to TOML-compatible strings.
    None becomes "None" string for TOML compatibility.
    """
    match default_val:
        case None:
            return "None"
        case str():
            return default_val
        case bool() | int() | float():
            return str(default_val)
        case list() | dict():
            import json

            return json.dumps(default_val)
        case _:
            return repr(default_val)
