from __future__ import annotations
import json
from pathlib import Path
from typing import Any, Dict, List, Optional
from codestar.core.schemas import ProblemMeta

DEFAULT_PROBLEMS_DIR = Path(__file__).resolve().parents[2] / "data" / "problems"


class ProblemRepository:
    def __init__(self, problems_dir: Path = DEFAULT_PROBLEMS_DIR):
        self.problems_dir = problems_dir
        self._cache: Dict[str, ProblemMeta] = {}
        self._list_items_cache: List[Dict[str, Any]] = []
        self._by_category_cache: Dict[str, List[Dict[str, Any]]] = {}
        self._featured_cache: List[Dict[str, Any]] = []
        self._difficulty_breakdowns: Dict[str, Dict[str, int]] = {}
        self._test_cases: Dict[str, List[Dict[str, Any]]] = {}
        self.reload()

    def reload(self):
        self._cache.clear()
        self._list_items_cache.clear()
        self._by_category_cache.clear()
        self._featured_cache.clear()
        self._difficulty_breakdowns.clear()
        self._test_cases.clear()

        if not self.problems_dir.exists():
            self.problems_dir.mkdir(parents=True, exist_ok=True)
            return

        for path in sorted(self.problems_dir.glob("*.json")):
            try:
                data = json.loads(path.read_text(encoding="utf-8"))
                problem = ProblemMeta(**data)
                self._cache[problem.slug] = problem

                # Extract 2D difficulty breakdown with fallback calibration
                breakdown = data.get("difficulty_breakdown")
                if not breakdown or not isinstance(breakdown, dict):
                    diff_str = str(problem.difficulty).capitalize()
                    if diff_str == "Easy":
                        breakdown = {"concept": 1, "implementation": 1}
                    elif diff_str == "Hard":
                        breakdown = {"concept": 4, "implementation": 4}
                    else:
                        breakdown = {"concept": 3, "implementation": 3}
                self._difficulty_breakdowns[problem.slug] = breakdown
                object.__setattr__(problem, "difficulty_breakdown", breakdown)

                # Extract test_cases contract
                tcases = data.get("test_cases", [])
                self._test_cases[problem.slug] = tcases
                object.__setattr__(problem, "test_cases", tcases)
                
                # Pre-construct list item representation for high-speed API responses
                item_dict = {
                    "slug": problem.slug,
                    "title": problem.title,
                    "difficulty": problem.difficulty,
                    "difficulty_breakdown": breakdown,
                    "category": problem.category,
                    "tags": problem.tags,
                    "time_limit": problem.time_limit,
                    "memory_limit_mb": problem.memory_limit_mb,
                    "is_featured": problem.is_featured,
                }
                self._list_items_cache.append(item_dict)

                cat_key = problem.category.lower()
                if cat_key not in self._by_category_cache:
                    self._by_category_cache[cat_key] = []
                self._by_category_cache[cat_key].append(item_dict)

                if problem.is_featured:
                    self._featured_cache.append(item_dict)

            except Exception as e:
                print(f"[Warning] Failed to load problem from {path}: {e}")

    def get_list_items(self, category: Optional[str] = None) -> List[Dict[str, Any]]:
        if not category or category.strip() in ("", "全部"):
            return self._list_items_cache
        cat_lower = category.lower().strip()
        if cat_lower in ("featured", "常见手撕", "高频手撕", "🔥 常见手撕"):
            return self._featured_cache
        return self._by_category_cache.get(cat_lower, [])

    def get(self, slug: str) -> Optional[ProblemMeta]:
        return self._cache.get(slug)

    def get_difficulty_breakdown(self, slug: str) -> Optional[Dict[str, int]]:
        return self._difficulty_breakdowns.get(slug)

    def get_test_cases(self, slug: str) -> List[Dict[str, Any]]:
        return self._test_cases.get(slug, [])

    def list(self) -> List[ProblemMeta]:
        return list(self._cache.values())

    def save(self, problem: ProblemMeta):
        self.problems_dir.mkdir(parents=True, exist_ok=True)
        file_path = self.problems_dir / f"{problem.slug}.json"
        data = problem.model_dump()
        if hasattr(problem, "difficulty_breakdown") and problem.difficulty_breakdown:
            data["difficulty_breakdown"] = problem.difficulty_breakdown
        if hasattr(problem, "test_cases") and problem.test_cases:
            data["test_cases"] = problem.test_cases
        file_path.write_text(
            json.dumps(data, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        self._cache[problem.slug] = problem
        if hasattr(problem, "difficulty_breakdown"):
            self._difficulty_breakdowns[problem.slug] = problem.difficulty_breakdown
        if hasattr(problem, "test_cases"):
            self._test_cases[problem.slug] = problem.test_cases


repo = ProblemRepository()