import json
import os
import re
import time
from bs4 import BeautifulSoup, Comment
import requests

# ─────────────────────────────────────────────────────────────────────────────
# Model fallback chain  (primary → fallback1 → fallback2)
# ─────────────────────────────────────────────────────────────────────────────
_FREE_MODEL_CHAIN = [
    "google/gemma-4-31b-it:free",
    "qwen/qwen3-next-80b-a3b-instruct:free",
    "qwen/qwen3-coder:free",
    "meta-llama/llama-3.3-8b-instruct:free"
]

# Files larger than this get the old tag-by-tag approach to stay within context limits
_MAX_WHOLE_DOC_CHARS = 12_000


def _openrouter_headers():
    key = os.environ.get("OPENROUTER_API_KEY", "")
    return {
        "Authorization": f"Bearer {key}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://webcraft.com",
        "X-Title": "WebCraft AI Backend",
    }


def _call_openrouter(model, messages, timeout=90, label=""):
    """Make a single OpenRouter API call. Returns response or None."""
    try:
        return requests.post(
            url="https://openrouter.ai/api/v1/chat/completions",
            headers=_openrouter_headers(),
            data=json.dumps({"model": model, "messages": messages}),
            timeout=timeout,
        )
    except requests.exceptions.Timeout:
        print(f"  OpenRouter Timeout ({label or model})")
        return None
    except Exception as e:
        print(f"  OpenRouter Request Error ({label or model}): {e}")
def _call_huggingface(messages, timeout=90):
    """Make a single Hugging Face Inference API call as an ultimate fallback."""
    hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_TOKEN", "")
    if not hf_token:
        print("  HF_TOKEN not found for Hugging Face failover.")
        return None
        
    model = "microsoft/Phi-3.5-mini-instruct"
    print(f"  [Ultimate Failover] Trying Hugging Face Inference: {model}")
    
    try:
        from huggingface_hub import InferenceClient
        client = InferenceClient(api_key=hf_token, timeout=30)
        reply = client.chat_completion(model=model, messages=messages, max_tokens=2000)
        return reply.choices[0].message.content.strip()
    except Exception as e:
        print(f"  Hugging Face Request Error: {e}")
        return None

def _try_models(messages, label=""):
    """
    Try each model in the fallback chain.
    Returns the raw text content on success, or raises RuntimeError if all fail.
    """
    for global_attempt, model in enumerate(_FREE_MODEL_CHAIN):
        for retry in range(3):
            if retry > 0:
                print(f"  [Retry {retry}/3] Waiting 5s before retrying {model}...")
                time.sleep(5)
            elif global_attempt > 0 and retry == 0:
                print(f"  Moving to fallback {global_attempt+1}/{len(_FREE_MODEL_CHAIN)}: {model}")
                time.sleep(2)

            resp = _call_openrouter(model, messages, label=label or f"attempt {global_attempt+1}")
            if resp is None:
                continue

            if resp.status_code == 200:
                try:
                    return resp.json()["choices"][0]["message"]["content"].strip()
                except (KeyError, Exception) as e:
                    print(f"  Parse error ({model}): {e} — trying next")
                    continue
            else:
                tag = "Primary" if global_attempt == 0 else "Fallback"
                print(f"  OpenRouter {tag} Error ({model}): Code {resp.status_code}")
                try:
                    print(f"  Reason: {resp.json().get('error', {}).get('message', resp.text)}")
                except:
                    print(f"  Reason: {resp.text}")

    # ==========================
    # ULTIMATE FAILOVER: Hugging Face
    # ==========================
    for retry in range(3):
        if retry > 0:
            print(f"  [Retry {retry}/3] Waiting 5s before retrying Hugging Face...")
            time.sleep(5)

        hf_resp_text = _call_huggingface(messages)
        if hf_resp_text:
            return hf_resp_text

    raise RuntimeError("AI_EXHAUSTED")


