"""Natural-language API for bloodwork search."""

import logging
from datetime import UTC, datetime, timedelta
from functools import lru_cache
from typing import Iterable

from pydantic import BaseModel

from app.repositories import test_repository

from .dspy_api import DSPYService
from .test_query import PossibleTest, TestQuery

logger = logging.getLogger(__name__)

class InvalidSearchQuery(ValueError):
    """Raised when query text does not match the expected intent."""


class ParsedSearchQuery(BaseModel):
    """Normalized representation of a user bloodwork search query."""
    raw_query: str
    test_queries: list[TestQuery]
    out_of_reference_only: bool

dspy_service = DSPYService()

class MLSearch:
    """Search API that parses natural-language bloodwork requests."""
    def __init__(
        self,
        known_test_name_fetch_interval: timedelta = timedelta(minutes=5),
        auto_refresh_test_names: bool = True,
    ):
        self.auto_refresh_known_test_names = auto_refresh_test_names
        self._known_tests_by_patient: dict[str, list[PossibleTest]] = {}
        self._known_tests_refreshed_at_by_patient: dict[str, datetime] = {}

        self.refresh_known_test_names_every = known_test_name_fetch_interval
        self._search_cached = lru_cache(maxsize=256)(self._search_uncached)


    def _fetch_known_test_names(self, patient_id: str, force: bool = False) -> list[PossibleTest]:
        """Refresh known test records for a patient based on TTL."""
        if not self.auto_refresh_known_test_names and not force:
            return self._known_tests_by_patient.get(patient_id, [])

        now = datetime.now(UTC)
        refreshed_at = self._known_tests_refreshed_at_by_patient.get(patient_id)
        if (
            not force
            and refreshed_at is not None
            and now - refreshed_at < self.refresh_known_test_names_every
        ):
            return self._known_tests_by_patient.get(patient_id, [])

        try:
            tests = test_repository.find_by_patient_id(patient_id)
            possible_tests: list[PossibleTest] = []
            for test in tests:
                possible_tests.append(
                    PossibleTest(
                        test_id=str(test.id),
                        test_name=test.name,
                        subtest_names=[subtest.name for subtest in test.subtests if subtest.name],
                    )
                )

            possible_tests = self._dedupe_possible_tests(possible_tests)

            self._known_tests_by_patient[patient_id] = possible_tests
            self._known_tests_refreshed_at_by_patient[patient_id] = now
            return possible_tests
        except Exception as exc:
            logger.warning("Failed to refresh known test names for patient %s: %s", patient_id, exc)
            return self._known_tests_by_patient.get(patient_id, [])

    def search(self, query: str, patient_id: str) -> ParsedSearchQuery:
        known_tests = self._fetch_known_test_names(patient_id)

        return self._search_cached(
            query.strip(),
            self._known_tests_cache_key(known_tests),
        )

    def _search_uncached(
        self,
        query: str,
        known_tests_key: tuple[tuple[str, str, tuple[str, ...]], ...],
    ) -> ParsedSearchQuery:
        known_tests = [
            PossibleTest(test_id=test_id, test_name=test_name, subtest_names=list(subtest_names))
            for test_id, test_name, subtest_names in known_tests_key
        ]

        result = dspy_service.parse_search_query(
            query=query,
            possible_tests=known_tests,
        )

        # Support both new and legacy model keys while transition is in progress.
        test_queries = result.get("test_queries") or []

        if not test_queries:
            return ParsedSearchQuery(
                raw_query=query,
                test_queries=[],
                out_of_reference_only=bool(result.get("out_of_reference_only", False)),
            )

        test_queries = [TestQuery(**dr) for dr in test_queries]
        if not any(tq.test_parameter.strip() for tq in test_queries):
            raise InvalidSearchQuery("No test names found in query. Please specify at least one test to search for.")

        test_queries = self._dedupe_test_queries(test_queries)
        
        return ParsedSearchQuery(
            raw_query=query,
            test_queries=test_queries,
            out_of_reference_only=bool(result.get("out_of_reference_only", False)),
        )

    @staticmethod
    def _known_tests_cache_key(
        known_tests: list[PossibleTest],
    ) -> tuple[tuple[str, str, tuple[str, ...]], ...]:
        return tuple(
            (
                test.test_id.strip(),
                test.test_name.strip(),
                tuple(subtest.strip() for subtest in test.subtest_names if subtest.strip()),
            )
            for test in known_tests
        )

    @staticmethod
    def _dedupe(values: Iterable[str]) -> list[str]:
        """
        Deduplicate while preserving order,
        ignoring case and surrounding whitespace.
        """
        seen: set[str] = set()
        ordered: list[str] = []
        for value in values:
            normalized = value.strip()
            if not normalized:
                continue
            key = normalized.casefold()
            if key in seen:
                continue
            seen.add(key)
            ordered.append(normalized)
        return ordered

    @staticmethod
    def _dedupe_possible_tests(values: list[PossibleTest]) -> list[PossibleTest]:
        seen: set[str] = set()
        ordered: list[PossibleTest] = []
        for value in values:
            key = value.test_id.strip()
            if not key or key in seen:
                continue
            seen.add(key)
            ordered.append(value)
        return ordered

    @staticmethod
    def _dedupe_test_queries(values: list[TestQuery]) -> list[TestQuery]:
        seen: set[tuple[str, str, str]] = set()
        ordered: list[TestQuery] = []

        for value in values:
            key = (
                value.test_parameter.strip().casefold(),
                value.start.isoformat() if value.start else None,
                value.end.isoformat() if value.end else None,
            )
            if key in seen:
                continue

            seen.add(key)
            ordered.append(value)

        return ordered