from __future__ import annotations

import difflib
import json
import os
import re
import urllib.error
import urllib.request
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from .base import Patcher, PatchResult


def _normalize_base_url(base_url: str) -> str:
    base = base_url.rstrip("/")
    # Already has a versioned path (e.g. /v1, /v4) — don't append /v1.
    if base.endswith(("/v1", "/v2", "/v3", "/v4")):
        return base
    return base + "/v1"


def _extract_unified_diff(text: str) -> str | None:
    s = _strip_think(text)
    s = s.strip("\n")
    if not s.strip():
        return None

    # Prefer fenced blocks if present.
    fence = re.search(r"```(?:diff|patch)?\n(.*?)\n```", s, re.DOTALL | re.IGNORECASE)
    if fence:
        candidate = fence.group(1).strip("\n")
        extracted = _extract_unified_diff(candidate)
        if extracted:
            return extracted

    lines = s.splitlines()

    # If it already looks like a git diff, keep from the first header.
    git = _extract_git_diff_blocks(lines)
    if git is not None:
        return git

    # Otherwise try to normalize "unified diff" blocks into git-apply format.
    normalized = _normalize_unified_blocks(lines)
    return normalized


def _strip_think(text: str) -> str:
    # Some reasoning models emit <think>...</think>; strip it.
    return re.sub(r"<think>.*?</think>\n?", "", text, flags=re.DOTALL | re.IGNORECASE)


def _normalize_unified_blocks(lines: list[str]) -> str | None:
    out_lines: list[str] = []
    i = 0
    n = len(lines)
    while i < n:
        line = lines[i]
        if not line.startswith("--- "):
            i += 1
            continue
        if i + 1 >= n or not lines[i + 1].startswith("+++ "):
            i += 1
            continue

        old_raw = lines[i][4:].strip()
        new_raw = lines[i + 1][4:].strip()
        old_path = _clean_path_token(old_raw)
        new_path = _clean_path_token(new_raw)
        if old_path is None or new_path is None:
            i += 2
            continue

        # Avoid accidental renames in model output (common mistake).
        if old_path != "/dev/null" and new_path != old_path:
            new_path = old_path

        out_lines.append(f"diff --git a/{old_path} b/{new_path}")
        out_lines.append(f"--- a/{old_path}")
        out_lines.append(f"+++ b/{new_path}")

        # Copy hunks until next file header, but stop on non-diff junk.
        i += 2
        seen_hunk = False
        while i < n:
            nxt = lines[i]
            if nxt.startswith("--- ") and (i + 1) < n and lines[i + 1].startswith("+++ "):
                break
            if nxt.startswith("@@"):
                seen_hunk = True
                out_lines.append(nxt)
                i += 1
                continue
            if seen_hunk and (
                nxt.startswith(" ")
                or nxt.startswith("+")
                or nxt.startswith("-")
                or nxt.startswith("\\")
            ):
                out_lines.append(nxt)
                i += 1
                continue
            # End of diff content (model started explaining, etc).
            break
        # Ensure a blank line between file diffs for readability.
        if out_lines and out_lines[-1] != "":
            out_lines.append("")

    if not out_lines:
        return None

    text = "\n".join(out_lines).rstrip("\n") + "\n"
    if "diff --git " not in text:
        return None
    return text


def _clean_path_token(raw: str) -> str | None:
    if not raw:
        return None
    # Common weird outputs:
    # - "--- buggy.py ---" (extra trailing token)
    # - "--- a/foo.py\t2026-..." (timestamps)
    token = raw.split()[0]
    if token in {"a/", "b/"}:
        return None
    token = token.removeprefix("a/")
    token = token.removeprefix("b/")
    if token == "/dev/null":
        return token
    return token


def _extract_git_diff_blocks(lines: list[str]) -> str | None:
    starts = [i for i, line in enumerate(lines) if line.startswith("diff --git ")]
    if not starts:
        return None

    captured: list[str] = []
    i = starts[0]
    n = len(lines)

    # Capture only "diff-like" lines; stop when model starts narrating.
    while i < n:
        line = lines[i]

        if line.startswith("diff --git "):
            captured.append(line)
            i += 1
            continue

        allowed = (
            line.startswith("index ")
            or line.startswith("new file mode ")
            or line.startswith("deleted file mode ")
            or line.startswith("old mode ")
            or line.startswith("new mode ")
            or line.startswith("similarity index ")
            or line.startswith("rename from ")
            or line.startswith("rename to ")
            or line.startswith("copy from ")
            or line.startswith("copy to ")
            or line.startswith("--- a/")
            or line == "--- /dev/null"
            or line.startswith("+++ b/")
            or line == "+++ /dev/null"
            or line.startswith("@@")
            or line.startswith(" ")
            or line.startswith("+")
            or line.startswith("-")
            or line.startswith("\\")
            or line == ""
        )
        if not allowed:
            break

        captured.append(line)
        i += 1

    text = "\n".join(captured).strip("\n") + "\n"
    if "diff --git " not in text:
        return None
    if "--- a/" not in text and "--- /dev/null" not in text:
        return None
    if "+++ b/" not in text and "+++ /dev/null" not in text:
        return None
    return text


