"""
Core implementation of the @ai decorator and LLM synthesis engine.
"""

import inspect
import hashlib
import json
import logging
import os
import textwrap
import time
import typing
from dataclasses import dataclass, asdict
from pathlib import Path
from functools import wraps

import anthropic


logger = logging.getLogger("speq")


# ---------------------------------------------------------------------------
# Confidence scores
# ---------------------------------------------------------------------------

@dataclass
class ConfidenceScores:
    """Confidence assessment for a synthesized function."""
    spec_clarity: float          # 0-1: how unambiguous is the spec?
    implementation_confidence: float  # 0-1: how likely is the implementation correct?
    assumptions: list[str]       # things the LLM assumed that weren't in the spec
    ambiguities: list[str]       # vague parts of the spec
    missing_edge_cases: list[str]  # edge cases the spec didn't address


def _score_synthesis(func_name: str, signature: str, docstring: str,
                     generated_code: str, model: str) -> ConfidenceScores:
    """Ask the LLM to assess the spec clarity and implementation confidence."""
    client = anthropic.Anthropic()

    prompt = f"""You are a critical code reviewer evaluating a function spec and its generated implementation.

Function: def {func_name}{signature}

Specification (docstring):
\"\"\"{docstring}\"\"\"

Generated implementation:
```python
{generated_code}
```

STEP 1: List every assumption the implementation makes that the spec didn't state.
STEP 2: List every ambiguity in the spec (words or phrases that could mean multiple things).
STEP 3: List edge cases the spec doesn't address (empty inputs, zeros, negatives, boundaries, overflow, type edge cases).
STEP 4: Count: how many assumptions? how many ambiguities? how many missing edge cases?
STEP 5: Score using these calibration rules:

spec_clarity scoring (be harsh):
  0.95-1.0: spec is mathematically precise, covers ALL edge cases, zero ambiguity. Almost no spec qualifies.
  0.80-0.94: spec is clear with minor gaps. One or two unaddressed edge cases.
  0.60-0.79: spec is understandable but has several ambiguities or missing edge cases.
  0.40-0.59: spec requires significant interpretation. Multiple assumptions needed.
  0.20-0.39: spec is vague. "Handle errors appropriately" level.
  0.00-0.19: spec is nearly useless. "Do the right thing."

implementation_confidence scoring (be harsh):
  0.95-1.0: implementation is provably correct for all valid inputs. Trivial function.
  0.80-0.94: very likely correct, but one or two edge cases are hard to verify by inspection.
  0.60-0.79: probably correct for common cases, but edge cases may be wrong.
  0.40-0.59: significant uncertainty. Complex logic that's hard to verify.
  0.20-0.39: likely has bugs. Multiple risky assumptions.
  0.00-0.19: almost certainly wrong.

A function like "double x" should get 0.98/0.99. A function like "merge overlapping ranges" with a short spec should get 0.65/0.75. Vary your scores. Do NOT default to 0.8/0.9 for everything.

Respond with ONLY a JSON object (no markdown, no explanation):
{{
    "spec_clarity": <float>,
    "implementation_confidence": <float>,
    "assumptions": [<list of strings>],
    "ambiguities": [<list of strings>],
    "missing_edge_cases": [<list of strings>]
}}"""

    try:
        response = client.messages.create(
            model=model,
            max_tokens=1024,
            messages=[{"role": "user", "content": prompt}],
        )

        raw = response.content[0].text.strip()
        # strip markdown fences if present
        if raw.startswith("```"):
            raw = raw.split("\n", 1)[1] if "\n" in raw else raw[3:]
        if raw.endswith("```"):
            raw = raw[:-3]
        raw = raw.strip()

        data = json.loads(raw)
        return ConfidenceScores(
            spec_clarity=float(data.get("spec_clarity", 0.5)),
            implementation_confidence=float(data.get("implementation_confidence", 0.5)),
            assumptions=data.get("assumptions", []),
            ambiguities=data.get("ambiguities", []),
            missing_edge_cases=data.get("missing_edge_cases", []),
        )
    except Exception as e:
        logger.debug(f"speq: confidence scoring failed for '{func_name}': {e}")
        return ConfidenceScores(
            spec_clarity=-1,
            implementation_confidence=-1,
            assumptions=[],
            ambiguities=[],
            missing_edge_cases=["scoring failed"],
        )