# ─────────────────────────────────────────────────────────────────────────────
# STRATEGY A  — Whole-document annotation  (BEST quality, used by default)
# The AI sees the full HTML at once → context-aware, page-specific comments.
# ─────────────────────────────────────────────────────────────────────────────
_WHOLE_DOC_SYSTEM = """\
You are an expert Web Development Instructor. Your job is to add educational \
HTML comments to the code a student wrote.

RULES
1. Return the COMPLETE HTML file with your comments inserted — nothing else.
2. Add a structured comment block DIRECTLY ABOVE every meaningful structural \
element: <header>, <nav>, <main>, <section>, <article>, <aside>, <footer>, \
<h1>–<h3>, <form>, and important <div> blocks (ones with an id or a clear role).
3. Skip trivial wrappers, <span>, <br>, <link>, <meta>, <script src="...">.
4. Use THIS exact comment format (4-space indent inside the comment):

<!--
    [ELEMENT]
    WHAT  : What this specific element IS on THIS page (mention the actual content).
    STYLE : What its classes/attributes achieve visually (be specific about Tailwind/CSS).
    WHY   : The web-dev best practice or semantic reason behind this choice.
-->

5. Keep each field to ONE sentence. Be concrete and page-specific — never generic.
6. Do NOT alter the HTML itself in any way. Do NOT add markdown fences around the output.
"""


def annotate_whole_document(html_content):
    """
    Best-quality approach: send the full HTML to the AI and get it back pre-commented.
    Returns annotated HTML string, or None if all models fail.
    """
    messages = [
        {"role": "system", "content": _WHOLE_DOC_SYSTEM},
        {"role": "user", "content": html_content},
    ]
    result = _try_models(messages, label="whole-doc")
    if result is None:
        return None

    # Strip accidental markdown fences the model may have added
    result = re.sub(r"^```(?:html)?\s*", "", result, flags=re.IGNORECASE)
    result = re.sub(r"\s*```$", "", result)
    return result.strip()


# ─────────────────────────────────────────────────────────────────────────────
# STRATEGY B  — Tag-by-tag bulk JSON  (fallback for large files)
# ─────────────────────────────────────────────────────────────────────────────
_BULK_SYSTEM = """\
You are an expert Web Development Instructor creating structured HTML documentation.
I will give you a JSON object where each key maps to a code snippet.
For EACH key produce a JSON object with exactly three fields:
  "what":  One sentence — what this element IS and its role on this specific page.
  "style": One sentence — what its classes/attributes achieve visually. \
If there are none, write "No special attributes."
  "why":   One sentence — the best-practice or semantic reason for this pattern.
Output ONE valid JSON object:
{ "key": { "what": "...", "style": "...", "why": "..." }, ... }
No markdown, no preamble, no text outside the JSON.
"""


def _bulk_ai_comments(snippets_dict, context="HTML"):
    """Tag-by-tag bulk call — used when the document is too large for whole-doc mode."""
    if not snippets_dict:
        return {}

    messages = [
        {"role": "system", "content": _BULK_SYSTEM.replace("HTML", context)},
        {"role": "user", "content": json.dumps(snippets_dict)},
    ]

    raw = _try_models(messages, label="bulk-tags")
    if raw is None:
        print("  All models exhausted — using placeholder comments.")
        return {k: {"what": "HTML element.", "style": "", "why": "Standard semantic element."}
                for k in snippets_dict}

    # Strip markdown fences
    if raw.startswith("```json"): raw = raw[7:]
    if raw.startswith("```"):     raw = raw[3:]
    if raw.endswith("```"):       raw = raw[:-3]

    try:
        parsed = json.loads(raw.strip())
        # Normalise flat-string responses
        result = {}
        for k, v in parsed.items():
            if isinstance(v, dict) and "what" in v:
                result[k] = v
            else:
                result[k] = {"what": str(v), "style": "", "why": ""}
        return result
    except json.JSONDecodeError as e:
        print(f"  JSON parse error in bulk response: {e}")
        # Re-raise so it hits the 503 instead of placeholders
        raise RuntimeError("AI_EXHAUSTED")


