"""
NLS Lockfile - Determinism guarantees for NLS compilation

Generates and verifies .nl.lock files to ensure reproducible builds.
"""

import hashlib
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

from . import __version__
from .schema import NLFile, ANLU


@dataclass
class ANLULock:
    """Lock entry for a single ANLU"""
    source_hash: str
    output_hash: str
    output_lines: int
    generated_code: str = ""  # Inline generated code for zero-LLM rebuild


@dataclass
class ModuleLock:
    """Lock entry for a module"""
    source_hash: str
    anlus: dict[str, ANLULock] = field(default_factory=dict)


@dataclass
class TargetLock:
    """Lock entry for a generated target file"""
    file: str
    hash: str
    lines: int


@dataclass
class Lockfile:
    """Complete lockfile structure"""
    schema_version: int = 1
    generated_at: str = ""
    compiler_version: str = __version__
    llm_backend: str = "mock"

    modules: dict[str, ModuleLock] = field(default_factory=dict)
    targets: dict[str, TargetLock] = field(default_factory=dict)

    def __post_init__(self):
        if not self.generated_at:
            self.generated_at = datetime.now(timezone.utc).isoformat()


def hash_content(content: str) -> str:
    """Generate SHA256 hash of content"""
    return f"sha256:{hashlib.sha256(content.encode('utf-8')).hexdigest()[:12]}"


def hash_anlu(anlu: ANLU) -> str:
    """Generate deterministic hash of an ANLU"""
    # Create canonical representation
    canonical = f"{anlu.identifier}|{anlu.purpose}|{anlu.returns}"
    canonical += "|" + ",".join(f"{i.name}:{i.type}" for i in anlu.inputs)
    canonical += "|" + ",".join(anlu.depends)
    return hash_content(canonical)


def extract_function_code(generated_code: str, func_name: str) -> str:
    """
    Extract a single function's code from the generated output.

    Args:
        generated_code: Full generated Python code
        func_name: Python function name to extract

    Returns:
        The function's code including docstring, or empty string if not found
    """
    import re

    # Pattern to match function definition including decorators
    # Matches: optional decorators, def line, docstring, body until next def or end
    pattern = rf'((?:@\w+.*\n)*def {func_name}\([^)]*\)[^:]*:.*?)(?=\n(?:@|\ndef |\Z))'

    match = re.search(pattern, generated_code, re.DOTALL)
    if match:
        return match.group(1).rstrip()

    return ""


def generate_lockfile(
    nl_file: NLFile,
    generated_code: str,
    output_path: str,
    llm_backend: str = "mock"
) -> Lockfile:
    """
    Generate a lockfile for a compilation.

    Args:
        nl_file: The parsed NLFile
        generated_code: The generated Python code
        output_path: Path to the generated file
        llm_backend: Name of LLM backend used

    Returns:
        Lockfile object
    """
    lockfile = Lockfile(llm_backend=llm_backend)

    # Create module lock
    source_content = ""
    if nl_file.source_path:
        try:
            source_content = Path(nl_file.source_path).read_text()
        except (FileNotFoundError, OSError):
            pass

    module_lock = ModuleLock(source_hash=hash_content(source_content))

    # Lock each ANLU with its generated code
    for anlu in nl_file.anlus:
        func_name = anlu.python_name
        anlu_code = extract_function_code(generated_code, func_name)

        module_lock.anlus[anlu.identifier] = ANLULock(
            source_hash=hash_anlu(anlu),
            output_hash=hash_content(anlu_code) if anlu_code else hash_content(func_name),
            output_lines=anlu_code.count("\n") + 1 if anlu_code else 0,
            generated_code=anlu_code,
        )

    lockfile.modules[nl_file.module.name] = module_lock

    # Lock target
    output_lines = generated_code.count("\n") + 1
    lockfile.targets["python"] = TargetLock(
        file=output_path,
        hash=hash_content(generated_code),
        lines=output_lines
    )

    return lockfile