@dataclass(frozen=True)
class OpenAICompatPatcher(Patcher):
    """OpenAI API-compatible chat-completions patcher.

    Tested against local OpenAI-compatible servers (vLLM/LM Studio/etc).
    """

    base_url: str
    model: str
    api_key: str | None = None
    timeout_sec: float = 120.0
    temperature: float = 0.0
    max_tokens: int = 2048
    response_mode: str = "edits"  # "edits" | "diff"
    force_patch_attempt: bool = False
    seed: int | None = None
    name: str = "openai"

    def propose_patch(self, *, prompt: str, context_dir: Path, cwd: Path) -> PatchResult:
        base_url = _normalize_base_url(self.base_url)
        url = f"{base_url}/chat/completions"

        api_key = self.api_key
        if api_key is None:
            api_key = (
                os.getenv("OPENAI_API_KEY") or os.getenv("LLMDEBUG_EVAL_OPENAI_API_KEY") or "local"
            )

        if self.response_mode == "edits":
            result = _propose_patch_via_edits(
                url=url,
                api_key=api_key,
                model=self.model,
                prompt=prompt,
                cwd=cwd,
                temperature=self.temperature,
                max_tokens=self.max_tokens,
                timeout_sec=self.timeout_sec,
                seed=self.seed,
            )
        else:
            result = _propose_patch_via_diff(
                url=url,
                api_key=api_key,
                model=self.model,
                prompt=prompt,
                temperature=self.temperature,
                max_tokens=self.max_tokens,
                timeout_sec=self.timeout_sec,
                seed=self.seed,
            )

        result_meta = dict(result.meta or {})
        result_meta.setdefault("base_url", base_url)
        result_meta.setdefault("model", self.model)
        result_meta.setdefault("temperature", self.temperature)
        result_meta.setdefault("max_tokens", self.max_tokens)
        result_meta.setdefault("response_mode", self.response_mode)
        result_meta.setdefault("api_compat_seed_requested", self.seed)
        result_meta.setdefault(
            "api_compat_seed_acknowledged", False if self.seed is not None else None
        )
        result_meta.setdefault("api_compat_seed_response", None)
        result = PatchResult(diff=result.diff, meta=result_meta)

        if result.diff is not None or not self.force_patch_attempt:
            return result

        forced = _build_forced_attempt_diff(prompt=prompt, cwd=cwd)
        if forced is None:
            return result

        meta = dict(result.meta or {})
        meta["forced_patch_attempt"] = True
        meta["forced_patch_target"] = forced[0]
        meta["forced_patch_reason"] = "no_model_diff"
        return PatchResult(diff=forced[1], meta=meta)


def _propose_patch_via_diff(
    *,
    url: str,
    api_key: str,
    model: str,
    prompt: str,
    temperature: float,
    max_tokens: int,
    timeout_sec: float,
    seed: int | None,
) -> PatchResult:
    base_messages = [
        {
            "role": "system",
            "content": (
                "You fix failing tests by returning a single unified diff. "
                "Return ONLY the diff; no explanations. Do not use Markdown fences."
            ),
        },
        {"role": "user", "content": prompt},
    ]

    # Try twice: second attempt primes the assistant to start with a diff header.
    attempts: list[list[dict[str, str]]] = [
        base_messages,
        [*base_messages, {"role": "assistant", "content": "diff --git a/"}],
    ]

    last_meta: dict[str, object] | None = None
    for idx, messages in enumerate(attempts, start=1):
        raw, meta = _chat_completion(
            url=url,
            api_key=api_key,
            model=model,
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens,
            timeout_sec=timeout_sec,
            seed=seed,
        )
        last_meta = {"attempt": idx, **meta}
        if raw is None:
            continue

        content = _get_first_choice_content(raw)
        if content is None:
            last_meta["error"] = "no_choice_content"
            last_meta["response"] = _safe_trunc(raw)
            continue

        diff = _extract_unified_diff(content)
        if diff is None:
            last_meta["error"] = "no_diff_found"
            last_meta["raw_response"] = content[:2000]
            continue

        fixed = _fix_hunk_counts(diff)
        return PatchResult(diff=fixed, meta=last_meta)

    return PatchResult(diff=None, meta=last_meta or {"error": "unknown"})


