#!/usr/bin/env python3
"""
MasterMind CLI

This script will:

1. Load the list of questions from `Master Mind/questions.json`.
2. Load the system prompt from a user-provided file (default: `system_prompt.txt`).
3. Load the proposal and IDIOT method documents.
4. For each question:
   - Send a request to a local AI model.
   - Store the response in a file inside a project-named folder.
   - Log the activity to the console.

A very small helper is used to talk to a local model; it can be replaced with a real
LLM client (Llama.cpp, Ollama, FastLLM, etc.) by editing the ``fetch_response`` function.

Usage:
    mastermind_cli.py \\
      --project ProjectName \\
      --proposal path/to/ProjectProposal.txt \\
      --idiot path/to/IDIOTMethod.txt \\
      [--prompt-file path/to/system_prompt.txt] \\
      [--model "llama3"]              # optional flag for the local model name
      [--ollama-url "http://..."]     # optional Ollama endpoint (default: OLLAMA_URL env)
      [--questions-file path.json]    # optional questions file (default: questions.json)
      [--output-dir path]             # optional output directory
      [--api-key "secret"]            # optional API key for Ollama auth
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import pathlib
import re
import sys
import urllib.request
from typing import List

logging.basicConfig(level=logging.INFO, format="%(message)s")

# Sliding window size for accumulated context (issue #8)
MAX_CONTEXT_WINDOW = 5


def validate_project_name(name: str) -> str:
    """Validate project name contains only safe characters. Prevents path traversal."""
    if not re.match(r"^[a-zA-Z0-9_-]+$", name):
        logging.error(
            "Invalid project name '%s'. Only alphanumeric, hyphens, and underscores allowed.",
            name,
        )
        sys.exit(1)
    return name


def validate_no_traversal(path: pathlib.Path, label: str) -> pathlib.Path:
    """Ensure resolved path does not escape its intended base directory."""
    resolved = path.resolve()
    base = pathlib.Path(".").resolve()
    try:
        resolved.relative_to(base)
    except ValueError:
        logging.error(
            "Invalid %s '%s': path escapes the allowed base directory.",
            label,
            path,
        )
        sys.exit(1)
    return resolved


def sanitize_phase_name(name: str) -> str:
    """Reduce a phase name to a safe filename slug. Prevents path traversal."""
    slug = re.sub(r"[^a-z0-9_-]+", "_", name.strip().lower()).strip("_")
    return slug or "phase"


def iter_question_texts(questions: List[dict]) -> List[str]:
    """Collect question texts recursively, including nested subQuestions."""
    texts: List[str] = []
    for q in questions or []:
        text = (q.get("text") or "").strip()
        if text:
            texts.append(text)
        texts.extend(iter_question_texts(q.get("subQuestions") or []))
    return texts


def read_file(path: pathlib.Path) -> str:
    """Return the contents of a text file."""
    try:
        return path.read_text(encoding="utf-8")
    except Exception as exc:
        logging.error("Could not read %s: %s", path, exc)
        sys.exit(1)


def load_questions(json_path: pathlib.Path) -> List[dict]:
    """Load questions from the JSON file."""
    text = read_file(json_path)
    try:
        return json.loads(text)
    except json.JSONDecodeError as exc:
        logging.error("Invalid JSON in %s: %s", json_path, exc)
        sys.exit(1)


def fetch_response(
    system_prompt: str,
    user_prompt: str,
    context: List[str],
    model: str = "gpt-oss:20b",
    ollama_url: str = "http://localhost:11434/api/generate",
    api_key: str | None = None,
    timeout: int = 120,
) -> str:
    """
    Send a request to a local AI model and return the response.

    This is a placeholder implementation.  Replace the body of this function
    with actual calls to your local model (e.g., llama.cpp, Ollama, FastLLM, etc.).

    Parameters
    ----------
    system_prompt : str
        The system-level instruction to the AI.
    user_prompt : str
        The question or user message.
    context : List[str]
        Optional additional context to prepend to the request.
    model : str
        The model name to use.
    ollama_url : str
        The Ollama API endpoint URL.
    api_key : str | None
        Optional API key for authentication.
    timeout : int
        Request timeout in seconds.

    Returns
    -------
    str
        The AI's reply.
    """
    prompt_parts = [system_prompt]
    prompt_parts.extend(f"{ctx}" for ctx in context)
    prompt_parts.append(f"Question: {user_prompt}")
    prompt = "\n\n".join(prompt_parts) + "\n"

    payload = json.dumps({
        "model": model,
        "prompt": prompt,
        "stream": False,
    }).encode("utf-8")
    headers = {"Content-Type": "application/json"}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"

    req = urllib.request.Request(
        ollama_url, data=payload, headers=headers, method="POST"
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            data = resp.read().decode("utf-8")
            result = json.loads(data)
            return result.get("response", "").strip()
    except Exception as exc:
        logging.error("Ollama request failed: %s", exc)
        sys.exit(1)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Send MasterMind questions to a local AI and log responses."
    )
    parser.add_argument(
        "--project",
        required=True,
        help="Name of the project – creates a folder with this name to store outputs.",
    )
    parser.add_argument(
        "--proposal",
        required=True,
        help="Path to the Project Proposal document (text file).",
    )
    parser.add_argument(
        "--idiot",
        required=True,
        help="Path to the IDIOT Method document (text file).",
    )
    parser.add_argument(
        "--prompt-file",
        default="system_prompt.txt",
        help="Path to the file that contains the system prompt.",
    )
    parser.add_argument(
        "--model",
        default="gpt-oss:20b",
        choices=[
            "gpt-oss:20b",
            "qwen3:30b",
            "devstral:24b",
            "llama3.3:70b",
        ],
        help="LLM model to use.",
    )
    parser.add_argument(
        "--output-dir",
        default="",
        help="Custom output directory for the project folder (default: current dir).",
    )
    parser.add_argument(
        "--ollama-url",
        default=None,
        help="Ollama API endpoint URL (default: $OLLAMA_URL or http://localhost:11434/api/generate).",
    )
    parser.add_argument(
        "--api-key",
        default=None,
        help="API key for Ollama authentication (default: $OLLAMA_API_KEY).",
    )
    parser.add_argument(
        "--questions-file",
        default=None,
        help="Path to questions JSON file (default: questions.json next to script).",
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=120,
        help="Request timeout in seconds for Ollama calls (default: 120).",
    )

    args = parser.parse_args()

    # Resolve Ollama URL: CLI arg > env var > default
    ollama_url = args.ollama_url or os.environ.get(
        "OLLAMA_URL", "http://localhost:11434/api/generate"
    )
    api_key = args.api_key or os.environ.get("OLLAMA_API_KEY")

    # Warn if using HTTP (cleartext)
    if ollama_url.startswith("http://"):
        logging.warning(
            "Using unencrypted HTTP for Ollama endpoint. "
            "Sensitive project data will be transmitted in cleartext. "
            "Use https:// or ensure network isolation."
        )

    # Validate project name against path traversal
    validate_project_name(args.project)

    base_dir = pathlib.Path.cwd()

    # Resolve project output directory with traversal protection
    if args.output_dir:
        out_base = pathlib.Path(args.output_dir).expanduser().resolve()
        validate_no_traversal(out_base, "--output-dir")
    else:
        out_base = base_dir
    project_dir = out_base / args.project

    try:
        project_dir.mkdir(parents=True, exist_ok=True)
    except Exception as exc:
        logging.error("Could not create project directory %s: %s", project_dir, exc)
        sys.exit(1)

    # Load documents
    system_prompt_path = pathlib.Path(args.prompt_file)
    proposal_path = pathlib.Path(args.proposal)
    idiot_path = pathlib.Path(args.idiot)

    project_prompt = read_file(system_prompt_path)
    proposal_text = read_file(proposal_path)
    idiot_text = read_file(idiot_path)

    # Resolve questions file: CLI arg > script-relative > project root
    if args.questions_file:
        questions_path = pathlib.Path(args.questions_file)
    else:
        script_dir = pathlib.Path(__file__).resolve().parent
        questions_path = script_dir.parent / "questions.json"
    questions = load_questions(questions_path)

    # Merge context into system prompt for easier reuse
    full_system_prompt = (
        f"{project_prompt}\n\n"
        f"--- PROPOSAL ---\n{proposal_text}\n\n"
        f"--- IDIOT METHOD ---\n{idiot_text}\n\n"
    )

    logging.info("Project folder: %s", project_dir)
    logging.info("Processing %d phases of questions.", len(questions))

    accumulated_summaries: List[str] = []
    experiment_results: List[dict] = []
    counter = 1

    for phase in questions:
        phase_name = phase.get("phase", "Unnamed Phase")
        phase_texts = iter_question_texts(phase.get("questions", []))
        logging.info("Phase: %s – %d questions", phase_name, len(phase_texts))

        # Sliding window: only keep last N summaries to bound context growth
        phase_context: List[str] = accumulated_summaries[-MAX_CONTEXT_WINDOW:]

        for question_text in phase_texts:
            logging.info("Q%d: %s", counter, question_text)
            response = fetch_response(
                full_system_prompt,
                question_text,
                phase_context,
                args.model,
                ollama_url,
                api_key,
                args.timeout,
            )
            output_file = project_dir / f"question_{counter:03d}.txt"
            try:
                output_file.write_text(response, encoding="utf-8")
            except Exception as exc:
                logging.error("Could not write %s: %s", output_file, exc)
                continue
            logging.info("Saved response to %s", output_file)

            accumulated_summaries.append(response)
            counter += 1

        # After finishing the phase, write a short summary file for this phase
        phase_summary_file = project_dir / f"{sanitize_phase_name(phase_name)}_summary.txt"
        try:
            summary = fetch_response(
                "Summarize the following input concisely, in at most 500 tokens",
                "\n\n".join(accumulated_summaries),
                [],
                args.model,
                ollama_url,
                api_key,
                args.timeout,
            )

            phase_summary_file.write_text(summary, encoding="utf-8")
            logging.info("Saved per-phase summary to %s", phase_summary_file)
        except SystemExit:
            raise
        except Exception as exc:
            logging.error("Could not write %s: %s", phase_summary_file, exc)

        # Record result for every phase (was incorrectly inside except block)
        experiment_results.append({
            "phase": phase_name,
            "summary_file": str(phase_summary_file),
        })

    logging.info("All done! Total responses: %d", counter - 1)

    results_file = project_dir / "experiment_results.json"
    try:
        results_file.write_text(
            json.dumps(experiment_results, indent=2), encoding="utf-8"
        )
        logging.info("Saved experiment results to %s", results_file)
    except Exception as exc:
        logging.error("Could not write experiment results: %s", exc)


if __name__ == "__main__":
    main()