from __future__ import annotations

import asyncio
import json
import logging
import re
from collections.abc import Coroutine
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, cast

import sentry_sdk
from ninja import NinjaAPI, Schema
from openai.types.chat import ChatCompletionSystemMessageParam, ChatCompletionUserMessageParam

from aiservice.analytics.posthog import ph
from aiservice.common_utils import validate_trace_id
from aiservice.env_specific import debug_log_sensitive_data
import stamina

from aiservice.llm import LLMOutputUnparseable, llm_client
from aiservice.llm_models import LLM, RANKING_MODEL
from authapp.auth import AuthenticatedRequest
from core.log_features.log_event import update_optimization_cost
from core.log_features.log_features import log_features

if TYPE_CHECKING:
    from openai.types.chat import ChatCompletionMessageParam


ranker_api = NinjaAPI(urls_namespace="ranker")

# Regex patterns for fallback parsing (legacy format)
rank_regex_pattern = re.compile(r"<rank>(.*)<\/rank>", re.DOTALL | re.IGNORECASE)
explain_regex_pattern = re.compile(r"<explain>(.*)<\/explain>", re.DOTALL | re.IGNORECASE)
scores_regex_pattern = re.compile(r"<scores>(.*)<\/scores>", re.DOTALL | re.IGNORECASE)

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

# Scoring dimensions and their weights (higher weight = more important)
# Weights should sum to 1.0 for normalized scoring
SCORING_DIMENSIONS = {
    "speedup_quality": 0.30,  # Quality of the speedup (not just magnitude, but how it's achieved)
    "readability": 0.25,  # Code readability and maintainability
    "diff_precision": 0.20,  # How focused/minimal is the change
    "correctness_confidence": 0.25,  # Confidence the optimization is correct/safe
}

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]]


@dataclass
class ParsedRankingResponse:
    """Structured response from ranking LLM."""

    ranking: list[int]  # 1-indexed candidate order (best first)
    scores: CandidateScores | None  # {candidate_idx: {dimension: score}}
    explanation: str


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

    Handles:
    - Raw JSON
    - JSON in markdown code blocks
    - JSON with surrounding text
    """
    # Try direct JSON parse first
    try:
        return cast(dict[str, Any], json.loads(content.strip()))
    except json.JSONDecodeError:
        pass

    # Try to extract from markdown code block or find JSON object
    match = json_pattern.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

    # Try to find any JSON-like structure with "ranking" key
    # Look for the outermost braces containing "ranking"
    try:
        start = content.find("{")
        if start != -1:
            # Find matching closing brace
            depth = 0
            for i, char in enumerate(content[start:], start):
                if char == "{":
                    depth += 1
                elif char == "}":
                    depth -= 1
                    if depth == 0:
                        json_str = content[start : i + 1]
                        return cast(dict[str, Any], json.loads(json_str))
    except json.JSONDecodeError:
        pass

    return None


def _parse_json_response(content: str, num_candidates: int) -> ParsedRankingResponse | None:
    """Parse structured JSON response from LLM.

    Expected format:
    {
        "scores": {"1": {"speedup_quality": 7, ...}, "2": {...}},
        "ranking": [2, 1, 3],
        "explanation": "..."
    }
    """
    data = _extract_json_from_response(content)
    if data is None:
        return None

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

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

        # Parse scores
        scores: CandidateScores | None = None
        scores_data = data.get("scores")
        if isinstance(scores_data, dict):
            scores = {}
            for candidate_key, dim_scores in scores_data.items():
                try:
                    candidate_idx = int(candidate_key)
                    if isinstance(dim_scores, dict):
                        parsed_scores: dict[str, float] = {}
                        for dim in SCORING_DIMENSIONS:
                            if dim in dim_scores:
                                score = float(dim_scores[dim])
                                # Clamp to valid range
                                parsed_scores[dim] = max(1.0, min(10.0, score))
                        if all(dim in parsed_scores for dim in SCORING_DIMENSIONS):
                            scores[candidate_idx] = parsed_scores
                except (ValueError, TypeError):
                    continue

            # Validate we have scores for all candidates
            if len(scores) != num_candidates or not all(i in scores for i in range(1, num_candidates + 1)):
                scores = None  # Invalid, fall back to no scores

        # Parse explanation
        explanation = str(data.get("explanation", ""))

        return ParsedRankingResponse(ranking=ranking, scores=scores, explanation=explanation)

    except (KeyError, TypeError, ValueError) as e:
        logging.warning("Failed to parse JSON response: %s", e)
        return None


def _parse_scores(scores_text: str, num_candidates: int) -> CandidateScores | None:
    """Parse the scores block from LLM output.

    Expected format:
    1: speedup_quality=7, readability=8, diff_precision=9, correctness_confidence=8
    2: speedup_quality=9, readability=6, diff_precision=7, correctness_confidence=7
    ...

    Returns dict mapping candidate index (1-indexed) to dimension scores, or None if parsing fails.
    """
    scores: CandidateScores = {}
    lines = scores_text.strip().split("\n")

    for line in lines:
        line = line.strip()
        if not line:
            continue

        # Parse "N: dimension=score, dimension=score, ..."
        try:
            idx_part, scores_part = line.split(":", 1)
            candidate_idx = int(idx_part.strip())

            candidate_scores: dict[str, float] = {}
            for pair in scores_part.split(","):
                pair = pair.strip()
                if "=" not in pair:
                    continue
                dimension, score_str = pair.split("=", 1)
                dimension = dimension.strip()
                score = float(score_str.strip())
                # Clamp to valid range
                score = max(1.0, min(10.0, score))
                candidate_scores[dimension] = score

            # Verify all dimensions are present
            if all(dim in candidate_scores for dim in SCORING_DIMENSIONS):
                scores[candidate_idx] = candidate_scores
        except (ValueError, IndexError):
            continue

    # Verify we got scores for all candidates
    if len(scores) == num_candidates and all(i in scores for i in range(1, num_candidates + 1)):
        return scores
    return None


def _compute_weighted_score(candidate_scores: dict[str, float]) -> float:
    """Compute weighted score from 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]:
    """Convert candidate scores to a ranking (best first, 1-indexed)."""
    # Calculate weighted score for each candidate
    weighted_scores = {
        candidate_idx: _compute_weighted_score(dim_scores) for candidate_idx, dim_scores in scores.items()
    }
    # Sort by weighted score descending (higher is better)
    sorted_candidates = sorted(weighted_scores.keys(), key=lambda c: weighted_scores[c], reverse=True)
    return sorted_candidates