def _propose_patch_via_edits(
    *,
    url: str,
    api_key: str,
    model: str,
    prompt: str,
    cwd: Path,
    temperature: float,
    max_tokens: int,
    timeout_sec: float,
    seed: int | None,
) -> PatchResult:
    messages = [
        {
            "role": "system",
            "content": (
                "Return only JSON with this schema: "
                '{"edits":[{"path":"relative/path.py","find":"exact old text","replace":"new text","count":1}]}. '
                "No prose. No markdown. "
                "Use only exact existing paths from the prompt's available file list. "
                "Never use placeholder paths like relative/path.py."
            ),
        },
        {"role": "user", "content": prompt},
    ]

    raw, meta = _chat_completion(
        url=url,
        api_key=api_key,
        model=model,
        messages=messages,
        temperature=temperature,
        max_tokens=max_tokens,
        timeout_sec=timeout_sec,
        seed=seed,
        extra_payload={"response_format": {"type": "json_object"}},
    )
    if raw is None:
        # Some OpenAI-compatible backends do not support response_format.
        err = str(meta.get("error", ""))
        body = str(meta.get("body", ""))
        if err == "http_error" and "response_format" in body:
            raw, meta = _chat_completion(
                url=url,
                api_key=api_key,
                model=model,
                messages=messages,
                temperature=temperature,
                max_tokens=max_tokens,
                timeout_sec=timeout_sec,
                seed=seed,
            )
    if raw is None:
        return PatchResult(diff=None, meta=meta)

    content = _get_first_choice_content(raw)
    if content is None:
        return PatchResult(
            diff=None, meta={"error": "no_choice_content", "response": _safe_trunc(raw), **meta}
        )

    payload = _extract_json_object(content)
    if payload is None:
        diff_fallback = _propose_patch_via_diff(
            url=url,
            api_key=api_key,
            model=model,
            prompt=prompt,
            temperature=temperature,
            max_tokens=max_tokens,
            timeout_sec=timeout_sec,
            seed=seed,
        )
        if diff_fallback.diff is not None:
            fallback_meta = dict(diff_fallback.meta or {})
            fallback_meta["fallback_from"] = "edits"
            fallback_meta["fallback_reason"] = "no_json_payload"
            return PatchResult(diff=diff_fallback.diff, meta=fallback_meta)
        return PatchResult(
            diff=None,
            meta={
                "error": "no_json_payload",
                "raw_response": content[:2000],
                **meta,
            },
        )

    edits = payload.get("edits")
    if not isinstance(edits, list) or not edits:
        diff_fallback = _propose_patch_via_diff(
            url=url,
            api_key=api_key,
            model=model,
            prompt=prompt,
            temperature=temperature,
            max_tokens=max_tokens,
            timeout_sec=timeout_sec,
            seed=seed,
        )
        if diff_fallback.diff is not None:
            fallback_meta = dict(diff_fallback.meta or {})
            fallback_meta["fallback_from"] = "edits"
            fallback_meta["fallback_reason"] = "missing_edits"
            return PatchResult(diff=diff_fallback.diff, meta=fallback_meta)
        return PatchResult(
            diff=None, meta={"error": "missing_edits", "payload": _safe_trunc(payload), **meta}
        )

    changed: list[tuple[str, str, str]] = []
    for entry in edits:
        if not isinstance(entry, dict):
            return PatchResult(
                diff=None, meta={"error": "bad_edit_entry", "entry": _safe_trunc(entry), **meta}
            )

        path = entry.get("path")
        find = entry.get("find", entry.get("old"))
        replace = entry.get("replace", entry.get("new"))
        count_raw = entry.get("count", 1)
        count = count_raw if isinstance(count_raw, int) else 1

        if not isinstance(path, str) or not path.strip():
            return PatchResult(
                diff=None, meta={"error": "bad_path", "entry": _safe_trunc(entry), **meta}
            )
        if not isinstance(find, str) or not isinstance(replace, str):
            return PatchResult(
                diff=None, meta={"error": "bad_find_replace", "entry": _safe_trunc(entry), **meta}
            )
        if count < 0:
            count = 0

        file_path = (cwd / path).resolve()
        try:
            file_path.relative_to(cwd.resolve())
        except Exception:
            return PatchResult(diff=None, meta={"error": "path_outside_case", "path": path, **meta})
        if not file_path.exists():
            return PatchResult(diff=None, meta={"error": "file_not_found", "path": path, **meta})

        before = file_path.read_text(encoding="utf-8")
        if find not in before:
            return PatchResult(
                diff=None,
                meta={"error": "find_not_found", "path": path, "find_preview": find[:200], **meta},
            )

        after = before.replace(find, replace, count or -1)
        if after == before:
            continue
        changed.append((path, before, after))

    if not changed:
        return PatchResult(diff=None, meta={"error": "no_effective_changes", **meta})

    diff = _build_diff_from_changes(changed)
    return PatchResult(diff=diff, meta=meta)


