import json
import re
from typing import Optional, Dict, Any
from pydantic import BaseModel, Field, validator
from huggingface_hub import InferenceClient
from huggingface_hub.errors import HfHubHTTPError
from variables import *

class LLMResponse(BaseModel):
    initial_prompt_evaluation: str = Field(..., description="Evaluation of the initial prompt")
    refined_prompt: str = Field(..., description="The refined version of the prompt")
    explanation_of_refinements: str = Field(..., description="Explanation of the refinements made")
    response_content: Optional[Dict[str, Any]] = Field(None, description="Raw response content")

    @validator('initial_prompt_evaluation', 'refined_prompt', 'explanation_of_refinements')
    def clean_text_fields(cls, v):
        if isinstance(v, str):
            return v.strip().replace('\\n', '\n').replace('\\"', '"')
        return v

class PromptRefiner:
    def __init__(self, api_token: str, meta_prompts):
        self.client = InferenceClient(token=api_token, timeout=120)
        self.meta_prompts = meta_prompts
        
    def refine_prompt(self, prompt: str, meta_prompt_choice: str) -> tuple:
        try:
            selected_meta_prompt = self.meta_prompts.get(
                meta_prompt_choice, 
                self.meta_prompts["star"]
            )
            
            messages = [
                {
                    "role": "system", 
                    "content": 'You are an expert at refining and extending prompts. Given a basic prompt, provide a more relevant and detailed prompt.'
                },
                {
                    "role": "user", 
                    "content": selected_meta_prompt.replace("[Insert initial prompt here]", prompt)
                }
            ]
            
            response = self.client.chat_completion(
                model=prompt_refiner_model,
                messages=messages,
                max_tokens=3000,
                temperature=0.8
            )
            
            response_content = response.choices[0].message.content.strip()
            result = self._parse_response(response_content)
            
            # Create and validate LLMResponse
            llm_response = LLMResponse(**result)
            
            return (
                llm_response.initial_prompt_evaluation,
                llm_response.refined_prompt,
                llm_response.explanation_of_refinements,
                llm_response.dict()
            )

        except HfHubHTTPError as e:
            return self._create_error_response("Model timeout. Please try again later.")
        except Exception as e:
            return self._create_error_response(f"Unexpected error: {str(e)}")

    def _create_error_response(self, error_message: str) -> tuple:
        error_response = LLMResponse(
            initial_prompt_evaluation=f"Error: {error_message}",
            refined_prompt="The selected model is currently unavailable.",
            explanation_of_refinements="An error occurred during processing.",
            response_content={"error": error_message}
        )
        return (
            error_response.initial_prompt_evaluation,
            error_response.refined_prompt,
            error_response.explanation_of_refinements,
            error_response.dict()
        )

    def _parse_response(self, response_content: str) -> dict:
        try:
            # First attempt: Try to extract JSON from <json> tags
            json_match = re.search(r'<json>\s*(.*?)\s*</json>', response_content, re.DOTALL)
            if json_match:
                json_str = json_match.group(1)
                json_str = re.sub(r'\n\s*', ' ', json_str)
                json_str = json_str.replace('"', '\\"')
                json_output = json.loads(f'"{json_str}"')
                
                if isinstance(json_output, str):
                    json_output = json.loads(json_output)
                
                return {
                    "initial_prompt_evaluation": json_output.get("initial_prompt_evaluation", ""),
                    "refined_prompt": json_output.get("refined_prompt", ""),
                    "explanation_of_refinements": json_output.get("explanation_of_refinements", ""),
                    "response_content": json_output
                }
            
            # Second attempt: Try to extract fields using regex
            output = {}
            for key in ["initial_prompt_evaluation", "refined_prompt", "explanation_of_refinements"]:
                pattern = rf'"{key}":\s*"(.*?)"(?:,|\}})'
                match = re.search(pattern, response_content, re.DOTALL)
                output[key] = match.group(1) if match else ""
            
            output["response_content"] = response_content
            return output

        except (json.JSONDecodeError, ValueError) as e:
            print(f"Error parsing response: {e}")
            print(f"Raw content: {response_content}")
            return {
                "initial_prompt_evaluation": "Error parsing response",
                "refined_prompt": "",
                "explanation_of_refinements": str(e),
                "response_content": str(e)
            }

    def apply_prompt(self, prompt: str, model: str) -> str:
        try:
            messages = [
                {
                    "role": "system",
                    "content": """You are a markdown formatting expert. Format your responses with proper spacing and structure following these rules:
                        
                        1. Paragraph Spacing:
                        - Add TWO blank lines between major sections (##)
                        - Add ONE blank line between subsections (###)
                        - Add ONE blank line between paragraphs within sections
                        - Add ONE blank line before and after lists
                        - Add ONE blank line before and after code blocks
                        - Add ONE blank line before and after blockquotes
                        
                        2. Section Formatting:
                        # Title
                        
                        ## Major Section
                        
                        [blank line]
                        Content paragraph 1
                        [blank line]
                        Content paragraph 2
                        [blank line]"""
                },
                {
                    "role": "user",
                    "content": prompt
                }
            ]
    
            response = self.client.chat_completion(
                model=model,
                messages=messages,
                max_tokens=3000,
                temperature=0.8,
                stream=True
            )
            
            full_response = ""
            for chunk in response:
                if chunk.choices[0].delta.content is not None:
                    full_response += chunk.choices[0].delta.content
                    
            return full_response.replace('\n\n', '\n').strip()
                
        except Exception as e:
            return f"Error: {str(e)}"