# ---------------------------------------------------------------------------
# Test generation
# ---------------------------------------------------------------------------

def _generate_tests(func_name: str, signature: str, docstring: str,
                    generated_code: str, model: str) -> str:
    """Generate test cases from the spec, not the implementation."""
    client = anthropic.Anthropic()

    prompt = f"""Generate pytest test cases for this function based ONLY on its specification (docstring).
Do NOT look at the implementation to write tests — derive them purely from the spec.

Function: def {func_name}{signature}

Specification (docstring):
\"\"\"{docstring}\"\"\"

Write a complete, runnable pytest file. Requirements:
- Import the function as: from __function__ import {func_name}
- Test every example mentioned in the docstring
- Test edge cases implied by the docstring (empty inputs, zeros, boundaries)
- Test at least 5 cases total
- Each test should be a simple assert statement
- Use descriptive test function names
- No markdown fences, no explanation, just the Python code
- Do NOT import speq or use the @ai decorator in tests"""

    try:
        response = client.messages.create(
            model=model,
            max_tokens=2048,
            messages=[{"role": "user", "content": prompt}],
        )

        code = _clean_code(response.content[0].text)
        logger.debug(f"speq: [{func_name}] generated {len(code)} chars of tests")
        return code
    except Exception as e:
        logger.debug(f"speq: test generation failed for '{func_name}': {e}")
        return ""


# ---------------------------------------------------------------------------
# Cache
# ---------------------------------------------------------------------------

_CACHE_DIR_NAME = ".speq_cache"
_memory_cache: dict[str, typing.Callable] = {}


def _get_cache_dir(source_file: str) -> Path:
    """Return the .speq_cache directory next to the source file."""
    source_dir = Path(source_file).parent
    cache_dir = source_dir / _CACHE_DIR_NAME
    cache_dir.mkdir(exist_ok=True)
    return cache_dir


def _spec_hash(func_name: str, signature: str, docstring: str) -> str:
    """Hash the function spec so we can detect when it changes."""
    content = f"{func_name}|{signature}|{docstring}"
    return hashlib.sha256(content.encode()).hexdigest()[:16]


def _cache_prefix(func_name: str, spec_id: str) -> str:
    """Build the cache filename prefix: funcname_spechash."""
    return f"{func_name}_{spec_id}"


def _load_from_disk_cache(cache_dir: Path, func_name: str, spec_id: str) -> str | None:
    """Try to load a previously synthesized implementation from disk."""
    cache_file = cache_dir / f"{_cache_prefix(func_name, spec_id)}.py"
    if cache_file.exists():
        return cache_file.read_text()
    return None


def _load_scores_from_disk_cache(cache_dir: Path, func_name: str, spec_id: str) -> ConfidenceScores | None:
    """Try to load previously computed confidence scores from disk."""
    meta_file = cache_dir / f"{_cache_prefix(func_name, spec_id)}.meta.json"
    if meta_file.exists():
        try:
            data = json.loads(meta_file.read_text())
            if "scores" in data:
                return ConfidenceScores(**data["scores"])
        except (json.JSONDecodeError, TypeError, KeyError):
            pass
    return None


def _save_to_disk_cache(cache_dir: Path, func_name: str, spec_id: str, code: str,
                        scores: ConfidenceScores | None = None,
                        model: str | None = None,
                        tests: str | None = None) -> None:
    """Save a synthesized implementation, scores, and tests to disk."""
    prefix = _cache_prefix(func_name, spec_id)
    cache_file = cache_dir / f"{prefix}.py"
    cache_file.write_text(code)

    meta = {
        "spec_id": spec_id,
        "function": func_name,
        "created_by": "speq",
        "model": model,
    }
    if scores:
        meta["scores"] = asdict(scores)

    meta_file = cache_dir / f"{prefix}.meta.json"
    meta_file.write_text(json.dumps(meta, indent=2))

    if tests:
        test_file = cache_dir / f"{prefix}_test.py"
        test_file.write_text(tests)


