import json
import os
import asyncio
from utils.logger import setup_logger

logger = setup_logger("reasoning_loop")

class ReasoningLoop:
    def __init__(self, llm_service, threshold=6.5, max_retries=2, log_path="memory/reasoning_log.json"):
        self.llm = llm_service
        self.threshold = threshold
        self.max_retries = max_retries
        self.log_path = log_path

    def _log_event(self, event: dict):
        try:
            logs = []
            if os.path.exists(self.log_path):
                with open(self.log_path, "r") as f:
                    logs = json.load(f)
            logs.append(event)
            # Maintain max 500 entries as per rules
            if len(logs) > 500:
                logs = logs[-500:]
            with open(self.log_path, "w") as f:
                json.dump(logs, f, indent=2)
        except Exception as e:
            logger.error(f"Failed to log reasoning event: {e}")

    async def evaluate(self, user_input: str, ai_response: str) -> dict:
        prompt = f"""
You are a response quality critic for an AI assistant.
Evaluate the response below against the user's question.

User Question: {user_input}
AI Response: {ai_response}

Rate on these dimensions (1-10):
- Accuracy: Is it factually correct?
- Completeness: Does it fully answer the question?
- Relevance: Is it on-topic?
- Tone: Is it appropriate and natural?

Return ONLY this JSON:
{{
  "accuracy": <score>,
  "completeness": <score>,
  "relevance": <score>,
  "tone": <score>,
  "overall": <average>,
  "needs_correction": <true/false>,
  "correction_hint": "<what to fix>"
}}
"""
        try:
            response_text, _ = await self.llm.generate_response(prompt)
            import re
            json_match = re.search(r'\{.*\}', response_text, re.DOTALL)
            if json_match:
                return json.loads(json_match.group())
        except Exception as e:
            logger.error(f"Error in reasoning evaluation: {e}")
        return {"needs_correction": False, "overall": 10.0}

    async def correct(self, user_input: str, ai_response: str, hint: str) -> str:
        prompt = f"""
Your previous response had issues: {hint}
Rewrite the response to fix these issues.
Be concise, accurate, and natural.
Original Question: {user_input}
Improved Response:
"""
        try:
            new_response, _ = await self.llm.generate_response(prompt)
            return new_response.strip()
        except Exception as e:
            logger.error(f"Error in reasoning correction: {e}")
            return ai_response

    async def run(self, user_input: str, initial_response: str) -> str:
        current_response = initial_response
        for i in range(self.max_retries):
            evaluation = await self.evaluate(user_input, current_response)
            
            self._log_event({
                "iteration": i + 1,
                "input": user_input,
                "response": current_response,
                "evaluation": evaluation
            })

            if not evaluation.get("needs_correction") or evaluation.get("overall", 10.0) >= self.threshold:
                break
            
            logger.info(f"Self-correction triggered (Attempt {i+1}). Hint: {evaluation.get('correction_hint')}")
            current_response = await self.correct(user_input, current_response, evaluation.get("correction_hint"))
            
        return current_response