import time
from pathlib import Path

from openai import OpenAI
from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from rich.markdown import Markdown
from dotenv import load_dotenv
load_dotenv()
console = Console()
client = OpenAI()

# OpenAI pricing per 1K tokens (as of 2024, update as needed)
PRICING = {
    "gpt-4o-mini": {"input": 0.000150, "output": 0.000600},
    "gpt-4o":      {"input": 0.002500, "output": 0.010000},
}

SYSTEM_PROMPT = """You are Direction, a project-aware coding assistant.
You have been given relevant code chunks retrieved from the project's vector database.
Use these chunks as your primary context to answer the user's question accurately.

Rules:
- Reference specific files and line context when relevant
- If the retrieved context is insufficient, say so clearly
- Do not hallucinate code that isn't in the context
- Be concise but complete
- If suggesting changes, show the exact code to change"""


# ---------------------------------------------------------------------------
# Embedding
# ---------------------------------------------------------------------------

def embed_query(question: str, model: str) -> list[float]:
    """Embed the user's question for similarity search."""
    response = client.embeddings.create(
        input=question,
        model=model
    )
    return response.data[0].embedding


# ---------------------------------------------------------------------------
# Context assembly
# ---------------------------------------------------------------------------

def assemble_context(documents: list[str], metadatas: list[dict], distances: list[float]) -> str:
    """
    Build a structured context string from retrieved chunks.
    Groups chunks by file for clarity.
    """
    from direction.db import distance_to_score

    file_chunks: dict[str, list[dict]] = {}

    for doc, meta, dist in zip(documents, metadatas, distances):
        file = meta.get("file", "unknown")
        score = distance_to_score(dist)
        if file not in file_chunks:
            file_chunks[file] = []
        file_chunks[file].append({
            "content": doc,
            "chunk_index": meta.get("chunk_index", 0),
            "score": score
        })

    parts = []
    for file, chunks in file_chunks.items():
        chunks.sort(key=lambda x: x["chunk_index"])
        part = f"### {file}\n"
        for chunk in chunks:
            part += f"// relevance: {chunk['score']}\n"
            part += chunk["content"]
            part += "\n"
        parts.append(part)

    return "\n---\n".join(parts)


def estimate_full_context_tokens(config: dict) -> int:
    """
    Estimate tokens if we had loaded entire project into context.
    Used for token savings calculation.
    """
    from direction.db import get_collection, list_indexed_files
    import tiktoken

    enc = tiktoken.encoding_for_model("gpt-4o")
    collection = get_collection()
    files = list_indexed_files(collection)

    total = 0
    for file_path in files:
        path = Path(file_path)
        if path.exists():
            try:
                text = path.read_text(encoding="utf-8", errors="ignore")
                total += len(enc.encode(text))
            except Exception:
                pass

    return total


# ---------------------------------------------------------------------------
# Token counting
# ---------------------------------------------------------------------------

def count_tokens(text: str, model: str) -> int:
    """Count tokens in a string for a given model."""
    import tiktoken
    try:
        enc = tiktoken.encoding_for_model(model)
    except KeyError:
        enc = tiktoken.encoding_for_model("gpt-4o")
    return len(enc.encode(text))


def calculate_cost(input_tokens: int, output_tokens: int, model: str) -> float:
    """Calculate cost in USD for a query."""
    pricing = PRICING.get(model, PRICING["gpt-4o-mini"])
    input_cost = (input_tokens / 1000) * pricing["input"]
    output_cost = (output_tokens / 1000) * pricing["output"]
    return round(input_cost + output_cost, 6)


# ---------------------------------------------------------------------------
# LLM call with fallback
# ---------------------------------------------------------------------------