def _inject_tag_comments(soup, comment_level):
    """Tag-by-tag injection used for large files (Strategy B)."""
    for comment in soup.find_all(string=lambda t: isinstance(t, Comment)):
        comment.extract()

    target = ["header", "nav", "main", "section", "article", "aside", "footer",
              "h1", "h2", "h3", "form", "img", "div", "ul"]
    html_tags = soup.find_all(target)

    ai_queue = {}
    for i, tag in enumerate(html_tags):
        if comment_level != "concise":
            parts = [f"<{tag.name}"]
            for attr, val in tag.attrs.items():
                if isinstance(val, list): val = " ".join(val)
                parts.append(f' {attr}="{val}"')
            parts.append(">")
            ai_queue[f"tag_{i}"] = "".join(parts)

    ai_results = _bulk_ai_comments(ai_queue) if ai_queue else {}

    for i, tag in enumerate(html_tags):
        name = tag.name.upper()
        key  = f"tag_{i}"

        if comment_level == "concise":
            block = f" [{name}] "
        else:
            d = ai_results.get(key, {})
            what  = d.get("what", "").strip()  if isinstance(d, dict) else str(d)
            style = d.get("style", "").strip() if isinstance(d, dict) else ""
            why   = d.get("why", "").strip()   if isinstance(d, dict) else ""

            lines = [f"\n    [{name}]"]
            if what:  lines.append(f"    WHAT  : {what}")
            if style and style.lower() not in ("no special attributes.", "none", ""):
                lines.append(f"    STYLE : {style}")
            if why:   lines.append(f"    WHY   : {why}")
            lines.append("")
            block = "\n".join(lines)

        tag.insert_before(Comment(block))
        tag.insert_before(soup.new_string("\n"))


# ─────────────────────────────────────────────────────────────────────────────
# CSS  — rule-based (fast, no API call needed)
# ─────────────────────────────────────────────────────────────────────────────
def _annotate_css(soup, comment_level):
    for style_tag in soup.find_all("style"):
        css = style_tag.string or ""
        if not css.strip():
            continue

        if comment_level == "educational":
            header = (
                "/* ════════════════════════════════════════\n"
                "   CSS STYLES — Cascading Style Sheets\n"
                "   Styles defined here cascade down to all\n"
                "   matching child elements on the page.\n"
                "   ════════════════════════════════════════ */\n\n"
            )
        elif comment_level == "detailed":
            header = (
                "/* ════════════════════════════════════════\n"
                "   WebCraft — Generated Styles\n"
                "   ════════════════════════════════════════ */\n\n"
            )
        else:
            header = "/* Styles */\n\n"

        rules   = re.split(r"([\w\s.,#:\-\[\]=\"'()]+\{[^}]*\})", css)
        output  = header

        for rule in rules:
            rule = rule.strip()
            if not rule:
                continue
            if "{" in rule and "}" in rule:
                m = re.match(r"([\w\s.,#:\-\[\]=\"'()]+)\{", rule)
                if m and comment_level != "concise":
                    sel = m.group(1).strip()
                    if comment_level == "educational":
                        if sel.startswith("."):
                            output += f"/* CLASS  '{sel}' — reusable, applies to any element with this class */\n"
                        elif sel.startswith("#"):
                            output += f"/* ID     '{sel}' — unique, targets one specific element */\n"
                        elif sel == "*":
                            output += f"/* UNIVERSAL SELECTOR — resets default browser styles on every element */\n"
                        else:
                            output += f"/* ELEMENT '{sel}' */\n"
                    else:
                        output += f"/* {sel} */\n"
                output += rule + "\n\n"
            else:
                output += rule
        style_tag.string = output


# ─────────────────────────────────────────────────────────────────────────────
# JS  — rule-based per-function comments
# ─────────────────────────────────────────────────────────────────────────────
def _annotate_js(soup, comment_level):
    for script in soup.find_all("script"):
        if script.get("src"):
            continue
        js = script.string or ""
        if not js.strip():
            continue

        if comment_level == "educational":
            header = "// ═══ JAVASCRIPT — Interactive Logic ═══\n// Each function below handles a user interaction or page behaviour.\n\n"
        elif comment_level == "detailed":
            header = "// ═══ WebCraft — Generated JS ═══\n\n"
        else:
            header = "// Scripts\n\n"

        lines = js.split("\n")

        # Collect functions for bulk AI (keeps the one API call per document rule)
        ai_queue = {}
        for i, line in enumerate(lines):
            stripped = line.strip()
            if stripped.startswith("function "):
                m = re.search(r"function\s+(\w+)", stripped)
                if m and comment_level != "concise":
                    body = " ".join(l.strip() for l in lines[i : i + 6])
                    ai_queue[f"func_{i}"] = body

        ai_results = _bulk_ai_comments(ai_queue, context="JavaScript") if ai_queue else {}

        output = header
        for i, line in enumerate(lines):
            key = f"func_{i}"
            if key in ai_results:
                d = ai_results[key]
                what = d.get("what", "") if isinstance(d, dict) else str(d)
                why  = d.get("why",  "") if isinstance(d, dict) else ""
                output += f"// WHAT : {what}\n"
                if why:
                    output += f"// WHY  : {why}\n"
            output += line + "\n"

        script.string = output


