from __future__ import annotations

import json
import logging
import re
from typing import TYPE_CHECKING, Annotated, Any, cast

from fastapi import APIRouter, Depends, HTTPException, Request

from codeflash_api.auth._deps import (
    check_rate_limit,
    require_auth,
    track_usage,
)
from codeflash_api.auth.models import AuthenticatedUser  # noqa: TC001
from codeflash_api.llm._models import RANKING_MODEL
from codeflash_api.optimize._context import validate_trace_id
from codeflash_api.ranking.schemas import (
    RankErrorResponse,
    RankRequest,
    RankResponse,
)

if TYPE_CHECKING:
    from codeflash_api.llm._client import LLMClient

log = logging.getLogger(__name__)

router = APIRouter()

# Scoring dimensions and their weights (must sum to 1.0).
SCORING_DIMENSIONS: dict[str, float] = {
    "speedup_quality": 0.30,
    "readability": 0.25,
    "diff_precision": 0.20,
    "correctness_confidence": 0.25,
}

# Regex for legacy XML-tag fallback parsing.
_RANK_RE = re.compile(r"<rank>(.*)</rank>", re.DOTALL | re.IGNORECASE)
_EXPLAIN_RE = re.compile(r"<explain>(.*)</explain>", re.DOTALL | re.IGNORECASE)
_SCORES_RE = re.compile(r"<scores>(.*)</scores>", re.DOTALL | re.IGNORECASE)

# Pattern to extract JSON from markdown code blocks or raw JSON.
_JSON_RE = re.compile(
    r"```(?:json)?\s*(\{.*?\})\s*```"
    r'|(\{[^`]*"ranking"[^`]*\})',
    re.DOTALL,
)

SYSTEM_PROMPT = """\
You are an expert code reviewer who understands why programs run fast.

You are provided with a list of optimization candidates with their code \
diff with respect to the baseline code and speedup ratio information. \
Your task is to evaluate each candidate on multiple dimensions and \
provide scores.

You are also provided with the following information.
- python_version - The version of python the code would be executed on.
- function_references - Python markdown blocks with filename and \
references of some functions which call the function being optimized. \
The filenames and/or references could indicate if the function being \
optimized is in a hot path. The reference could have the function \
being called from a place that is important, for example in a loop, \
which means the effect of optimization might be important.

## Scoring Dimensions (score each from 1-10, where 10 is best):

1. **speedup_quality** - Quality of the performance improvement:
   - 10: Significant speedup achieved through clean, algorithmic \
improvements
   - 7-9: Good speedup with reasonable implementation
   - 4-6: Moderate speedup or achieved through acceptable \
micro-optimizations
   - 1-3: Minimal speedup or achieved through hacky/unacceptable means

2. **readability** - Code readability and maintainability:
   - 10: Code is cleaner or equally readable as original
   - 7-9: Minor readability trade-offs that are acceptable
   - 4-6: Noticeable readability impact but still maintainable
   - 1-3: Significantly harder to read/maintain
   - **IMPORTANT**: Use of `global`/`nonlocal` keywords should \
score 1-3
   - **IMPORTANT**: Replacing `isinstance()` with `type()` should \
score 1-3

3. **diff_precision** - How focused and minimal is the change:
   - 10: Minimal, surgical change that only touches what's necessary
   - 7-9: Focused change with minor extra modifications
   - 4-6: Moderate scope, some unnecessary changes
   - 1-3: Large, sprawling diff that touches many things

4. **correctness_confidence** - Confidence the optimization preserves \
correctness:
   - 10: Obviously correct, no edge cases affected
   - 7-9: Very likely correct, minimal risk
   - 4-6: Probably correct but has some risk
   - 1-3: Risky, may break edge cases or change behavior

## Rules:
- DISREGARD the fact that new dependencies could be introduced in \
the diff, these are already installed in the system.
- Prefer optimizations with higher speedup ratios. If the higher \
speedups happen due to strange hacks or micro-optimizations or \
something an expert won't write then prefer it less.
- Prefer optimizations which contain precise diffs unless the speedup \
provided is very high. Larger pull requests are typically harder to \
accept then more precise smaller pull requests.
- Introduction of the `global` and `nonlocal` keywords in \
optimizations is **HIGHLY DISCOURAGED** as it reduces code clarity \
and maintainability, introduces hidden dependencies, can cause subtle \
bugs and breaks modularity. **DO NOT** prefer such optimizations.
- Replacement of `isinstance()` checks with `type()` checks is \
**HIGHLY DISCOURAGED** as `isinstance()` correctly handles \
inheritance and subclasses, while `type()` checks are incorrect for \
subclass instances and represent a micro-optimization that should be \
avoided. Do not prefer such optimizations.
- If the only optimizations are micro-optimizations like inlining a \
function call, or localizing variables or methods (attribute lookup \
optimizations), do not prefer the optimizations. The performance \
improvements are minimal and come at a substantial cost to readability.
- Local variable caching of globals (e.g., `local_var = GLOBAL_VAR` \
before a loop) is a micro-optimization only relevant on Python \
<=3.10. On Python 3.11+ `LOAD_GLOBAL` uses adaptive specialization \
and is nearly as fast as `LOAD_FAST`, so the benefit is negligible \
-- prefer the simpler code WITHOUT the local cache.

## Response Format

You MUST respond with a valid JSON object. Do not include any text \
before or after the JSON.

```json
{
  "scores": {
    "1": {"speedup_quality": 7, "readability": 8, \
"diff_precision": 9, "correctness_confidence": 8},
    "2": {"speedup_quality": 9, "readability": 6, \
"diff_precision": 7, "correctness_confidence": 7}
  },
  "ranking": [2, 1],
  "explanation": "Candidate 2 has the best speedup quality despite \
slightly lower readability..."
}
```

The JSON schema:
- "scores": Object mapping candidate number (as string "1", "2", \
etc.) to dimension scores (integers 1-10)
- "ranking": Array of candidate numbers in decreasing order of \
viability (best first)
- "explanation": Brief explanation of the scoring and ranking decisions

Example for 3 candidates where candidate 2 is best:
```json
{
  "scores": {
    "1": {"speedup_quality": 7, "readability": 8, \
"diff_precision": 9, "correctness_confidence": 8},
    "2": {"speedup_quality": 9, "readability": 6, \
"diff_precision": 7, "correctness_confidence": 7},
    "3": {"speedup_quality": 5, "readability": 4, \
"diff_precision": 6, "correctness_confidence": 9}
  },
  "ranking": [2, 1, 3],
  "explanation": "Candidate 2 has the best speedup quality despite \
slightly lower readability. Candidate 1 offers good balance. \
Candidate 3 has poor readability."
}
```\
"""

