import json
import re

import httpx

from app.core.config import get_settings

settings = get_settings()

_JSON_ARRAY_RE = re.compile(r"\[.*\]", re.DOTALL)

GENERATION_SYSTEM_PROMPT = """You are an expert instructional designer. Given a source document, \
you generate high-quality instruction-tuning examples for training a domain-specific AI \
assistant.

Produce a single JSON object of the form {"examples": [ ... ]}, where "examples" is an array
and each item has exactly these fields:
- "example_type": one of "qa", "conversation", "scenario", "reasoning", "diagnostic", "follow_up"
- "difficulty": one of "beginner", "intermediate", "advanced"
- "instruction": the question or task a user would ask, grounded ONLY in the source document
- "input": optional extra context for the instruction, or null
- "output": a complete, accurate answer based ONLY on the source document

Rules:
- Ground every example strictly in the provided document. Do not invent facts.
- Include at least one example where the correct answer is to express uncertainty or ask a
  clarifying question, if the document does not fully cover a plausible user question.
- Vary phrasing between a complete beginner and an experienced practitioner.
- Return ONLY the JSON object described above, no prose, no markdown fences.
"""


def is_reachable() -> bool:
    try:
        resp = httpx.get(f"{settings.ollama_host}/api/tags", timeout=3)
        return resp.status_code == 200
    except httpx.HTTPError:
        return False


def list_models() -> list[dict]:
    resp = httpx.get(f"{settings.ollama_host}/api/tags", timeout=10)
    resp.raise_for_status()
    return resp.json().get("models", [])


def delete_model(tag: str) -> None:
    resp = httpx.request("DELETE", f"{settings.ollama_host}/api/delete", json={"model": tag}, timeout=30)
    resp.raise_for_status()


def pull_model_stream(tag: str):
    """Yields Ollama's NDJSON pull-progress dicts as they arrive."""
    with httpx.stream(
        "POST", f"{settings.ollama_host}/api/pull", json={"model": tag, "stream": True}, timeout=None
    ) as resp:
        resp.raise_for_status()
        for line in resp.iter_lines():
            if not line:
                continue
            yield json.loads(line)


def create_model(tag: str, modelfile: str) -> None:
    resp = httpx.post(
        f"{settings.ollama_host}/api/create", json={"model": tag, "modelfile": modelfile, "stream": False}, timeout=300
    )
    resp.raise_for_status()


def chat(
    model: str,
    messages: list[dict],
    json_mode: bool = False,
    temperature: float = 0.3,
    timeout: float = 120,
) -> dict:
    """Thin wrapper over /api/chat returning content plus Ollama's own latency/token
    counters (total_duration is nanoseconds; eval_count/prompt_eval_count are tokens)."""
    payload: dict = {
        "model": model,
        "messages": messages,
        "stream": False,
        "options": {"temperature": temperature, "num_ctx": settings.ollama_num_ctx},
    }
    if json_mode:
        payload["format"] = "json"

    resp = httpx.post(f"{settings.ollama_host}/api/chat", json=payload, timeout=timeout)
    resp.raise_for_status()
    data = resp.json()
    return {
        "content": data.get("message", {}).get("content", ""),
        "total_duration_ns": data.get("total_duration", 0),
        "prompt_eval_count": data.get("prompt_eval_count", 0),
        "eval_count": data.get("eval_count", 0),
        "done_reason": data.get("done_reason"),
    }


def _extract_json_array(text: str) -> list[dict]:
    text = text.strip()
    if text.startswith("```"):
        text = text.strip("`")
        text = text.split("\n", 1)[-1] if "\n" in text else text

    try:
        parsed = json.loads(text)
    except json.JSONDecodeError:
        match = _JSON_ARRAY_RE.search(text)
        if not match:
            raise
        parsed = json.loads(match.group(0))

    if isinstance(parsed, list):
        return parsed

    if isinstance(parsed, dict):
        if isinstance(parsed.get("examples"), list):
            return parsed["examples"]
        if "instruction" in parsed and "output" in parsed:
            return [parsed]
        for value in parsed.values():
            if isinstance(value, list):
                return value

    return []


EXPANSION_SYSTEM_PROMPT = """You are an expert instructional designer. Given one already-approved \
instruction-tuning example, you rewrite it into new variants \
along specific dimensions requested by the caller, while keeping every fact grounded in the \
original example (and source document excerpt, if provided). Do not invent new facts.

Produce a single JSON object of the form {"variants": [ ... ]}, where "variants" is an array with
exactly one item per requested variant type, each item having exactly these fields:
- "variant_type": echoed back exactly as requested
- "instruction": the reworded question or task
- "input": optional extra context for the instruction, or null
- "output": a complete, accurate answer consistent with the original example

Return ONLY the JSON object described above, no prose, no markdown fences.
"""