# ─────────────────────────────────────────────────────────────────────────────
# PUBLIC API
# ─────────────────────────────────────────────────────────────────────────────
def add_comments_to_webcraft_file(
    file_path,
    comment_level="concise",
    output_path=None,
    openai_api_key=None,   # kept for signature compatibility
    force_title=None,
    benchmark_mode=False,
):
    """
    Add AI-powered educational comments to a WebCraft HTML file.

    Strategy selection:
      • concise   — fast rule-based labels only (no API call)
      • detailed  — whole-document AI annotation if file ≤ 12 000 chars,
                    otherwise tag-by-tag bulk JSON
      • educational — same as detailed, with richer CSS/JS headers

    Args:
        file_path     : Path to the HTML file
        comment_level : 'concise' | 'detailed' | 'educational'
        output_path   : Write output here (None = overwrite original)
        force_title   : If set, updates the <title> tag to this string

    Returns:
        Path to the commented file
    """
    print(f"Adding '{comment_level}' comments to WebCraft file...")

    # ── Read ──────────────────────────────────────────────────────────────────
    html_content = None
    for enc in ["utf-8", "latin-1", "cp1252", "iso-8859-1"]:
        try:
            with open(file_path, "r", encoding=enc) as f:
                html_content = f.read()
            break
        except Exception:
            continue
    if not html_content:
        with open(file_path, "rb") as f:
            html_content = f.read().decode("utf-8", errors="ignore")

    # ── HTML annotation ───────────────────────────────────────────────────────
    if comment_level == "concise":
        # Rule-based only — no API needed
        soup = BeautifulSoup(html_content, "html.parser")
        _inject_tag_comments(soup, comment_level)
        _annotate_css(soup, comment_level)
        _annotate_js(soup, comment_level)
        annotated_html = str(soup)

    else:
        # detailed / educational  →  try whole-doc AI first
        use_whole_doc = len(html_content) <= _MAX_WHOLE_DOC_CHARS

        try:
            if use_whole_doc:
                print("  Strategy: whole-document AI annotation (best quality)...")
                annotated_html = annotate_whole_document(html_content)

            if not use_whole_doc:
                print("  Strategy: tag-by-tag bulk JSON annotation...")
                soup = BeautifulSoup(html_content, "html.parser")
                _inject_tag_comments(soup, comment_level)
                annotated_html = str(soup)
        except RuntimeError as e:
            if str(e) == "AI_EXHAUSTED":
                if benchmark_mode:
                    raise e
                else:
                    print("  Production Mode: All AI exhausted. Rendering placeholder block.")
                    soup = BeautifulSoup(html_content, "html.parser")
                    msg = Comment(" AI comments temporarily unavailable. \n       Please try again in a few minutes. ")
                    body = soup.find("body")
                    if body:
                        body.insert(0, msg)
                    else:
                        soup.insert(0, msg)
                    annotated_html = str(soup)
            else:
                raise e

        # For CSS/JS we always use the fast rule-based pass
        # (parse the already-annotated HTML so we don't lose the HTML comments)
        soup = BeautifulSoup(annotated_html, "html.parser")
        _annotate_css(soup, comment_level)
        _annotate_js(soup, comment_level)
        annotated_html = str(soup)

    # ── Force title ───────────────────────────────────────────────────────────
    if force_title:
        soup = BeautifulSoup(annotated_html, "html.parser")
        title_tag = soup.find("title")
        if title_tag:
            title_tag.string = force_title
            print(f"  Updated <title> to: {force_title}")
        else:
            head = soup.find("head")
            if head:
                new_title = soup.new_tag("title")
                new_title.string = force_title
                head.insert(0, new_title)
                print(f"  Created <title> tag: {force_title}")
        annotated_html = str(soup)

    # ── Write ─────────────────────────────────────────────────────────────────
    if output_path is None:
        output_path = file_path

    with open(output_path, "w", encoding="utf-8") as f:
        f.write(annotated_html)

    print(f"Comments added successfully to {output_path}")
    return output_path