"""Core solve logic shared between CLI and MCP server."""

import re
import time

from gpai_cli.client import GPAIClient

D3_VISUAL_RE = re.compile(r'<d3-visual[^>]*data-visual-id="([^"]*)"[^>]*>.*?</d3-visual>')


def clean_content(content: str, visuals: list[dict]) -> str:
    """Replace <d3-visual> tags with image URLs from visuals list."""
    visual_map = {v["id"]: v for v in visuals}

    def replace_visual(m: re.Match) -> str:
        visual_id = m.group(1)
        v = visual_map.get(visual_id)
        if v and v.get("imageUrl"):
            return f'[Visual: {v["imageUrl"]}]'
        return ""

    return D3_VISUAL_RE.sub(replace_visual, content)


def run_solve(client: GPAIClient, image_path: str) -> dict:
    """Run the full solve pipeline: upload → create task → wait → get solution.

    Returns a dict with keys: answer, solution, visuals, image_url, task_id, problem_id
    """
    # Upload
    file_id = client.upload_file(image_path)

    # Create task
    task_resp = client.create_task(file_id)
    task_id = task_resp["task"]["id"]
    problems = task_resp["task"].get("problems", [])
    problem_id = problems[0]["id"] if problems else None

    if not problem_id:
        raise ValueError("No problems found in task response")

    # Wait for task completion
    client.wait_for_completion(task_id)

    # Get solution
    resp = client.get_problem_solution(task_id, problem_id)

    # Wait for visuals (max 60s)
    resp = _wait_for_visuals(client, task_id, problem_id, resp)

    # Parse response
    problem = resp.get("problem", {})
    solution = problem.get("solution", {})
    if isinstance(solution, list):
        solution = solution[0] if solution else {}

    visuals = problem.get("visuals", [])
    raw_content = solution.get("content", "")
    content = clean_content(raw_content, visuals)

    visual_urls = [v["imageUrl"] for v in visuals if v.get("imageUrl")]

    return {
        "answer": solution.get("answer", ""),
        "solution": content,
        "visuals": visual_urls,
        "image_url": problem.get("cdnUrl", ""),
        "task_id": task_id,
        "problem_id": problem_id,
        "raw_response": resp,
    }


def _wait_for_visuals(
    client: GPAIClient, task_id: str, problem_id: str, resp: dict, timeout: float = 60.0
) -> dict:
    """Poll until all visuals are done."""
    visuals = resp.get("problem", {}).get("visuals", [])
    if not any(v.get("status") == "processing" for v in visuals):
        return resp

    deadline = time.time() + timeout
    while time.time() < deadline:
        time.sleep(2)
        try:
            resp = client.get_problem_solution(task_id, problem_id)
            visuals = resp.get("problem", {}).get("visuals", [])
            if not any(v.get("status") == "processing" for v in visuals):
                return resp
        except Exception:
            return resp
    return resp