# ---------------------------------------------------------------------------
# Synthesis
# ---------------------------------------------------------------------------

def _build_prompt(func_name: str, signature: str, docstring: str, 
                  type_hints: dict) -> str:
    """Build the LLM prompt for synthesizing an implementation."""
    hints_str = ""
    if type_hints:
        formatted = []
        for param, hint in type_hints.items():
            if param == "return":
                formatted.append(f"    returns: {hint}")
            else:
                formatted.append(f"    {param}: {hint}")
        hints_str = "\nType annotations:\n" + "\n".join(formatted)

    return f"""Implement this Python function. Return ONLY the complete function 
definition including the def line and body. No explanation, no markdown fences, 
no imports (assume all necessary imports are available in the calling scope).

def {func_name}{signature}:
    \"\"\"{docstring}\"\"\"
    # implement this
{hints_str}

Requirements:
- The implementation must satisfy everything described in the docstring.
- Keep it correct and simple, not clever.
- Handle edge cases mentioned or implied in the docstring.
- Use only the Python standard library unless the docstring specifies otherwise."""


def _synthesize(func_name: str, signature: str, docstring: str,
                type_hints: dict, model: str) -> str:
    """Call the LLM to synthesize an implementation."""
    client = anthropic.Anthropic()
    prompt = _build_prompt(func_name, signature, docstring, type_hints)

    logger.debug(f"speq: synthesizing '{func_name}' with model '{model}'")
    logger.debug(f"speq: prompt length: {len(prompt)} chars")

    start = time.time()
    response = client.messages.create(
        model=model,
        max_tokens=4096,
        messages=[{"role": "user", "content": prompt}],
    )
    elapsed = time.time() - start

    code = _clean_code(response.content[0].text)

    logger.debug(f"speq: synthesis of '{func_name}' completed in {elapsed:.1f}s")
    logger.debug(f"speq: generated {len(code)} chars of code")

    return code


def _clean_code(code: str) -> str:
    """Strip markdown fences and other LLM artifacts from generated code."""
    code = code.strip()
    if code.startswith("```python"):
        code = code[len("```python"):]
    elif code.startswith("```"):
        code = code[3:]
    if code.endswith("```"):
        code = code[:-3]
    return code.strip()


def _compile_function(code: str, func_name: str, 
                      source_globals: dict) -> typing.Callable:
    """Compile generated code and extract the function."""
    namespace = {}
    for key, value in source_globals.items():
        if isinstance(value, type) or inspect.ismodule(value) or key.startswith("__") is False:
            namespace[key] = value

    try:
        exec(code, namespace)
    except Exception as e:
        raise RuntimeError(
            f"speq: failed to compile synthesized code for '{func_name}'.\n"
            f"Error: {e}\n"
            f"Generated code:\n{code}"
        ) from e

    if func_name not in namespace:
        raise RuntimeError(
            f"speq: synthesized code does not define '{func_name}'.\n"
            f"Generated code:\n{code}"
        )

    return namespace[func_name]


# ---------------------------------------------------------------------------
# Confidence logging
# ---------------------------------------------------------------------------