def call_llm(
    question: str,
    context: str,
    model: str,
    fallback_model: str
) -> dict:
    """
    Call OpenAI with assembled context.
    Falls back to stronger model if response seems insufficient.
    Returns response dict with content, tokens, model used.
    """
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"## Retrieved Context\n\n{context}\n\n## Question\n\n{question}"}
    ]

    response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=0.2  # low temp for code accuracy
    )

    answer = response.choices[0].message.content
    usage = response.usage

    # Simple fallback heuristic: if answer is very short and contains uncertainty phrases
    uncertainty_phrases = [
        "i don't have enough", "insufficient context",
        "not enough information", "cannot determine"
    ]
    needs_fallback = (
        len(answer.split()) < 30 and
        any(phrase in answer.lower() for phrase in uncertainty_phrases) and
        model != fallback_model
    )

    if needs_fallback:
        console.print(f"[yellow]Falling back to {fallback_model}...[/yellow]")
        response = client.chat.completions.create(
            model=fallback_model,
            messages=messages,
            temperature=0.2
        )
        answer = response.choices[0].message.content
        usage = response.usage
        model = fallback_model

    return {
        "answer": answer,
        "model": model,
        "input_tokens": usage.prompt_tokens,
        "output_tokens": usage.completion_tokens,
        "total_tokens": usage.total_tokens
    }


# ---------------------------------------------------------------------------
# Main query runner
# ---------------------------------------------------------------------------

def run_query(
    question: str,
    top_k: int,
    model: str,
    show_sources: bool,
    config: dict
):
    """
    Full query pipeline:
    1. Embed question
    2. Retrieve chunks
    3. Assemble context
    4. Call LLM
    5. Display result
    6. Log analytics
    """
    from direction.db import get_collection, query_chunks, distance_to_score
    from direction.analytics import log_query

    start_time = time.time()

    # Step 1: Embed question
    with console.status("Embedding query..."):
        query_embedding = embed_query(question, config["embedding_model"])

    # Step 2: Retrieve chunks
    with console.status("Retrieving context..."):
        collection = get_collection()
        results = query_chunks(collection, query_embedding, top_k=top_k)

    if not results["documents"]:
        console.print("[red]No relevant context found. Try running `direction update` first.[/red]")
        return

    # Step 3: Assemble context
    context = assemble_context(
        results["documents"],
        results["metadatas"],
        results["distances"]
    )

    context_tokens = count_tokens(context, model)

    # Step 4: Call LLM
    with console.status(f"Calling {model}..."):
        llm_result = call_llm(
            question=question,
            context=context,
            model=model,
            fallback_model=config["fallback_model"]
        )

    elapsed = round(time.time() - start_time, 2)

    # Step 5: Display result
    console.print(Panel(
        Markdown(llm_result["answer"]),
        title=f"[bold green]Direction[/bold green] [dim]({llm_result['model']})[/dim]",
        border_style="green"
    ))

    # Show sources if requested
    if show_sources:
        console.print("\n[bold]Retrieved Sources:[/bold]")
        source_table = Table(show_header=True)
        source_table.add_column("File", style="cyan")
        source_table.add_column("Chunk", justify="right")
        source_table.add_column("Relevance", justify="right")

        for meta, dist in zip(results["metadatas"], results["distances"]):
            score = distance_to_score(dist)
            source_table.add_row(
                meta.get("file", "unknown"),
                str(meta.get("chunk_index", 0)),
                f"{score:.2%}"
            )
        console.print(source_table)

    # Show token analytics inline
    cost = calculate_cost(
        llm_result["input_tokens"],
        llm_result["output_tokens"],
        llm_result["model"]
    )

    # Estimate savings
    full_tokens = estimate_full_context_tokens(config)
    tokens_used = llm_result["input_tokens"]
    savings_pct = round((1 - tokens_used / full_tokens) * 100, 1) if full_tokens > 0 else 0

    console.print(
        f"\n[dim]Tokens: {llm_result['total_tokens']} | "
        f"Cost: ${cost:.4f} | "
        f"Time: {elapsed}s | "
        f"Saved ~{savings_pct}% vs full context[/dim]"
    )

    # Step 6: Log to analytics
    avg_score = sum(
        distance_to_score(d) for d in results["distances"]
    ) / len(results["distances"])

    log_query(
        question=question,
        model=llm_result["model"],
        input_tokens=llm_result["input_tokens"],
        output_tokens=llm_result["output_tokens"],
        total_tokens=llm_result["total_tokens"],
        cost=cost,
        response_time=elapsed,
        retrieval_score=avg_score,
        chunks_retrieved=len(results["documents"]),
        full_context_tokens=full_tokens,
        tokens_saved=full_tokens - tokens_used
    )