def _chat_completion(
    *,
    url: str,
    api_key: str,
    model: str,
    messages: list[dict[str, str]],
    temperature: float,
    max_tokens: int,
    timeout_sec: float,
    seed: int | None,
    extra_payload: dict[str, Any] | None = None,
) -> tuple[dict[str, Any] | None, dict[str, object]]:
    payload: dict[str, Any] = {
        "model": model,
        "messages": messages,
        "temperature": temperature,
        "max_tokens": max_tokens,
    }
    if seed is not None:
        payload["seed"] = seed
    if extra_payload:
        payload.update(extra_payload)

    base_meta: dict[str, object] = {
        "model": model,
        "base_url": url.rsplit("/v1/", 1)[0] + "/v1",
        "temperature": temperature,
        "max_tokens": max_tokens,
        "api_compat_seed_requested": seed,
        "api_compat_seed_acknowledged": False if seed is not None else None,
        "api_compat_seed_response": None,
    }

    headers = {
        "Content-Type": "application/json",
        "Accept": "application/json",
        "Authorization": f"Bearer {api_key}",
    }

    req = urllib.request.Request(
        url=url,
        data=json.dumps(payload).encode("utf-8"),
        headers=headers,
        method="POST",
    )

    raw_text = None
    try:
        with urllib.request.urlopen(req, timeout=timeout_sec) as resp:
            raw_text = resp.read().decode("utf-8", errors="replace")
    except urllib.error.HTTPError as e:
        body = ""
        try:
            body = e.read().decode("utf-8", errors="replace")
        except Exception:
            body = ""
        return None, {**base_meta, "error": "http_error", "status": e.code, "body": body[:2000]}
    except Exception as e:
        return None, {**base_meta, "error": "request_failed", "exc": f"{type(e).__name__}: {e}"}

    try:
        data = json.loads(raw_text) if raw_text is not None else {}
    except Exception:
        return None, {**base_meta, "error": "bad_json_response", "body": (raw_text or "")[:2000]}

    meta: dict[str, object] = dict(base_meta)
    usage = data.get("usage")
    if isinstance(usage, dict):
        meta["usage"] = usage
    if seed is not None:
        response_seed = _extract_response_seed(data)
        meta["api_compat_seed_response"] = response_seed
        meta["api_compat_seed_acknowledged"] = response_seed == seed
    return data if isinstance(data, dict) else None, meta


def _get_first_choice_content(resp: dict[str, Any]) -> str | None:
    choices = resp.get("choices")
    if not isinstance(choices, list) or not choices:
        return None
    c0 = choices[0]
    if not isinstance(c0, dict):
        return None
    msg = c0.get("message")
    if isinstance(msg, dict):
        content = msg.get("content")
        return content if isinstance(content, str) else None
    # Some servers return `text` (completions-like).
    text = c0.get("text")
    return text if isinstance(text, str) else None


def _safe_trunc(x: object, limit: int = 2000) -> str:
    try:
        s = json.dumps(x, ensure_ascii=False, default=str)
    except Exception:
        s = repr(x)
    return s[:limit]


def _extract_response_seed(resp: dict[str, Any]) -> int | None:
    candidates: list[Any] = [
        resp.get("seed"),
        resp.get("response_seed"),
    ]
    metadata = resp.get("metadata")
    if isinstance(metadata, dict):
        candidates.append(metadata.get("seed"))
        candidates.append(metadata.get("response_seed"))
    for candidate in candidates:
        if isinstance(candidate, int):
            return candidate
        if isinstance(candidate, str):
            try:
                return int(candidate)
            except ValueError:
                continue
    return None


