import json
import logging

import litellm

from app.config import settings

logger = logging.getLogger(__name__)


def _proxy_model(model_id: str) -> str:
    """Prefix model with openai/ if using a LiteLLM proxy and no provider is specified."""
    if settings.LITELLM_API_BASE and "/" not in model_id:
        return f"openai/{model_id}"
    return model_id

EXTRACTION_PROMPT = """You are extracting questions from a PREP (Pediatric Review and Education Program) exam PDF.

These PDFs have questions with a clinical vignette, followed by five answer options A-E.
Some PDFs include "Correct Answer: X" or "Preferred Response: X" right after the options.
Some PDFs have the correct answers only at the END of the document — in that case correct_answer will be null.

Return ONLY a JSON object:

{{"questions": [
  {{
    "item_number": "<item number like '193' or '21' — digits only, null if not found>",
    "question_text": "<full question stem including the clinical vignette>",
    "question_type": "mcq",
    "options": ["<option A text>", "<option B text>", "<option C text>", "<option D text>", "<option E text>"],
    "correct_answer": "<full text of the correct option, NOT the letter — or null if not found>",
    "explanation": "<explanation/critique/content specs if present, else empty string>",
    "page_reference": {page_ref}
  }}
]}}

RULES:
- A new question starts when you see "Item NNN" or "ltem NNN" (OCR artifact for Item).
- Extract question_text as the full vignette + the "Of the following..." or "Which of the following..." stem.
- Options are labeled A. B. C. D. E. — extract just the text, not the letter.
- If "Correct Answer: X" or "Preferred Response: X" appears after the options, resolve it to full option text.
- CRITICAL: Content immediately after an image reference (e.g. "Item C123A", "Item C123B", figure captions,
  table data) is NOT a new question. Skip it and wait for the next "Item NNN" number.
- If no correct answer line is found for a question, set correct_answer to null — do NOT guess.
- Do NOT omit questions — extract all questions found in the content, even if partial.
- Return ONLY the JSON — no markdown, no preamble.

Content from page(s) {page_info}:
{content}"""

ANSWER_KEY_PROMPT = """Extract the answer key from this PREP exam content.

The answer key lists items with their correct answer letters, like:
  "Item 193  Preferred Response: D"
  "Item 194  Preferred Response: A"

Return ONLY a JSON object mapping item numbers to correct letters:
{{"answers": {{"193": "D", "194": "A", "211": "C"}}}}

Rules:
- "Preferred Response: X" or "Correct Answer: X" — X is the letter.
- Item numbers may appear as "ltemXXX" (OCR artifact — l is actually I).
- Only include items where you find a clear correct answer letter.
- Return ONLY the JSON — no markdown, no preamble.

Content from page(s) {page_info}:
{content}"""


def get_model_for_task(db, task: str = "extraction") -> tuple[str, str | None]:
    """Get the configured model for a specific task from DB, or fall back to settings."""
    try:
        from app.models.ai_model_config import AIModelConfig
        config = db.query(AIModelConfig).filter(
            AIModelConfig.task == task,
            AIModelConfig.is_active == True,
            AIModelConfig.is_default == True,
        ).first()
        if config:
            return config.model_id, config.api_key
    except Exception:
        pass
    return settings.LITELLM_MODEL, settings.LITELLM_API_KEY or None


def _truncate_content(content: str, max_chars: int = 100000) -> str:
    if len(content) <= max_chars:
        return content
    half = max_chars // 2
    return content[:half] + "\n\n... [content truncated] ...\n\n" + content[-half:]


