import click
import llm
import json
import asyncio
from typing import List, Dict, Any, Optional
import sys
import pathlib
import sqlite_utils
import os

def user_dir() -> pathlib.Path:
    """Get or create user directory for storing application data."""
    llm_user_path = os.environ.get("LLM_USER_PATH")
    if llm_user_path:
        path = pathlib.Path(llm_user_path)
    else:
        path = pathlib.Path(click.get_app_dir("io.datasette.llm"))
    path.mkdir(exist_ok=True, parents=True)
    return path

def logs_db_path() -> pathlib.Path:
    """Get path to logs database."""
    return user_dir() / "logs.db"

class MultiModelComparator:
    def __init__(
        self,
        models: List[str],
        evaluator: Optional[str] = None,
        score_threshold: float = 0.7
    ):
        self.models = models
        self.evaluator = evaluator or "gpt-4"
        self.score_threshold = score_threshold

    async def compare(self, prompt: str) -> Dict[str, Any]:
        # Get responses from all models
        model_responses = await self._get_model_responses(prompt)
        
        # Have evaluator analyze and score responses
        evaluation = await self._evaluate_responses(prompt, model_responses)
        
        return {
            "prompt": prompt,
            "responses": model_responses,
            "evaluation": evaluation,
            "metadata": {
                "models": self.models,
                "evaluator": self.evaluator,
                "score_threshold": self.score_threshold
            }
        }

    async def _get_model_responses(self, prompt: str) -> List[Dict[str, Any]]:
        tasks = [self._get_single_response(model, prompt) for model in self.models]
        return await asyncio.gather(*tasks)

    async def _get_single_response(self, model: str, prompt: str) -> Dict[str, Any]:
        try:
            db = sqlite_utils.Database(logs_db_path())
            response = llm.get_model(model).prompt(prompt)
            response.log_to_db(db)
            
            return {
                "model": model,
                "response": response.text(),
                "response_json": response.response_json
            }
        except Exception as e:
            return {"model": model, "error": str(e)}

    async def _evaluate_responses(self, prompt: str, responses: List[Dict[str, Any]]) -> Dict[str, Any]:
        evaluator = llm.get_model(self.evaluator)
        
        evaluation_prompt = f"""Analyze and compare these model responses to the prompt:

Prompt: {prompt}

Responses:
{self._format_responses(responses)}

Please provide:
1. A detailed comparison of the responses
2. Score each response (0-1)
3. Highlight key agreements and disagreements
4. Synthesize a combined best answer

Format your response as JSON with these keys:
- comparison
- scores
- agreements
- disagreements 
- synthesis
"""

        db = sqlite_utils.Database(logs_db_path())
        response = evaluator.prompt(evaluation_prompt)
        response.log_to_db(db)

        try:
            return json.loads(response.text())
        except json.JSONDecodeError:
            return {
                "comparison": response.text(),
                "scores": {},
                "agreements": [],
                "disagreements": [],
                "synthesis": ""
            }

    def _format_responses(self, responses: List[Dict[str, Any]]) -> str:
        formatted = []
        for r in responses:
            if "error" in r:
                formatted.append(f"{r['model']}: Error - {r['error']}")
            else:
                formatted.append(f"{r['model']}: {r['response']}")
        return "\n\n".join(formatted)

@llm.hookimpl
def register_commands(cli):
    @cli.command()
    @click.argument("prompt")
    @click.option(
        "-m",
        "--models",
        multiple=True,
        help="Models to include in comparison (can specify multiple)",
        default=["gpt-4", "claude-3-sonnet", "gemini-pro"]
    )
    @click.option(
        "--evaluator",
        help="Model to use as evaluator",
        default="gpt-4"
    )
    @click.option(
        "--score-threshold",
        type=float,
        help="Minimum score threshold",
        default=0.7
    )
    @click.option(
        "--output",
        type=click.Path(dir_okay=False, writable=True),
        help="Save full results to this JSON file"
    )
    @click.option(
        "--stdin/--no-stdin",
        default=True,
        help="Read additional input from stdin"
    )
    def compare(
        prompt: str,
        models: List[str],
        evaluator: str,
        score_threshold: float,
        output: Optional[str],
        stdin: bool
    ):
        """Compare responses from multiple models and evaluate them.
        
        Reads additional input from stdin if provided and --stdin flag is enabled.
        """
        if stdin and not sys.stdin.isatty():
            stdin_content = sys.stdin.read().strip()
            prompt = f"{prompt}\n\n{stdin_content}"
        
        comparator = MultiModelComparator(
            models=list(models),
            evaluator=evaluator,
            score_threshold=score_threshold
        )
        
        try:
            result = asyncio.run(comparator.compare(prompt))
            
            if output:
                with open(output, 'w') as f:
                    json.dump(result, f, indent=2)
            
            # Print formatted results
            click.echo("\nModel Responses:\n")
            for response in result["responses"]:
                click.echo(f"\n{response['model']}:")
                click.echo(response.get('response', f"Error: {response.get('error')}"))
            
            click.echo("\nEvaluation:\n")
            eval_result = result["evaluation"]
            
            click.echo("Comparison:")
            click.echo(eval_result["comparison"])
            
            click.echo("\nScores:")
            for model, score in eval_result["scores"].items():
                click.echo(f"{model}: {score}")
            
            click.echo("\nKey Agreements:")
            for agreement in eval_result["agreements"]:
                click.echo(f"- {agreement}")
            
            click.echo("\nKey Disagreements:")
            for disagreement in eval_result["disagreements"]:
                click.echo(f"- {disagreement}")
            
            click.echo("\nSynthesized Answer:")
            click.echo(eval_result["synthesis"])
            
        except Exception as e:
            raise click.ClickException(str(e))

@llm.hookimpl
def register_models(register):
    pass

@llm.hookimpl
def register_prompts(register):
    pass