from datetime import date
from functools import lru_cache

import dspy
from dspy.signatures import InputField, OutputField, Signature

from app import Config

from .test_query import PossibleTest, TestQuery


class SearchQueryResponse(Signature):
    """Parse bloodwork search queries into structured test selections and date ranges.

    Instructions:
    1. Map user query to specific tests from 'possible_tests'.
    2. Extract date ranges using 'current_date' for relative terms (e.g., 'last year').
    3. If no specific date range is mentioned, leave start and end empty.
    4. Set 'out_of_reference_only' to True ONLY if the user specifically asks for abnormal, high, low, or out-of-range values.
    """

    query: str = InputField(description="User's natural language search query")
    current_date: str = InputField(description="Today's date (ISO format) for relative time calculations")
    possible_tests: list[PossibleTest] = InputField(description="List of valid tests to choose from")
    
    test_queries: list[TestQuery] = OutputField(description="List of identified test requirements with date ranges")
    out_of_reference_only: bool = OutputField(description="Whether to filter for out-of-reference results only")


class DSPYService:
    """Service class that wraps the DSPY model for parsing search queries."""

    def __init__(self):
        dspy.configure_cache(enable_memory_cache=True, enable_disk_cache=True)

        self.lm = dspy.LM(
            "groq/llama-3.3-70b-versatile",
            api_key=Config.GROQ_API_KEY,
            base_url="https://api.groq.com/openai/v1",
            cache=True,
            num_retries=4,
            temperature=0.4,
            max_tokens=2048,
        )
        dspy.configure(lm=self.lm)
        self._initialized = True
        self.model = dspy.Predict(SearchQueryResponse)
        self._parse_search_query_cached = lru_cache(maxsize=256)(
            self._parse_search_query_uncached
        )

    def parse_search_query(
        self, query: str, possible_tests: list[PossibleTest]
    ) -> dict:
        """Parse a natural-language search query into structured parameters."""

        normalized_query = query.strip()
        current_date = date.today().isoformat()
        possible_tests_key = self._possible_tests_cache_key(possible_tests)

        return self._parse_search_query_cached(
            normalized_query, current_date, possible_tests_key
        )

    def _parse_search_query_uncached(
        self,
        query: str,
        current_date: str,
        possible_tests_key: tuple[tuple[str, str, tuple[str, ...]], ...],
    ) -> dict:
        possible_tests = [
            PossibleTest(
                test_id=test_id, test_name=test_name, subtest_names=list(subtest_names)
            )
            for test_id, test_name, subtest_names in possible_tests_key
        ]

        response = self.model(
            query=query, possible_tests=possible_tests, current_date=current_date
        ).toDict()

        return response

    @staticmethod
    def _possible_tests_cache_key(
        possible_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 possible_tests
        )