@stamina.retry(on=LLMOutputUnparseable, attempts=2)
async def rank_optimizations(  # noqa: D417
    user_id: str, data: RankInputSchema, rank_model: LLM = RANKING_MODEL
) -> tuple[RankResponseSchema, float]:
    """Rank optimization candidates using multi-dimensional scoring.

    Parameters
    ----------
    - speedups list[str]: list of speedups of optimized candidates.
    - diffs list[str]: list of diffs of optimized candidates.
    - python_version (tuple[int, int, int]): The python version to use. Default is (3,12,9).

    Returns: - List[Tuple[Union[str, None], Union[str, None]]]: A list of tuples where the first element is the
    optimized code and the second is the explanation.
    :param optimization_ids:

    """
    debug_log_sensitive_data(f"Generating a ranking for {user_id}")
    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",
    )
    system_message = ChatCompletionSystemMessageParam(role="system", content=SYSTEM_PROMPT)
    user_message = ChatCompletionUserMessageParam(role="user", content=user_prompt)
    debug_log_sensitive_data(f"{SYSTEM_PROMPT}{user_prompt}")
    messages: list[ChatCompletionMessageParam] = [system_message, user_message]

    try:
        output = await llm_client.call(
            llm=rank_model,
            messages=messages,
            call_type="ranking",
            trace_id=data.trace_id,
            user_id=user_id,
            context={
                "num_candidates": len(data.diffs),
                "speedups": data.speedups,
                "python_version": data.python_version,
            },
        )
    except Exception as exc:
        logging.exception("Ranking failed for trace_id=%s", data.trace_id)
        raise LLMOutputUnparseable("Failed to rank optimizations") from exc

    debug_log_sensitive_data(f"AIClient optimization response:\n{output}")
    if output.raw_response.usage is not None:
        ph(
            user_id,
            "aiservice-optimize-openai-usage",
            properties={"model": rank_model.name, "n": 1, "usage": output.raw_response.usage.model_dump_json()},
        )

    num_candidates = len(data.diffs)

    # Try JSON parsing first (preferred structured format)
    json_response = _parse_json_response(output.content, num_candidates)
    if json_response is not None:
        logging.info("Successfully parsed JSON response")
        return (
            RankResponseSchema(
                ranking=_scores_to_ranking(json_response.scores)
                if json_response.scores is not None
                else json_response.ranking,
                explanation=json_response.explanation,
                scores=json_response.scores,
            ),
            output.cost,
        )

    # Fall back to regex parsing (legacy XML-tag format)
    logging.info("JSON parsing failed, falling back to regex")

    # Parse scores (optional - we can fall back to ranking if scores fail)
    scores: CandidateScores | None = None
    try:
        scores_match = re.search(scores_regex_pattern, output.content)
        if scores_match is not None:
            scores = _parse_scores(scores_match.group(1), num_candidates)
            if scores is None:
                logging.warning("Failed to parse scores")
    except Exception:  # noqa: BLE001
        logging.warning("Error parsing scores")

    # Parse explanation
    explanation = ""
    try:
        explanation_match = re.search(explain_regex_pattern, output.content)
        if explanation_match is not None:
            explanation = explanation_match.group(1).strip()
    except Exception:  # noqa: BLE001
        pass

    # Parse ranking
    ranking: list[int] | None = None
    try:
        ranking_match = re.search(rank_regex_pattern, output.content)
        if ranking_match is not None:
            ranking = list(map(int, ranking_match.group(1).strip().split(",")))
    except Exception:  # noqa: BLE001
        pass

    # Validate or derive ranking
    if ranking is None or sorted(ranking) != list(range(1, num_candidates + 1)):
        # If ranking is invalid but we have scores, derive from scores
        if scores is not None:
            ranking = _scores_to_ranking(scores)
            logging.info("Derived ranking from scores")
        else:
            logging.warning("No valid ranking found")
            raise LLMOutputUnparseable("No valid ranking found", cost=output.cost)

    return RankResponseSchema(ranking=ranking, explanation=explanation, scores=scores), output.cost


