from groq import Groq
from utils.config import settings
import json
import re

class QueryDecomposerService:
    def __init__(self):
        self.provider = "groq"  # Default
        
        if settings.GROQ_API_KEY:
            self.provider = "groq"
            self.client = Groq(api_key=settings.GROQ_API_KEY)
            self.model_name = getattr(settings, 'GROQ_FAST_MODEL', settings.GROQ_MODEL)
            print(f"QueryDecomposerService initialized with Groq model: {self.model_name}")
        else:
            raise ValueError("GROQ_API_KEY is not set.")

    def _get_json_response(self, prompt: str) -> dict:
        """Robust method to get JSON response, handling API errors and Markdown"""
        import re
        import time
        
        for attempt in range(3):
            try:
                # Primary Strategy: Strict JSON Mode
                try:
                    chat_completion = self.client.chat.completions.create(
                        messages=[
                            {"role": "system", "content": "You are a helpful assistant that outputs JSON."},
                            {"role": "user", "content": prompt}
                        ],
                        model=self.model_name,
                        temperature=0,
                        response_format={"type": "json_object"}
                    )
                    content = chat_completion.choices[0].message.content
                    return json.loads(content.strip())
                except Exception as api_err:
                    # Fallback Strategy: Text Mode if JSON validation fails
                    if "400" in str(api_err) or "json_validate_failed" in str(api_err):
                        print(f"JSON Mode failed. Falling back to Text Mode (Attempt {attempt+1})...")
                        chat_completion = self.client.chat.completions.create(
                            messages=[
                                {"role": "system", "content": "You are a helpful assistant that outputs JSON."},
                                {"role": "user", "content": prompt + "\n\nOUTPUT RAW JSON ONLY. NO MARKDOWN."}
                            ],
                            model=self.model_name,
                            temperature=0
                        )
                        content = chat_completion.choices[0].message.content
                        
                        # Strip Markdown Code Blocks
                        match = re.search(r"```(?:json)?(.*?)```", content, re.DOTALL)
                        if match:
                            content = match.group(1)
                        
                        return json.loads(content.strip())
                    raise api_err
                    
            except Exception as e:
                print(f"Error in LLM request (Attempt {attempt+1}): {e}")
                if attempt < 2:
                    time.sleep(2)
                else:
                    return {}
        return {}

    def decompose(self, query: str) -> dict:
        prompt = f"""
You are a constitutional query classifier. Analyze the user's question and extract structured information.

User Query: "{query}"

Your task:
1. Identify the query TYPE (one of: simple_article, amendment_impact, multi_hop, temporal_evolution, cross_reference, comparison)
2. Extract ENTITIES: article numbers, amendment numbers, years mentioned
3. Identify INTENT: what user wants (explain, compare, list, trace, relate)
4. Estimate COMPLEXITY: simple/medium/complex

Return ONLY valid JSON, no other text:

{{
  "query_type": "amendment_impact|simple_article|multi_hop|temporal_evolution|cross_reference|comparison",
  "intent": "explain|compare|list|trace|relate",
  "entities": {{
    "articles": [list of article numbers as strings],
    "amendments": [list of amendment numbers],
    "years": [list of years]
  }},
  "complexity": "simple|medium|complex",
  "raw_query": "{query}"
}}
        """
        
        result = self._get_json_response(prompt)
        
        if not result:
            return {
                "query_type": "simple_article",
                "intent": "explain",
                "entities": {"articles": [], "amendments": [], "years": []},
                "complexity": "simple",
                "raw_query": query,
                "error": "Error in query decomposition"
            }
            
        return result