def write_lockfile(lockfile: Lockfile, path: Path) -> None:
    """Write lockfile to disk as YAML-like format with inline generated code"""
    lines = [
        "# DO NOT EDIT - Generated by nlsc",
        f"# {lockfile.generated_at}",
        "",
        f"schema_version: {lockfile.schema_version}",
        f"generated_at: {lockfile.generated_at}",
        f"compiler_version: {lockfile.compiler_version}",
        f"llm_backend: {lockfile.llm_backend}",
        "",
        "modules:"
    ]

    for mod_name, mod_lock in lockfile.modules.items():
        lines.append(f"  {mod_name}:")
        lines.append(f"    source_hash: {mod_lock.source_hash}")
        lines.append("    anlus:")
        for anlu_id, anlu_lock in mod_lock.anlus.items():
            lines.append(f"      {anlu_id}:")
            lines.append(f"        source_hash: {anlu_lock.source_hash}")
            lines.append(f"        output_hash: {anlu_lock.output_hash}")
            lines.append(f"        output_lines: {anlu_lock.output_lines}")
            # Write generated_code as YAML multiline literal block
            if anlu_lock.generated_code:
                lines.append("        generated_code: |")
                for code_line in anlu_lock.generated_code.split("\n"):
                    lines.append(f"          {code_line}")

    lines.append("")
    lines.append("targets:")
    for target_name, target_lock in lockfile.targets.items():
        lines.append(f"  {target_name}:")
        lines.append(f"    file: {target_lock.file}")
        lines.append(f"    hash: {target_lock.hash}")
        lines.append(f"    lines: {target_lock.lines}")

    path.write_text("\n".join(lines), encoding="utf-8")


def read_lockfile(path: Path) -> Optional[Lockfile]:
    """Read an existing lockfile (enhanced YAML-like parser with generated_code support)"""
    if not path.exists():
        return None

    content = path.read_text(encoding="utf-8")
    lockfile = Lockfile()

    current_module_name = None
    current_anlu_id = None
    current_anlu_data: dict = {}
    in_generated_code = False
    generated_code_lines: list[str] = []

    lines = content.split("\n")
    i = 0
    while i < len(lines):
        line = lines[i]
        line_stripped = line.rstrip()

        # Handle generated_code multiline block continuation FIRST
        # (must check before skipping empty lines - code blocks can have empty lines)
        if in_generated_code:
            # Check if still in generated_code block by looking at raw indentation
            raw_indent = len(line) - len(line.lstrip()) if line.strip() else len(line)
            if raw_indent >= 10 or (not line.strip() and raw_indent >= 10):
                # Still in generated_code block (indented at least 10 spaces)
                # Remove the 10-space prefix
                code_line = line[10:].rstrip() if len(line) > 10 else ""
                generated_code_lines.append(code_line)
                i += 1
                continue
            elif not line.strip():
                # Empty line at lower indent - end of block
                if current_anlu_id and current_anlu_data:
                    current_anlu_data["generated_code"] = "\n".join(generated_code_lines)
                in_generated_code = False
                generated_code_lines = []
                i += 1
                continue
            else:
                # End of generated_code block (non-empty line at lower indent)
                if current_anlu_id and current_anlu_data:
                    current_anlu_data["generated_code"] = "\n".join(generated_code_lines)
                in_generated_code = False
                generated_code_lines = []

        # Skip comments and empty lines
        if line_stripped.startswith("#") or not line_stripped.strip():
            i += 1
            continue

        # Count indentation
        indent = len(line_stripped) - len(line_stripped.lstrip())

        # Parse key-value pairs
        if ":" in line_stripped:
            key, _, value = line_stripped.strip().partition(":")
            value = value.strip()

            if indent == 0:
                # Top-level: schema_version, generated_at, modules, targets
                if key == "schema_version":
                    lockfile.schema_version = int(value) if value else 1
                elif key == "generated_at":
                    lockfile.generated_at = value
                elif key == "compiler_version":
                    lockfile.compiler_version = value
                elif key == "llm_backend":
                    lockfile.llm_backend = value
                elif key == "modules":
                    pass  # Section header
                elif key == "targets":
                    # Save current module before switching to targets
                    if current_module_name and current_anlu_id and current_anlu_data:
                        _save_anlu_to_lockfile(lockfile, current_module_name, current_anlu_id, current_anlu_data)
                    current_module_name = None
                    current_anlu_id = None
                    current_anlu_data = {}

            elif indent == 2:
                # Module name (under modules:) or target entry
                if key not in ("source_hash", "anlus", "file", "hash", "lines"):
                    # Save previous ANLU if exists
                    if current_module_name and current_anlu_id and current_anlu_data:
                        _save_anlu_to_lockfile(lockfile, current_module_name, current_anlu_id, current_anlu_data)
                        current_anlu_id = None
                        current_anlu_data = {}
                    current_module_name = key
                    # Initialize module if not exists
                    if current_module_name not in lockfile.modules:
                        lockfile.modules[current_module_name] = ModuleLock(source_hash="")

            elif indent == 4:
                # Module properties: source_hash, anlus
                if key == "source_hash" and current_module_name:
                    lockfile.modules[current_module_name].source_hash = value
                elif key == "anlus":
                    pass  # ANLUs section header

            elif indent == 6 and current_module_name:
                # ANLU ID - save previous ANLU first
                if current_anlu_id and current_anlu_data:
                    _save_anlu_to_lockfile(lockfile, current_module_name, current_anlu_id, current_anlu_data)
                current_anlu_id = key
                current_anlu_data = {}

            elif indent == 8 and current_anlu_id:
                # ANLU property
                if key == "generated_code" and value == "|":
                    # Start of multiline block
                    in_generated_code = True
                    generated_code_lines = []
                else:
                    current_anlu_data[key] = value

        i += 1

    # Save last ANLU
    if current_module_name and current_anlu_id and current_anlu_data:
        _save_anlu_to_lockfile(lockfile, current_module_name, current_anlu_id, current_anlu_data)

    return lockfile