USER_PROMPT = """\
Here is a numbered list of optimization candidates' code diffs and \
their speedup ratios.

{ranking_context}

Here is the python version
{python_version}

Here are the function references
{function_references}\
"""

# Type alias for candidate scores: {candidate_idx: {dimension: score}}
CandidateScores = dict[int, dict[str, float]]


def _try_direct_json(content: str) -> dict[str, Any] | None:
    """
    Attempt to parse *content* as raw JSON.
    """
    try:
        return cast("dict[str, Any]", json.loads(content.strip()))
    except json.JSONDecodeError:
        return None


def _try_regex_json(content: str) -> dict[str, Any] | None:
    """
    Extract JSON from a markdown code block or embedded object.
    """
    match = _JSON_RE.search(content)
    if match:
        json_str = match.group(1) or match.group(2)
        if json_str:
            try:
                return cast(
                    "dict[str, Any]",
                    json.loads(json_str.strip()),
                )
            except json.JSONDecodeError:
                pass
    return None


def _try_brace_match_json(
    content: str,
) -> dict[str, Any] | None:
    """
    Find the outermost ``{...}`` in *content* and parse it as JSON.
    """
    start = content.find("{")
    if start == -1:
        return None
    depth = 0
    for i, char in enumerate(content[start:], start):
        if char == "{":
            depth += 1
        elif char == "}":
            depth -= 1
            if depth == 0:
                try:
                    return cast(
                        "dict[str, Any]",
                        json.loads(content[start : i + 1]),
                    )
                except json.JSONDecodeError:
                    return None
    return None


def _extract_json_from_response(
    content: str,
) -> dict[str, Any] | None:
    """
    Extract a JSON object from LLM response text.

    Tries direct ``json.loads``, then markdown code-block extraction,
    then brace-matching as a last resort.
    """
    return (
        _try_direct_json(content)
        or _try_regex_json(content)
        or _try_brace_match_json(content)
    )