def extract_questions(
    content: str,
    page_info: str = "unknown",
    page_ref: int | None = None,
    model_id: str | None = None,
    api_key: str | None = None,
) -> list[dict]:
    """Extract quiz questions from PDF content using LiteLLM."""
    content = _truncate_content(content)

    prompt = EXTRACTION_PROMPT.format(
        content=content,
        page_info=page_info,
        page_ref=page_ref if page_ref else "null",
    )

    use_model = _proxy_model(model_id or settings.LITELLM_MODEL)
    use_key = api_key or settings.LITELLM_API_KEY

    last_error = None
    for attempt in range(3):
        try:
            kwargs = {
                "model": use_model,
                "messages": [{"role": "user", "content": prompt}],
                "temperature": 0.1,  # low temp for faithful extraction
            }
            if use_key:
                kwargs["api_key"] = use_key
            if settings.LITELLM_API_BASE:
                kwargs["api_base"] = settings.LITELLM_API_BASE

            # Don't force JSON mode — let the model respond naturally and we parse it
            response = litellm.completion(**kwargs)

            response_text = response.choices[0].message.content
            logger.info(f"Model raw response (first 500 chars): {response_text[:500]!r}")

            # Try to parse JSON, handle markdown code blocks
            text = response_text.strip()
            if text.startswith("```"):
                text = text.split("\n", 1)[1] if "\n" in text else text[3:]
                if text.endswith("```"):
                    text = text[:-3]
                text = text.strip()

            data = json.loads(text)

            # Handle all common response shapes
            if isinstance(data, list):
                questions = data
            elif isinstance(data, dict):
                # Try common keys
                for key in ("questions", "items", "results", "data"):
                    if isinstance(data.get(key), list):
                        questions = data[key]
                        break
                else:
                    # Maybe the dict itself is a single question
                    if "question_text" in data:
                        questions = [data]
                    else:
                        raise ValueError(f"Unexpected response shape: {list(data.keys())}")
            else:
                raise ValueError("Response is not a list of questions")

            validated = []
            skipped = []
            for q in questions:
                if "question_text" not in q:
                    continue
                correct = q.get("correct_answer")
                if not correct:
                    skipped.append(q.get("question_text", "")[:120])
                    continue
                qtype = q.get("question_type", "mcq")
                if qtype not in ("mcq", "true_false", "fill_blank"):
                    qtype = "mcq"
                validated.append({
                    "question_text": q["question_text"],
                    "question_type": qtype,
                    "options": q.get("options"),
                    "correct_answer": correct,
                    "explanation": q.get("explanation", ""),
                    "page_reference": q.get("page_reference"),
                    "skipped": [],
                })

            # Attach skipped list to first question so caller can surface it
            if validated and skipped:
                validated[0]["skipped"] = skipped

            if validated:
                return validated

            raise ValueError("No valid questions extracted from content")

        except Exception as e:
            last_error = e
            logger.warning(f"Extraction attempt {attempt + 1} failed: {e!r}")

    raise RuntimeError(f"Failed to extract questions after 3 attempts: {last_error}")


def _call_model(prompt: str, model_id: str | None, api_key: str | None) -> str:
    """Call the configured LLM and return raw text response."""
    use_model = _proxy_model(model_id or settings.LITELLM_MODEL)
    use_key = api_key or settings.LITELLM_API_KEY
    kwargs = {
        "model": use_model,
        "messages": [{"role": "user", "content": prompt}],
        "temperature": 0.1,
    }
    if use_key:
        kwargs["api_key"] = use_key
    if settings.LITELLM_API_BASE:
        kwargs["api_base"] = settings.LITELLM_API_BASE
    response = litellm.completion(**kwargs)
    return response.choices[0].message.content


def extract_questions_no_answers(
    content: str,
    page_info: str = "unknown",
    page_ref: int | None = None,
    model_id: str | None = None,
    api_key: str | None = None,
) -> list[dict]:
    """Extract questions allowing null correct_answer — for PDFs where answers are at the end."""
    content = _truncate_content(content)
    prompt = EXTRACTION_PROMPT.format(
        content=content,
        page_info=page_info,
        page_ref=page_ref if page_ref else "null",
    )
    last_error = None
    for attempt in range(3):
        try:
            text = _call_model(prompt, model_id, api_key)
            text = text.strip()
            if text.startswith("```"):
                text = text.split("\n", 1)[1] if "\n" in text else text[3:]
                if text.endswith("```"):
                    text = text[:-3]
                text = text.strip()
            data = json.loads(text)
            questions = []
            if isinstance(data, list):
                questions = data
            elif isinstance(data, dict):
                for key in ("questions", "items", "results", "data"):
                    if isinstance(data.get(key), list):
                        questions = data[key]
                        break
                else:
                    if "question_text" in data:
                        questions = [data]

            result = []
            for q in questions:
                if "question_text" not in q:
                    continue
                qtype = q.get("question_type", "mcq")
                if qtype not in ("mcq", "true_false", "fill_blank"):
                    qtype = "mcq"
                result.append({
                    "item_number": str(q.get("item_number") or "").strip().lstrip("0") or None,
                    "question_text": q["question_text"],
                    "question_type": qtype,
                    "options": q.get("options"),
                    "correct_answer": q.get("correct_answer"),  # may be null
                    "explanation": q.get("explanation", ""),
                    "page_reference": q.get("page_reference"),
                })
            if result:
                return result
            raise ValueError("No questions found in content")
        except Exception as e:
            last_error = e
            logger.warning(f"No-answer extraction attempt {attempt + 1} failed: {e!r}")
    raise RuntimeError(f"Failed after 3 attempts: {last_error}")


def extract_answer_key(
    content: str,
    page_info: str = "unknown",
    model_id: str | None = None,
    api_key: str | None = None,
) -> dict[str, str]:
    """Extract answer key from end-of-document content. Returns {item_number: letter}."""
    content = _truncate_content(content, max_chars=80000)
    prompt = ANSWER_KEY_PROMPT.format(content=content, page_info=page_info)
    last_error = None
    for attempt in range(3):
        try:
            text = _call_model(prompt, model_id, api_key)
            text = text.strip()
            if text.startswith("```"):
                text = text.split("\n", 1)[1] if "\n" in text else text[3:]
                if text.endswith("```"):
                    text = text[:-3]
                text = text.strip()
            data = json.loads(text)
            answers = data.get("answers", data) if isinstance(data, dict) else {}
            # Normalize: strip leading zeros, uppercase letters
            return {str(k).strip().lstrip("0"): str(v).strip().upper() for k, v in answers.items() if v}
        except Exception as e:
            last_error = e
            logger.warning(f"Answer key extraction attempt {attempt + 1} failed: {e!r}")
    logger.warning(f"Answer key extraction failed: {last_error}")
    return {}