def _save_anlu_to_lockfile(lockfile: Lockfile, module_name: str, anlu_id: str, data: dict) -> None:
    """Helper to save an ANLU lock entry to the lockfile"""
    if module_name not in lockfile.modules:
        lockfile.modules[module_name] = ModuleLock(source_hash="")

    lockfile.modules[module_name].anlus[anlu_id] = ANLULock(
        source_hash=data.get("source_hash", ""),
        output_hash=data.get("output_hash", ""),
        output_lines=int(data.get("output_lines", 0)),
        generated_code=data.get("generated_code", ""),
    )


def verify_lockfile(lockfile: Lockfile, nl_file: NLFile) -> list[str]:
    """
    Verify that a lockfile matches current sources.

    Returns:
        List of verification errors (empty if valid)
    """
    errors = []

    mod_lock = lockfile.modules.get(nl_file.module.name)
    if not mod_lock:
        errors.append(f"Module {nl_file.module.name} not in lockfile")
        return errors

    for anlu in nl_file.anlus:
        anlu_lock = mod_lock.anlus.get(anlu.identifier)
        if not anlu_lock:
            errors.append(f"ANLU {anlu.identifier} not in lockfile")
            continue

        current_hash = hash_anlu(anlu)
        if current_hash != anlu_lock.source_hash:
            errors.append(f"ANLU {anlu.identifier} has changed since lock")

    return errors


@dataclass
class RebuildResult:
    """Result of rebuilding from lockfile"""
    code: str
    used_cache: bool
    llm_calls: int
    cached_anlus: list[str]
    regenerated_anlus: list[str]


def rebuild_from_lockfile(nl_file: NLFile, lockfile: Lockfile) -> Optional[RebuildResult]:
    """
    Rebuild generated code from lockfile, using cached code where possible.

    This enables zero-LLM rebuilds when sources haven't changed.

    Args:
        nl_file: The parsed NLFile
        lockfile: The lockfile with cached generated code

    Returns:
        RebuildResult with cached/regenerated code, or None if rebuild impossible
    """
    mod_lock = lockfile.modules.get(nl_file.module.name)
    if not mod_lock:
        return None

    cached_anlus = []
    regenerated_anlus = []
    all_code_pieces = []
    llm_calls = 0

    for anlu in nl_file.anlus:
        anlu_lock = mod_lock.anlus.get(anlu.identifier)
        current_hash = hash_anlu(anlu)

        if anlu_lock and anlu_lock.source_hash == current_hash and anlu_lock.generated_code:
            # Use cached code
            cached_anlus.append(anlu.identifier)
            all_code_pieces.append(anlu_lock.generated_code)
        else:
            # Would need to regenerate (but in mock mode, we can't)
            regenerated_anlus.append(anlu.identifier)
            # For now, mark as needing regeneration
            llm_calls += 1

    if regenerated_anlus:
        # Some ANLUs need regeneration - can't do zero-LLM rebuild
        # In a real implementation, this would call the LLM
        return RebuildResult(
            code="",
            used_cache=False,
            llm_calls=llm_calls,
            cached_anlus=cached_anlus,
            regenerated_anlus=regenerated_anlus,
        )

    # Full cache hit - zero-LLM rebuild
    # Reconstruct the full module code
    module_header = f'''"""
Generated by nlsc from lockfile
Module: {nl_file.module.name}
"""

'''
    full_code = module_header + "\n\n".join(all_code_pieces)

    return RebuildResult(
        code=full_code,
        used_cache=True,
        llm_calls=0,
        cached_anlus=cached_anlus,
        regenerated_anlus=[],
    )