class RankInputSchema(Schema):
    trace_id: str
    speedups: list[float]
    diffs: list[str]
    optimization_ids: list[str]  # which diff corresponded to which opt candidate
    python_version: str | None = None
    function_references: str | None = None


class RankResponseSchema(Schema):
    ranking: list[int]
    explanation: str
    scores: dict[int, dict[str, float]] | None = None  # {candidate_idx: {dimension: score}}


class RankErrorResponseSchema(Schema):
    error: str


@ranker_api.post(
    "/",
    response={
        200: RankResponseSchema,
        400: RankErrorResponseSchema,
        422: RankErrorResponseSchema,
        500: RankErrorResponseSchema,
    },
)
async def rank(
    request: AuthenticatedRequest, data: RankInputSchema
) -> tuple[int, RankResponseSchema | RankErrorResponseSchema]:
    ph(request.user, "aiservice-rank-called")
    if not validate_trace_id(data.trace_id):
        return 400, RankErrorResponseSchema(error="Invalid trace ID. Please provide a valid UUIDv4.")
    try:
        ranking_response, llm_cost = await rank_optimizations(request.user, data)
    except LLMOutputUnparseable:
        ph(request.user, "Invalid Ranking, fallback to default")
        debug_log_sensitive_data("No valid ranking was generated")
        return 422, RankErrorResponseSchema(error="Error generating ranking. Internal server error.")
    ph(request.user, "ranking generated", properties={"ranking": ranking_response})
    ranking_0_idx = [x - 1 for x in ranking_response.ranking]
    coros: list[Coroutine[Any, Any, Any]] = [
        update_optimization_cost(trace_id=data.trace_id, cost=llm_cost, user_id=request.user)
    ]
    if hasattr(request, "should_log_features") and request.should_log_features:
        ranked_opt_ids = [data.optimization_ids[i] for i in ranking_0_idx]
        coros.append(
            log_features(
                trace_id=data.trace_id,
                user_id=request.user,
                ranking={"ranking": ranked_opt_ids, "explanation": ranking_response.explanation},
            )
        )
    results = await asyncio.gather(*coros, return_exceptions=True)
    for result in results:
        if isinstance(result, BaseException):
            sentry_sdk.capture_exception(result)
    return 200, RankResponseSchema(
        ranking=ranking_0_idx, explanation=ranking_response.explanation, scores=ranking_response.scores
    )  # we don't really use `explanation` and `score` but still returning it for future use in CLI