def _parse_score_line(
    raw_line: str,
) -> tuple[int, dict[str, float]] | None:
    """
    Parse one legacy score line like ``1: dim=val, dim=val, ...``.

    Returns ``(candidate_idx, dimension_scores)`` or *None*.
    """
    stripped = raw_line.strip()
    if not stripped:
        return None
    try:
        idx_part, scores_part = stripped.split(":", 1)
        candidate_idx = int(idx_part.strip())
        candidate_scores: dict[str, float] = {}
        for raw_pair in scores_part.split(","):
            token = raw_pair.strip()
            if "=" not in token:
                continue
            dimension, score_str = token.split("=", 1)
            score = float(score_str.strip())
            candidate_scores[dimension.strip()] = max(1.0, min(10.0, score))
        if all(dim in candidate_scores for dim in SCORING_DIMENSIONS):
            return candidate_idx, candidate_scores
    except (ValueError, IndexError):
        pass
    return None


def _parse_scores_text(
    scores_text: str,
    num_candidates: int,
) -> CandidateScores | None:
    """
    Parse a legacy text-based scores block from LLM output.

    Expected format per line::

        1: speedup_quality=7, readability=8, ...
    """
    scores: CandidateScores = {}
    for raw_line in scores_text.strip().split("\n"):
        result = _parse_score_line(raw_line)
        if result is not None:
            scores[result[0]] = result[1]

    if len(scores) == num_candidates and all(
        i in scores for i in range(1, num_candidates + 1)
    ):
        return scores
    return None


def _parse_candidate_scores(
    scores_data: dict[str, Any],
    num_candidates: int,
) -> CandidateScores | None:
    """
    Parse scores from the JSON ``"scores"`` object.

    Returns validated scores or *None* if incomplete.
    """
    scores: CandidateScores = {}
    for candidate_key, dim_scores in scores_data.items():
        candidate_idx, parsed = _parse_one_candidate_score(
            candidate_key, dim_scores
        )
        if candidate_idx is not None and parsed is not None:
            scores[candidate_idx] = parsed

    if len(scores) != num_candidates or not all(
        i in scores for i in range(1, num_candidates + 1)
    ):
        return None
    return scores


def _parse_one_candidate_score(
    candidate_key: str,
    dim_scores: Any,
) -> tuple[int | None, dict[str, float] | None]:
    """
    Parse a single candidate's scores from JSON data.
    """
    try:
        candidate_idx = int(candidate_key)
    except (ValueError, TypeError):
        return None, None

    if not isinstance(dim_scores, dict):
        return None, None

    parsed: dict[str, float] = {}
    for dim in SCORING_DIMENSIONS:
        if dim in dim_scores:
            val = float(dim_scores[dim])
            parsed[dim] = max(1.0, min(10.0, val))

    if not all(dim in parsed for dim in SCORING_DIMENSIONS):
        return None, None
    return candidate_idx, parsed


def _parse_json_response(
    content: str,
    num_candidates: int,
) -> tuple[list[int], CandidateScores | None, str] | None:
    """
    Parse a structured JSON response from the LLM.

    Returns ``(ranking, scores, explanation)`` on success, or *None*
    when the response cannot be parsed into a valid ranking.
    """
    data = _extract_json_from_response(content)
    if data is None:
        return None

    try:
        ranking_data = data.get("ranking")
        if not isinstance(ranking_data, list):
            return None
        ranking = [int(x) for x in ranking_data]

        if sorted(ranking) != list(range(1, num_candidates + 1)):
            log.warning("Invalid ranking in JSON response: %s", ranking)
            return None

        scores: CandidateScores | None = None
        scores_data = data.get("scores")
        if isinstance(scores_data, dict):
            scores = _parse_candidate_scores(scores_data, num_candidates)

        explanation = str(data.get("explanation", ""))
    except (KeyError, TypeError, ValueError):
        log.warning(
            "Failed to parse JSON ranking response",
            exc_info=True,
        )
        return None
    else:
        return ranking, scores, explanation


def _compute_weighted_score(
    candidate_scores: dict[str, float],
) -> float:
    """
    Compute a single weighted score from per-dimension scores.
    """
    total = 0.0
    for dimension, weight in SCORING_DIMENSIONS.items():
        total += candidate_scores.get(dimension, 5.0) * weight
    return total


def _scores_to_ranking(scores: CandidateScores) -> list[int]:
    """
    Derive a best-first 1-indexed ranking from candidate scores.
    """
    weighted = {
        idx: _compute_weighted_score(dim_scores)
        for idx, dim_scores in scores.items()
    }
    return sorted(weighted, key=lambda c: weighted[c], reverse=True)