def _log_scores(func_name: str, scores: ConfidenceScores) -> None:
    """Log confidence scores at debug level."""
    if scores.spec_clarity < 0:
        logger.debug(f"speq: [{func_name}] confidence scoring unavailable")
        return

    logger.debug(f"speq: [{func_name}] spec clarity:      {scores.spec_clarity:.2f}")
    logger.debug(f"speq: [{func_name}] impl confidence:    {scores.implementation_confidence:.2f}")

    if scores.assumptions:
        logger.debug(f"speq: [{func_name}] assumptions made:")
        for a in scores.assumptions:
            logger.debug(f"speq: [{func_name}]   - {a}")

    if scores.ambiguities:
        logger.debug(f"speq: [{func_name}] ambiguities found:")
        for a in scores.ambiguities:
            logger.debug(f"speq: [{func_name}]   - {a}")

    if scores.missing_edge_cases:
        logger.debug(f"speq: [{func_name}] missing edge cases:")
        for e in scores.missing_edge_cases:
            logger.debug(f"speq: [{func_name}]   - {e}")

    # surface low scores as warnings so they're visible even without debug
    if scores.spec_clarity < 0.5:
        logger.warning(
            f"speq: [{func_name}] low spec clarity ({scores.spec_clarity:.2f}) "
            f"— consider adding more detail to the docstring"
        )
    if scores.implementation_confidence < 0.5:
        logger.warning(
            f"speq: [{func_name}] low implementation confidence "
            f"({scores.implementation_confidence:.2f}) — review the generated code"
        )


# ---------------------------------------------------------------------------
# Body detection
# ---------------------------------------------------------------------------

def _body_is_stub(func: typing.Callable) -> bool:
    """Check if the function body is just `...` or `pass` (needs synthesis)."""
    try:
        source = inspect.getsource(func)
    except OSError:
        return False

    lines = source.split("\n")
    in_docstring = False
    docstring_char = None
    past_def = False
    in_def = False
    past_docstring = False
    found_docstring = False
    body_lines = []

    for line in lines:
        stripped = line.strip()

        # skip decorators
        if not past_def and not in_def and stripped.startswith("@"):
            continue

        # handle multi-line def: wait for the closing ):
        if not past_def:
            if stripped.startswith("def ") or in_def:
                in_def = True
                if ":" in stripped and stripped.endswith(":"):
                    past_def = True
                    in_def = False
            continue

        # handle docstring
        if not found_docstring and not in_docstring:
            if stripped.startswith('"""') or stripped.startswith("'''"):
                docstring_char = stripped[:3]
                in_docstring = True
                found_docstring = True
                if stripped.count(docstring_char) >= 2:
                    in_docstring = False
                    past_docstring = True
                continue

        if in_docstring:
            if docstring_char in stripped:
                in_docstring = False
                past_docstring = True
            continue

        # collect body lines after docstring
        if (past_docstring or not found_docstring) and stripped:
            body_lines.append(stripped)

    if not body_lines:
        return True

    return all(line in ("...", "pass") for line in body_lines)


# ---------------------------------------------------------------------------
# The decorator
# ---------------------------------------------------------------------------

_DEFAULT_MODEL = os.environ.get("SPEQ_MODEL", "claude-sonnet-4-20250514")


