import json
import logging
import re
from typing import Dict, Optional, Tuple

import torch
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForCausalLM

# =========================================================
# CONFIG
# =========================================================
MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
MAX_NEW_TOKENS = 512
RETRY_ATTEMPTS = 2  # first deterministic, then sampled attempt if validation fails

app = FastAPI(title="Japanese Corrector API (Qwen Edition)")
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("japanese-corrector")

# =========================================================
# LOAD MODEL (CPU)
# =========================================================
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
# Ensure pad token exists
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float32,
    device_map="cpu"
)
model.eval()

class InferenceRequest(BaseModel):
    input: str

# =========================================================
# SYSTEM PROMPT (enhanced)
# =========================================================
SYSTEM_PROMPT = """
You are a certified Japanese language instructor (JLPT-focused).

Your task is to correct the learner's Japanese sentence into natural, grammatically correct Japanese,
while strictly preserving the original meaning and polarity.

You must follow JLPT N5–N3 pedagogy.

CRITICAL CONSISTENCY RULE:
- You may ONLY explain particle changes that actually occur in the final corrected sentence.
- If a particle is not changed, explicitly state that it remains unchanged and explain why it is correct as-is.
- NEVER assert a substitution (e.g., 'は → が') unless that substitution appears in the final 'kanji' output.

PARTICLE EXPLANATION REQUIREMENTS:
- When correcting particles (は・が・を・に), explicitly explain the grammatical reason using JLPT-friendly terminology.
- Use concise "Particle Logic: ... | Grammar: ..." format in the 'explanation' field.

OUTPUT FORMAT (STRICT):
Return ONLY valid JSON with this schema:
{
  "kanji": "...",
  "hiragana": "...",
  "explanation": "Particle Logic: ... | Grammar: ..."
}

IMPORTANT RULES:
- Do NOT change meaning or polarity.
- Do NOT invent particles or verb forms.
- Explanations must be precise, neutral, and instructional (teacher tone).
"""

# =========================================================
# JSON EXTRACTION (robust)
# =========================================================
def find_balanced_json(text: str) -> Optional[str]:
    """Find the first balanced { ... } JSON object in text (handles nested braces)."""
    start = None
    depth = 0
    for i, ch in enumerate(text):
        if ch == '{':
            if start is None:
                start = i
            depth += 1
        elif ch == '}':
            if depth > 0:
                depth -= 1
                if depth == 0 and start is not None:
                    return text[start:i + 1]
    return None

def extract_json(text: str) -> Optional[Dict]:
    """Clean and parse JSON from model output robustly."""
    # remove common markdown fences first
    text = re.sub(r"```(?:json)?\s*", "", text)
    text = re.sub(r"\s*```", "", text)
    candidate = find_balanced_json(text)
    if not candidate:
        return None
    try:
        return json.loads(candidate)
    except Exception as e:
        logger.error("JSON parse failed for candidate: %s; error: %s", candidate, e)
        return None

# =========================================================
# VALIDATION
# =========================================================
SUBSTITUTION_PATTERNS = [
    re.compile(r"は\s*(?:→|->)\s*が"),
    re.compile(r"が\s*(?:→|->)\s*を"),
    re.compile(r"を\s*(?:→|->)\s*が"),
    # optionally add Japanese wording patterns if you commonly see them
    re.compile(r"は\s*→\s*が"),
    re.compile(r"が\s*→\s*を"),
    re.compile(r"を\s*→\s*が"),
]

def mentions_explicit_substitution(explanation: str) -> bool:
    """Detect explicit substitution statements like 'は → が' in explanation."""
    for p in SUBSTITUTION_PATTERNS:
        if p.search(explanation):
            return True
    return False

def validate_explanation_consistency(kanji: str, explanation: str) -> Tuple[bool, str]:
    """
    Ensure model does not claim replacements that didn't occur.
    If explanation contains explicit substitution arrows, verify that the target particle exists
    in the final 'kanji' output.
    """
    if not explanation:
        return False, "Empty explanation"

    if mentions_explicit_substitution(explanation):
        # If explanation mentions は→が, verify 'が' present in kanji
        if re.search(r"は\s*(?:→|->)\s*が", explanation) and "が" not in kanji:
            return False, "Explanation claims は→が but final sentence lacks 'が'"
        if re.search(r"が\s*(?:→|->)\s*を", explanation) and "を" not in kanji:
            return False, "Explanation claims が→を but final sentence lacks 'を'"
        if re.search(r"を\s*(?:→|->)\s*が", explanation) and "が" not in kanji:
            return False, "Explanation claims を→が but final sentence lacks 'が'"
    # Passed basic checks
    return True, ""

