import json, re, io, contextlib, traceback
from typing import Dict, Any, Tuple
from functools import partial
import types

# ----- existing helper ---------------------------------------------------
def run_llm(system: str, messages, model_id: str = "gpt-4o") -> str:
    from AgentOccam.AgentOccam import (
        MODEL_FAMILIES,
        CALL_MODEL_WITH_MESSAGES_FUNCTION_MAP
    )
    model_family = [m for m in MODEL_FAMILIES if m in model_id][0]
    call_fn = partial(CALL_MODEL_WITH_MESSAGES_FUNCTION_MAP[model_family],
                      model_id=model_id)
    return call_fn(system_prompt=system, messages=messages)

# ----- util: extract python fenced block --------------------------------
_CODE_RE = re.compile(r"```python\s*(.*?)\s*```", re.S)

def _extract_code_block(text: str) -> str | None:
    m = _CODE_RE.search(text)
    return m.group(1) if m else None

def _exec_code_get_answer(code, extracted):
    exec_error, answer_code = None, None
    if code:
        sandbox = {"data": extracted}
        stdout = io.StringIO()
        try:
            with contextlib.redirect_stdout(stdout):
                exec(code, sandbox, sandbox)
            answer_code = sandbox.get("answer", None)
            if answer_code is None:
                answer_code = stdout.getvalue().strip() or "(no answer)"
        except Exception as exc:
            exec_error = traceback.format_exc().splitlines()[-1]
            answer_code = f"ERROR: {exec_error}"
    else:
        answer_code = "ERROR: code block missing"

    return answer_code

# ========================================================================
#  main entry point
# ========================================================================
def tri_phase_analyze(objective: str,
                      extracted: list[Dict[str, Any]],
                      model_id: str = "gpt-4o") -> Dict[str, Any]:
    """
    Returns   {"answer": <final text>, "source": "code"|"text"|"both", 
               "code_answer": ..., "text_answer": ...}
    """

    # ── 1) code LLM  ----------------------------------------------------
    system_code = (
        "You are an analysis assistant that MUST write Python code.\n"
        "You will be provided with objective and data samples (a small portion of all the data as a reference) for analysis as a reference.\n"
        "• The data is pre-loaded in a variable named `data`.\n"
        "• Assign your final answer to a variable named `answer`.\n"
        "Return only one fenced block:\n"
        "```python\n# code here\nanswer = ...\n```\n"
    )
    prompt_code = f"[ANALYSIS OBJECTIVE]\n{objective}\n\n[DATA]{extracted[:20]}"
    msg_code = [{"role": "user",
                 "content": [{"type": "text", "text": prompt_code}]}]
    raw_code = run_llm(system_code, msg_code, model_id)
    Py_code = _extract_code_block(raw_code)

    print("🧪 analyze code obtained")
    print(Py_code)

    answer_code = _exec_code_get_answer(Py_code, extracted)

    retry_limit = 3
    retry_count = 0
    while 'error' in str(answer_code).lower():
        msg_code.append({"role": "assistant",
                         "content": [{"type": "text", "text": raw_code}]})
        msg_code.append({"role": "user",
                         "content": [{"type": "text", "text":
                                      f"Your previous code had an error. Please fix it and return a new code block. The error information is given below:\n{answer_code}"}]})
        raw_code = run_llm(system_code, msg_code, model_id)
        Py_code = _extract_code_block(raw_code)
        print("🧪 analyze code obtained")
        print(Py_code)

        answer_code = _exec_code_get_answer(Py_code, extracted)
        retry_count += 1
        if retry_count >= retry_limit:
            break

    return {"answer": answer_code,
            "code": raw_code}