def ai(func: typing.Callable = None, *, model: str = _DEFAULT_MODEL):
    """
    Decorator that synthesizes a function implementation from its docstring.

    Usage:
        @ai
        def my_func(x: int) -> int:
            \"\"\"Double the input.\"\"\"
            ...

        @ai(model="claude-sonnet-4-20250514")
        def my_func(x: int) -> int:
            \"\"\"Double the input.\"\"\"
            ...

    Enable debug logging to see confidence scores:
        import logging
        logging.basicConfig(level=logging.DEBUG)
    """
    def decorator(fn: typing.Callable) -> typing.Callable:
        if not _body_is_stub(fn):
            return fn

        if not fn.__doc__:
            raise ValueError(
                f"speq: @ai function '{fn.__name__}' has no docstring. "
                f"Write a docstring describing what the function should do."
            )

        func_name = fn.__name__
        sig = str(inspect.signature(fn))
        docstring = fn.__doc__
        type_hints = typing.get_type_hints(fn)
        type_hints_str = {
            k: str(v) for k, v in type_hints.items()
        }

        try:
            source_file = inspect.getfile(fn)
        except (TypeError, OSError):
            source_file = None

        caller_globals = fn.__globals__ if hasattr(fn, "__globals__") else {}

        sid = _spec_hash(func_name, sig, docstring)

        # holder for scores — populated after first synthesis
        _scores_holder: dict[str, ConfidenceScores | None] = {"scores": None}

        @wraps(fn)
        def wrapper(*args, **kwargs):
            if sid in _memory_cache:
                return _memory_cache[sid](*args, **kwargs)

            cache_dir = None
            code = None
            scores = None

            if source_file:
                cache_dir = _get_cache_dir(source_file)
                code = _load_from_disk_cache(cache_dir, func_name, sid)
                if code is not None:
                    scores = _load_scores_from_disk_cache(cache_dir, func_name, sid)
                    logger.debug(f"speq: [{func_name}] loaded from disk cache")

            if code is None:
                logger.debug(f"speq: [{func_name}] synthesizing...")
                code = _synthesize(
                    func_name, sig, docstring, type_hints_str, model
                )

                # score the synthesis
                logger.debug(f"speq: [{func_name}] scoring confidence...")
                scores = _score_synthesis(
                    func_name, sig, docstring, code, model
                )

                # generate tests from the spec
                logger.debug(f"speq: [{func_name}] generating tests...")
                tests = _generate_tests(
                    func_name, sig, docstring, code, model
                )

                if cache_dir:
                    _save_to_disk_cache(
                        cache_dir, func_name, sid, code,
                        scores=scores, model=model, tests=tests
                    )

            # log scores
            if scores:
                _scores_holder["scores"] = scores
                _log_scores(func_name, scores)

            # compile and cache in memory
            impl = _compile_function(code, func_name, caller_globals)
            _memory_cache[sid] = impl

            return impl(*args, **kwargs)

        # attach metadata for introspection
        wrapper._speq_spec = docstring
        wrapper._speq_spec_id = sid
        wrapper._speq_model = model
        wrapper._speq_synthesized = True

        def get_scores() -> ConfidenceScores | None:
            """Get confidence scores. Available after first call."""
            return _scores_holder["scores"]

        wrapper.get_scores = get_scores

        return wrapper

    if func is not None:
        return decorator(func)
    return decorator


# ---------------------------------------------------------------------------
# Public API for CLI / programmatic use
# ---------------------------------------------------------------------------

@dataclass
class SynthesisResult:
    """Result of synthesizing a function from its spec."""
    func_name: str
    code: str
    scores: ConfidenceScores | None
    tests: str
    spec_id: str
    model: str
    cached: bool  # True if loaded from cache, False if freshly synthesized


def synthesize_from_spec(
    func_name: str,
    signature: str,
    docstring: str,
    source_file: str,
    model: str | None = None,
    force: bool = False,
) -> SynthesisResult:
    """
    Synthesize a function from its spec. Used by CLI commands.
    
    If force=True, re-synthesizes even if cached.
    Returns a SynthesisResult with the code, scores, and tests.
    """
    model = model or _DEFAULT_MODEL
    sid = _spec_hash(func_name, signature, docstring)
    cache_dir = _get_cache_dir(source_file)
    
    # check cache unless forcing
    if not force:
        code = _load_from_disk_cache(cache_dir, func_name, sid)
        if code is not None:
            scores = _load_scores_from_disk_cache(cache_dir, func_name, sid)
            # load tests
            tests = ""
            test_file = cache_dir / f"{_cache_prefix(func_name, sid)}_test.py"
            if test_file.exists():
                tests = test_file.read_text()
            return SynthesisResult(
                func_name=func_name, code=code, scores=scores,
                tests=tests, spec_id=sid, model=model, cached=True,
            )
    
    # synthesize
    type_hints_str = {}  # not available without importing the module
    code = _synthesize(func_name, signature, docstring, type_hints_str, model)
    scores = _score_synthesis(func_name, signature, docstring, code, model)
    tests = _generate_tests(func_name, signature, docstring, code, model)
    
    _save_to_disk_cache(
        cache_dir, func_name, sid, code,
        scores=scores, model=model, tests=tests,
    )
    
    return SynthesisResult(
        func_name=func_name, code=code, scores=scores,
        tests=tests, spec_id=sid, model=model, cached=False,
    )
