"""Core contract generation engine for SolAgent-FHE."""

import os
import time

from solagent_fhe.llm import chat, extract_solidity
from solagent_fhe.compiler import run_compile, write_contract


# Pre-defined FHE tasks for benchmarking
FHE_TASKS = {
    "voting": {
        "name": "ConfidentialVoting",
        "description": "Create a confidential voting contract. Features: admin sets options count, voters cast one vote each (public option ID), vote counts are encrypted with euint32. After admin closes voting, anyone can request decryption of results.",
    },
    "auction": {
        "name": "SealedBidAuction",
        "description": "Create a sealed-bid auction contract. Features: auctioneer deploys, bidders submit encrypted bids (InEuint64). Contract tracks highest bid (euint64) and highest bidder (eaddress) using FHE.gt/max/select. Auctioneer closes auction and results become decryptable.",
    },
    "token": {
        "name": "PrivateToken",
        "description": "Create a private ERC20-like token with encrypted balances (euint64). Features: transfer with encrypted amount (InEuint64), balance check uses FHE.gte + FHE.select (no revert on insufficient — just transfer 0). Owner can decrypt own balance. Operator approval (binary, no amount leak).",
    },
    "lottery": {
        "name": "BlindLottery",
        "description": "Create a blind lottery contract. Features: players enter by paying fixed ETH amount. When admin draws winner, contract generates random number with FHE.randomEuint32(), picks winner index via FHE.rem(). Winner address stored as eaddress, revealed after decryption.",
    },
    "threshold": {
        "name": "SecretThreshold",
        "description": "Create a secret threshold contract. Admin sets an encrypted threshold (euint32). Users submit encrypted values (InEuint32). Contract compares with FHE.gte() and stores result as ebool per user. Admin can decrypt individual results to check who passed.",
    },
    "payroll": {
        "name": "ConfidentialPayroll",
        "description": "Create a confidential payroll contract. Employer registers employees with encrypted salaries (euint64). Each employee can only view their own salary via FHE.allow(salary, msg.sender). Employer can pay all employees in one tx — each gets their encrypted amount transferred. Total payroll is tracked as euint64 (sum of all salaries). Employer can request decryption of total payroll only.",
    },
    "dex": {
        "name": "PrivateDEX",
        "description": "Create a private DEX swap contract. Users deposit token amounts as encrypted values (euint64). When swapping tokenA→tokenB, the swap amount is encrypted (InEuint64). Contract checks FHE.gte(balance, amount) before swap and uses FHE.select to update balances (no revert — just 0 swap if insufficient). Exchange rate is public (uint). Each user's balance of both tokens stays encrypted. FHE.allow(balance, msg.sender) so each user can decrypt own balances.",
    },
    "dao": {
        "name": "PrivateDAO",
        "description": "Create a private DAO governance contract. Members have encrypted voting power (euint32) set by admin. Proposals have title (string) and two encrypted counters: forVotes (euint32) and againstVotes (euint32). Members vote with FHE-encrypted weight: FHE.add updates the right counter. After voting closes, anyone can request decryption of both counters. quorumReached is ebool: FHE.gte(forVotes+againstVotes, encryptedQuorum). Each member can vote only once per proposal.",
    },
    "vault": {
        "name": "FHESavingsVault",
        "description": "Create an FHE savings vault contract. Users deposit ETH and receive encrypted balance (euint64) in wei units. Withdrawals: user provides encrypted amount (InEuint64), contract uses FHE.select to ensure no overdraft. Interest accrual: admin can add encrypted interest (euint64) to all balances via a mapping update. Each user can view only their own balance (FHE.allow). Admin can view total vault balance (euint64, sum of all deposits) after requesting decryption.",
    },
}


def _load_prompt(prompt_type: str) -> str:
    """Load system prompt from markdown file."""
    path = os.path.join(os.path.dirname(__file__), 'prompts', f'system_{prompt_type}.md')
    with open(path) as f:
        return f.read()