def _parse_legacy_xml(
    content: str,
    num_candidates: int,
) -> tuple[list[int] | None, CandidateScores | None, str]:
    """
    Parse legacy XML-tag format from LLM output.

    Returns ``(ranking, scores, explanation)``.  *ranking* may be
    *None* if it could not be extracted or validated.
    """
    scores: CandidateScores | None = None
    scores_match = _SCORES_RE.search(content)
    if scores_match is not None:
        scores = _parse_scores_text(scores_match.group(1), num_candidates)
        if scores is None:
            log.warning("Failed to parse scores from XML tags")

    explanation = ""
    explain_match = _EXPLAIN_RE.search(content)
    if explain_match is not None:
        explanation = explain_match.group(1).strip()

    ranking: list[int] | None = None
    rank_match = _RANK_RE.search(content)
    if rank_match is not None:
        ranking = list(map(int, rank_match.group(1).strip().split(",")))

    if ranking is None or sorted(ranking) != list(
        range(1, num_candidates + 1)
    ):
        if scores is not None:
            ranking = _scores_to_ranking(scores)
            log.info("Derived ranking from scores")
        else:
            ranking = None

    return ranking, scores, explanation


async def rank_optimizations(
    llm_client: LLMClient,
    user_id: str,
    data: RankRequest,
) -> tuple[list[int], CandidateScores | None, str, float]:
    """
    Rank optimization candidates via multi-dimensional LLM scoring.

    Returns ``(ranking_1idx, scores, explanation, cost)``.
    The ranking is **1-indexed** (best first).
    """
    ranking_context = ""
    for i, (diff, speedup) in enumerate(
        zip(data.diffs, data.speedups, strict=False)
    ):
        ranking_context += (
            f"{i + 1}. Diff:\n```diff\n{diff}\n```\nSpeedup: {speedup:.3f}\n"
        )

    user_prompt = USER_PROMPT.format(
        ranking_context=ranking_context,
        python_version=data.python_version or "Not available",
        function_references=(data.function_references or "Not available"),
    )

    messages: list[dict[str, Any]] = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_prompt},
    ]

    output = await llm_client.call(
        llm=RANKING_MODEL,
        messages=messages,
        call_type="ranking",
        trace_id=data.trace_id,
    )

    num_candidates = len(data.diffs)

    json_result = _parse_json_response(output.content, num_candidates)
    if json_result is not None:
        ranking, scores, explanation = json_result
        log.info("Successfully parsed JSON ranking response")
        if scores is not None:
            ranking = _scores_to_ranking(scores)
        return ranking, scores, explanation, output.cost

    log.info("JSON parsing failed, falling back to regex")

    ranking_fb, scores_fb, explanation_fb = _parse_legacy_xml(
        output.content, num_candidates
    )
    if ranking_fb is None:
        log.warning(
            "No valid ranking found for trace_id=%s",
            data.trace_id,
        )
        raise HTTPException(
            status_code=422,
            detail=("Error generating ranking. Internal server error."),
        )

    return ranking_fb, scores_fb, explanation_fb, output.cost


@router.post(
    "/ai/rank",
    response_model=RankResponse,
    responses={
        400: {"model": RankErrorResponse},
        422: {"model": RankErrorResponse},
        500: {"model": RankErrorResponse},
    },
)
async def rank(
    request: Request,
    data: RankRequest,
    user: Annotated[AuthenticatedUser, Depends(require_auth)],
    _rate: Annotated[None, Depends(check_rate_limit)],
    _usage: Annotated[None, Depends(track_usage)],
) -> RankResponse:
    """
    Rank optimization candidates by multi-dimensional scoring.
    """
    if not validate_trace_id(data.trace_id):
        raise HTTPException(
            status_code=400,
            detail=("Invalid trace ID. Please provide a valid UUIDv4."),
        )

    llm_client: LLMClient = request.app.state.llm_client

    try:
        ranking_1idx, scores, explanation, _cost = await rank_optimizations(
            llm_client, user.user_id, data
        )
    except HTTPException:
        raise
    except Exception as exc:
        log.exception("Ranking failed for trace_id=%s", data.trace_id)
        raise HTTPException(
            status_code=422,
            detail=("Error generating ranking. Internal server error."),
        ) from exc

    ranking_0idx = [x - 1 for x in ranking_1idx]

    return RankResponse(
        ranking=ranking_0idx,
        explanation=explanation,
        scores=scores,
    )