def generate_tts_audio(
    text: str,
    model_id: str | None = None,
    api_key: str | None = None,
) -> bytes | None:
    """Generate TTS audio. Supports OpenAI, ElevenLabs, Google Cloud TTS, and AWS Polly.

    model_id conventions:
      tts-1:alloy          → OpenAI TTS (voice after colon)
      tts-1-hd:nova        → OpenAI TTS HD
      elevenlabs/<voice>   → ElevenLabs
      google/<voice_name>  → Google Cloud TTS (e.g. google/en-US-Wavenet-D)
      polly/<VoiceId>      → AWS Polly Neural (e.g. polly/Joanna)
    """
    import httpx, base64

    use_model = model_id or "tts-1:alloy"

    # ── ElevenLabs ─────────────────────────────────────────────
    if use_model.startswith("elevenlabs/") or use_model.startswith("eleven_labs/"):
        voice = use_model.split("/", 1)[1]
        key = api_key or settings.ELEVENLABS_API_KEY
        if not key:
            logger.error("ElevenLabs API key not configured")
            return None
        try:
            resp = httpx.post(
                f"https://api.elevenlabs.io/v1/text-to-speech/{voice}",
                headers={"xi-api-key": key, "Content-Type": "application/json"},
                json={"text": text, "model_id": "eleven_turbo_v2_5"},
                timeout=30,
            )
            resp.raise_for_status()
            return resp.content
        except Exception as e:
            logger.error(f"ElevenLabs TTS failed: {e}")
            return None

    # ── Google Cloud TTS ────────────────────────────────────────
    if use_model.startswith("google/"):
        voice_name = use_model[len("google/"):]
        key = api_key or settings.GOOGLE_TTS_API_KEY
        if not key:
            logger.error("Google TTS API key not configured (GOOGLE_TTS_API_KEY)")
            return None
        # Parse language code from voice name (e.g. "en-US-Wavenet-D" → "en-US")
        parts = voice_name.split("-")
        lang_code = f"{parts[0]}-{parts[1]}" if len(parts) >= 2 else "en-US"
        try:
            resp = httpx.post(
                f"https://texttospeech.googleapis.com/v1/text:synthesize?key={key}",
                json={
                    "input": {"text": text},
                    "voice": {"languageCode": lang_code, "name": voice_name},
                    "audioConfig": {"audioEncoding": "MP3"},
                },
                timeout=30,
            )
            resp.raise_for_status()
            return base64.b64decode(resp.json()["audioContent"])
        except Exception as e:
            logger.error(f"Google TTS failed: {e}")
            return None

    # ── AWS Polly ───────────────────────────────────────────────
    if use_model.startswith("polly/"):
        voice_id = use_model[len("polly/"):]
        access_key = api_key or settings.AWS_ACCESS_KEY_ID
        secret_key = settings.AWS_SECRET_ACCESS_KEY
        region = settings.AWS_REGION or "us-east-1"
        if not access_key or not secret_key:
            logger.error("AWS credentials not configured (AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY)")
            return None
        try:
            import boto3
            polly = boto3.client(
                "polly",
                aws_access_key_id=access_key,
                aws_secret_access_key=secret_key,
                region_name=region,
            )
            response = polly.synthesize_speech(
                Text=text,
                OutputFormat="mp3",
                VoiceId=voice_id,
                Engine="neural",
            )
            return response["AudioStream"].read()
        except Exception as e:
            logger.error(f"AWS Polly TTS failed: {e}")
            return None

    # ── OpenAI (default) ────────────────────────────────────────
    # model_id may encode voice as "tts-1:nova", "tts-1-hd:alloy", etc.
    clean_model = use_model.replace("openai/", "")
    oai_voice = "alloy"
    if ":" in clean_model:
        clean_model, oai_voice = clean_model.split(":", 1)

    # Per-model key > OPENAI_API_KEY (direct) > LITELLM_API_KEY (proxy)
    if api_key:
        key = api_key
        base = (settings.LITELLM_API_BASE or "https://api.openai.com").rstrip("/")
    elif settings.OPENAI_API_KEY:
        key = settings.OPENAI_API_KEY
        base = "https://api.openai.com"
    else:
        key = settings.LITELLM_API_KEY
        base = (settings.LITELLM_API_BASE or "https://api.openai.com").rstrip("/")

    try:
        resp = httpx.post(
            f"{base}/v1/audio/speech",
            headers={"Authorization": f"Bearer {key}", "Content-Type": "application/json"},
            json={"model": clean_model, "input": text, "voice": oai_voice},
            timeout=60,
        )
        resp.raise_for_status()
        return resp.content
    except Exception as e:
        logger.error(f"OpenAI TTS failed: {e}")
        return None