def _extract_json_object(text: str) -> dict[str, Any] | None:
    s = _strip_think(text).strip()
    if not s:
        return None

    fence = re.search(r"```(?:json)?\n(.*?)\n```", s, flags=re.DOTALL | re.IGNORECASE)
    if fence:
        s = fence.group(1).strip()

    try:
        data = json.loads(s)
        return data if isinstance(data, dict) else None
    except Exception:
        pass

    decoder = json.JSONDecoder()
    for i, ch in enumerate(s):
        if ch != "{":
            continue
        try:
            obj, _end = decoder.raw_decode(s[i:])
            return obj if isinstance(obj, dict) else None
        except Exception:
            continue
    return None


def _build_diff_from_changes(changes: list[tuple[str, str, str]]) -> str:
    parts: list[str] = []
    for path, before, after in changes:
        udiff = difflib.unified_diff(
            before.splitlines(),
            after.splitlines(),
            fromfile=f"a/{path}",
            tofile=f"b/{path}",
            lineterm="",
        )
        file_lines = list(udiff)
        if not file_lines:
            continue
        parts.append(f"diff --git a/{path} b/{path}")
        parts.extend(file_lines)
        if parts and parts[-1] != "":
            parts.append("")
    return "\n".join(parts).rstrip("\n") + "\n"


def _build_forced_attempt_diff(*, prompt: str, cwd: Path) -> tuple[str, str] | None:
    """Build a deterministic minimal patch when model output is empty/invalid."""
    candidates = _extract_file_index_paths(prompt)
    if not candidates:
        return None

    target_rel: str | None = None
    for rel in candidates:
        if not rel.endswith(".py"):
            continue
        lower = rel.lower()
        name = Path(rel).name.lower()
        if name.startswith("test_") or "/test" in lower or "\\test" in lower:
            continue
        target_rel = rel
        break
    if target_rel is None:
        for rel in candidates:
            if rel.endswith(".py"):
                target_rel = rel
                break
    if target_rel is None:
        return None

    target_path = (cwd / target_rel).resolve()
    try:
        target_path.relative_to(cwd.resolve())
    except Exception:
        return None
    if not target_path.exists():
        return None

    before = target_path.read_text(encoding="utf-8")
    marker = "# llmdebug-force-attempt"
    if marker in before:
        marker = "# llmdebug-force-attempt-2"
    sep = "" if before.endswith("\n") else "\n"
    after = before + f"{sep}{marker}\n"
    if after == before:
        return None

    diff = _build_diff_from_changes([(target_rel, before, after)])
    return target_rel, diff


def _extract_file_index_paths(prompt: str) -> list[str]:
    lines = prompt.splitlines()
    start = None
    for i, line in enumerate(lines):
        if line.strip().startswith("Available files"):
            start = i + 1
            break
    if start is None:
        return []

    out: list[str] = []
    for line in lines[start:]:
        s = line.strip()
        if not s:
            if out:
                break
            continue
        if s.endswith(":"):
            break
        if s.startswith("Pytest output") or s.startswith("Project files:"):
            break
        out.append(s)
    return out


_HUNK_RE = re.compile(r"^@@ -(\d+)(?:,(\d+))? \+(\d+)(?:,(\d+))? @@(.*)$")


def _fix_hunk_counts(diff: str) -> str:
    """Make unified diff hunks syntactically consistent.

    LLMs often emit correct hunk bodies but incorrect `@@ -l,s +l,s @@` counts.
    Git treats those diffs as corrupt. We recompute the counts from the hunk body
    and rewrite the header counts, keeping the start line numbers as-is.
    """
    lines = diff.splitlines()
    out: list[str] = []
    i = 0
    n = len(lines)
    while i < n:
        line = lines[i]
        m = _HUNK_RE.match(line)
        if not m:
            out.append(line)
            i += 1
            continue

        old_start = max(1, int(m.group(1)))
        new_start = max(1, int(m.group(3)))
        suffix = m.group(5) or ""

        # Collect hunk body.
        body: list[str] = []
        i += 1
        while i < n:
            nxt = lines[i]
            if nxt.startswith("diff --git ") or nxt.startswith("@@ "):
                break
            if nxt.startswith("--- a/") or nxt.startswith("+++ b/"):
                break
            body.append(nxt)
            i += 1

        old_count = 0
        new_count = 0
        for b in body:
            if not b:
                continue
            c = b[0]
            if c == " ":
                old_count += 1
                new_count += 1
            elif c == "-":
                old_count += 1
            elif c == "+":
                new_count += 1
            else:
                # "\ No newline..." or other metadata does not affect counts.
                pass

        out.append(f"@@ -{old_start},{old_count} +{new_start},{new_count} @@{suffix}")
        out.extend(body)

    return "\n".join(out).rstrip("\n") + "\n"