VARIANT_META: dict[str, dict[str, str]] = {
    "beginner": {
        "difficulty": "beginner",
        "example_type": "qa",
        "guidance": "Rewrite for a complete beginner with no technical background — simple "
        "language, concrete steps, no jargon.",
    },
    "advanced": {
        "difficulty": "advanced",
        "example_type": "qa",
        "guidance": "Rewrite for an experienced practitioner — precise technical terminology and "
        "quantitative detail.",
    },
    "edge_case": {
        "difficulty": "advanced",
        "example_type": "scenario",
        "guidance": "Turn it into an edge-case scenario — unusual conditions, conflicting "
        "constraints, or a rare combination of factors.",
    },
    "troubleshooting": {
        "difficulty": "intermediate",
        "example_type": "diagnostic",
        "guidance": "Reframe as troubleshooting — something has gone wrong, diagnose the likely "
        "cause and recommend a fix.",
    },
    "comparative_reasoning": {
        "difficulty": "advanced",
        "example_type": "reasoning",
        "guidance": "Reframe as a comparison between two or more approaches or options, requiring "
        "reasoning about trade-offs.",
    },
}


def expand_instruction_variants(
    instruction: str,
    output: str,
    input_text: str | None,
    variant_types: list[str],
    grounding_text: str | None = None,
    model: str | None = None,
) -> list[dict]:
    model_name = model or settings.ollama_default_model
    requested = [v for v in variant_types if v in VARIANT_META]
    if not requested:
        return []

    guidance_lines = "\n".join(f'- "{v}": {VARIANT_META[v]["guidance"]}' for v in requested)
    grounding_block = (
        f"\n\nSource document excerpt for grounding facts:\n---\n{grounding_text[:6000]}\n---"
        if grounding_text
        else ""
    )

    user_prompt = (
        f"Original approved example:\n"
        f"Instruction: {instruction}\n"
        f"Input: {input_text or '(none)'}\n"
        f"Output: {output}"
        f"{grounding_block}\n\n"
        f'Produce one variant for EACH of the following variant types, as {{"variants": [...]}}:\n'
        f"{guidance_lines}\n"
        f'Each "variant_type" must be exactly one of: {requested}.'
    )

    resp = httpx.post(
        f"{settings.ollama_host}/api/chat",
        json={
            "model": model_name,
            "messages": [
                {"role": "system", "content": EXPANSION_SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt},
            ],
            "stream": False,
            "format": "json",
            "options": {"temperature": 0.8, "num_ctx": settings.ollama_num_ctx},
        },
        timeout=180,
    )
    resp.raise_for_status()
    content = resp.json()["message"]["content"]

    parsed = _extract_json_array(content)

    variants = []
    for item in parsed:
        if not isinstance(item, dict) or "instruction" not in item or "output" not in item:
            continue
        variant_type = item.get("variant_type")
        if variant_type not in VARIANT_META:
            continue
        meta = VARIANT_META[variant_type]
        input_value = item.get("input")
        if isinstance(input_value, (dict, list)):
            input_value = json.dumps(input_value, ensure_ascii=False)
        elif input_value is not None:
            input_value = str(input_value)
        variants.append(
            {
                "variant_type": variant_type,
                "difficulty": meta["difficulty"],
                "example_type": meta["example_type"],
                "instruction": str(item["instruction"]),
                "input": input_value,
                "output": str(item["output"]),
            }
        )
    return variants


def generate_instruction_examples(
    document_text: str, count: int, difficulty_mix: list[str], model: str | None = None
) -> list[dict]:
    model_name = model or settings.ollama_default_model
    excerpt = document_text[:12000]

    user_prompt = (
        f"Source document:\n---\n{excerpt}\n---\n\n"
        f'Generate {count} instruction-tuning examples as {{"examples": [...]}}. '
        f"Distribute difficulty across: {', '.join(difficulty_mix)}."
    )

    resp = httpx.post(
        f"{settings.ollama_host}/api/chat",
        json={
            "model": model_name,
            "messages": [
                {"role": "system", "content": GENERATION_SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt},
            ],
            "stream": False,
            "format": "json",
            "options": {"temperature": 0.7, "num_ctx": settings.ollama_num_ctx},
        },
        timeout=180,
    )
    resp.raise_for_status()
    content = resp.json()["message"]["content"]

    parsed = _extract_json_array(content)

    examples = []
    for item in parsed:
        if not isinstance(item, dict) or "instruction" not in item or "output" not in item:
            continue
        input_value = item.get("input")
        if isinstance(input_value, (dict, list)):
            input_value = json.dumps(input_value, ensure_ascii=False)
        elif input_value is not None:
            input_value = str(input_value)
        examples.append(
            {
                "example_type": item.get("example_type", "qa"),
                "difficulty": item.get("difficulty", "intermediate"),
                "instruction": str(item["instruction"]),
                "input": input_value,
                "output": str(item["output"]),
            }
        )
    return examples