async def generate_contract(
    task: str,
    contract_name: str = "Generated",
    provider: str = "deepseek",
    max_refine_rounds: int = 10,
    verbose: bool = True,
) -> dict:
    """
    Generate an FHE contract from natural language description.

    Args:
        task: Natural language description of what the contract should do
        contract_name: Name for the generated contract (no .sol extension)
        provider: LLM provider - "deepseek" or "groq"
        max_refine_rounds: Maximum compilation fix attempts
        verbose: Print progress to stdout

    Returns:
        dict with keys: name, code, compiled (bool), rounds, errors
    """
    system_code = _load_prompt("code")
    system_refine = _load_prompt("refine")

    if verbose:
        print(f"\n{'='*60}")
        print(f"  SolAgent-FHE — Generating: {contract_name}")
        print(f"  Provider: {provider}")
        print(f"{'='*60}\n")

    # Step 1: Code Agent generates initial contract
    if verbose:
        print("[1/3] Code Agent generating contract...")

    messages = [
        {"role": "system", "content": system_code},
        {"role": "user", "content": f"Contract name: {contract_name}\n\nTask: {task}"},
    ]

    t0 = time.time()
    response = await chat(messages, provider=provider)
    code = extract_solidity(response)
    gen_time = time.time() - t0

    if verbose:
        print(f"      Generated {len(code)} chars in {gen_time:.1f}s")

    # Step 2: Compile check
    if verbose:
        print("[2/3] Compiling...")

    write_contract(contract_name, code)
    result = run_compile()

    if result['success']:
        if verbose:
            print("      Compiled successfully on first try!")
        return {
            'name': contract_name,
            'code': code,
            'compiled': True,
            'rounds': 0,
            'errors': [],
        }

    # Step 3: Refine loop
    if verbose:
        print(f"      Compilation errors: {len(result['errors'])}")
        print("[3/3] Refine Agent fixing errors...")

    for round_num in range(1, max_refine_rounds + 1):
        error_text = result['output']

        refine_messages = [
            {"role": "system", "content": system_refine},
            {"role": "user", "content": (
                f"Contract name: {contract_name}\n\n"
                f"Current code:\n```solidity\n{code}\n```\n\n"
                f"Compilation errors:\n```\n{error_text}\n```\n\n"
                f"Fix all errors and return the complete corrected contract."
            )},
        ]

        t0 = time.time()
        response = await chat(refine_messages, provider=provider)
        code = extract_solidity(response)
        refine_time = time.time() - t0

        write_contract(contract_name, code)
        result = run_compile()

        if result['success']:
            if verbose:
                print(f"      Fixed in round {round_num} ({refine_time:.1f}s)")
            return {
                'name': contract_name,
                'code': code,
                'compiled': True,
                'rounds': round_num,
                'errors': [],
            }

        if verbose:
            print(f"      Round {round_num}: still {len(result['errors'])} errors ({refine_time:.1f}s)")

    if verbose:
        print(f"      Failed after {max_refine_rounds} rounds")

    return {
        'name': contract_name,
        'code': code,
        'compiled': False,
        'rounds': max_refine_rounds,
        'errors': result['errors'],
    }


async def run_benchmark(provider: str = "deepseek", tasks: list = None, verbose: bool = True):
    """Run generation on multiple FHE tasks and report results."""
    if tasks is None:
        tasks = list(FHE_TASKS.keys())

    results = []
    total_start = time.time()

    for task_key in tasks:
        task_info = FHE_TASKS[task_key]
        result = await generate_contract(
            task=task_info["description"],
            contract_name=task_info["name"],
            provider=provider,
            verbose=verbose,
        )
        results.append(result)

    total_time = time.time() - total_start

    compiled = sum(1 for r in results if r['compiled'])
    total = len(results)

    if verbose:
        print(f"\n{'='*60}")
        print(f"  BENCHMARK RESULTS")
        print(f"{'='*60}")
        print(f"  Provider: {provider}")
        print(f"  Compiled: {compiled}/{total} ({100*compiled/total:.0f}%)")
        print(f"  Total time: {total_time:.1f}s")
        print()
        for r in results:
            status = "OK" if r['compiled'] else "FAIL"
            rounds = f"(+{r['rounds']} refine)" if r['rounds'] > 0 else "(first try)"
            print(f"  {status} {r['name']} {rounds}")
        print(f"{'='*60}\n")

    return results