# =========================================================
# GENERATION (with retry and validation)
# =========================================================
def build_prompt(user_input: str) -> str:
    """Concatenate system + user into a single string prompt for non-chat tokenizers."""
    # Keep it explicit so LLM knows to return ONLY JSON
    return SYSTEM_PROMPT.strip() + "\n\nUser: Correct this Japanese: " + user_input.strip()

def generate_once(prompt_text: str, do_sample: bool = False, temperature: float = 0.0, top_p: float = 1.0) -> str:
    model_inputs = tokenizer(prompt_text, return_tensors="pt", padding=True)
    input_ids = model_inputs["input_ids"]
    attention_mask = model_inputs.get("attention_mask", None)

    with torch.no_grad():
        outputs = model.generate(
            input_ids=input_ids.to("cpu"),
            attention_mask=attention_mask.to("cpu") if attention_mask is not None else None,
            max_new_tokens=MAX_NEW_TOKENS,
            do_sample=do_sample,
            temperature=temperature,
            top_p=top_p,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

    # outputs is [batch, seq_len]; remove the prompt portion
    prompt_len = input_ids.shape[1]
    generated = outputs[:, prompt_len:]
    decoded = tokenizer.batch_decode(generated, skip_special_tokens=True)
    return decoded[0].strip()

def generate_response(user_input: str) -> Tuple[Optional[Dict], str]:
    """
    Try deterministic generation first. If explanation validation fails, retry once with light sampling.
    Returns (parsed_json_or_none, raw_output).
    """
    prompt = build_prompt(user_input)

    attempts = [
        {"do_sample": False, "temperature": 0.0, "top_p": 1.0},
        {"do_sample": True, "temperature": 0.25, "top_p": 0.9},
    ][:RETRY_ATTEMPTS]

    last_raw = ""
    for attempt_idx, attempt_cfg in enumerate(attempts, start=1):
        try:
            raw = generate_once(prompt, **attempt_cfg)
        except Exception as e:
            logger.exception("Model generation error on attempt %d: %s", attempt_idx, e)
            last_raw = f"GENERATION_ERROR: {e}"
            continue

        last_raw = raw
        logger.info("Generation attempt %d raw output: %s", attempt_idx, raw[:1000])

        parsed = extract_json(raw)
        if not parsed:
            logger.info("Attempt %d: could not parse JSON, will retry if attempts remain.", attempt_idx)
            continue

        # Validate schema minimally
        kanji = parsed.get("kanji", "")
        explanation = parsed.get("explanation", "")

        ok, reason = validate_explanation_consistency(kanji, explanation)
        if ok:
            return parsed, raw
        else:
            logger.warning("Attempt %d: Explanation consistency failed: %s", attempt_idx, reason)
            # continue to next attempt (if any) to try regenerate
            continue

    # If all attempts fail, return last parsed (if any) or None
    parsed_fallback = extract_json(last_raw) if last_raw else None
    return parsed_fallback, last_raw

# =========================================================
# ENDPOINT
# =========================================================
@app.post("/infer")
def infer(req: InferenceRequest):
    user_text = req.input.strip()
    if not user_text:
        return {"status": "error", "result": {"explanation": "Empty input"}}

    parsed, raw_output = generate_response(user_text)
    logger.info("Final raw output (truncated): %s", (raw_output or "")[:1200])

    if parsed:
        # final safety validate again and return helpful error if still inconsistent
        kanji = parsed.get("kanji", "")
        explanation = parsed.get("explanation", "")
        ok, reason = validate_explanation_consistency(kanji, explanation)
        if not ok:
            return {
                "status": "error",
                "result": {
                    "kanji": kanji,
                    "hiragana": parsed.get("hiragana", ""),
                    "explanation": f"Rejected by server-side consistency check: {reason}",
                    "raw_debug": raw_output,
                },
            }

        return {
            "status": "success",
            "result": {
                "kanji": kanji,
                "hiragana": parsed.get("hiragana", ""),
                "explanation": explanation,
            },
        }

    # no parsed JSON available
    return {
        "status": "error",
        "result": {
            "kanji": "",
            "hiragana": "",
            "explanation": "Could not parse model response as JSON after retries.",
            "raw_debug": raw_output,
        },
    }